macho: assign and cache section ordinals upon creation

then, when sorting sections within segments, clear and redo the
ordinals since we re-apply them to symbols anyway. It is vital
to have the ordinals consistent with parsing and resolving relocs
however.

3 files changed, 88 insertions, 39 deletions

diff --git a/src/link/MachO.zig b/src/link/MachO.zig
index d742e6ec12..b2f048d78b 100644
--- a/src/link/MachO.zig
+++ b/src/link/MachO.zig
@@ -167,6 +167,9 @@ strtab_needs_relocation: bool = false,
 has_dices: bool = false,
 has_stabs: bool = false,
 
+section_ordinals: std.ArrayListUnmanaged(MatchingSection) = .{},
+section_to_ordinal: std.AutoHashMapUnmanaged(MatchingSection, u8) = .{},
+
 pending_updates: std.ArrayListUnmanaged(struct {
     kind: enum {
         got,
@@ -925,6 +928,13 @@ fn linkWithZld(self: *MachO, comp: *Compilation) !void {
             else => unreachable,
         };
 
+        // Initialize section ordinals with null ordinal pointing at
+        // PAGEZERO segment.
+        try self.section_ordinals.append(self.base.allocator, .{
+            .seg = 0,
+            .sect = 0,
+        });
+
         try self.populateMetadata();
         try self.parseInputFiles(positionals.items, self.base.options.sysroot);
         try self.parseLibs(libs.items, self.base.options.sysroot);
@@ -1482,6 +1492,10 @@ pub fn getMatchingSection(self: *MachO, sect: macho.section_64) !?MatchingSectio
         }
     };
 
+    if (res) |match| {
+        try self.createSectionOrdinal(match);
+    }
+
     return res;
 }
 
@@ -1606,6 +1620,38 @@ fn sortSections(self: *MachO) !void {
         self.blocks.deinit(self.base.allocator);
         self.blocks = transient;
     }
+
+    {
+        // Create new section ordinals.
+        self.section_ordinals.clearRetainingCapacity();
+        self.section_to_ordinal.clearRetainingCapacity();
+        // First ordinal is always null
+        self.section_ordinals.appendAssumeCapacity(.{
+            .seg = 0,
+            .sect = 0,
+        });
+        const text_seg = self.load_commands.items[self.text_segment_cmd_index.?].Segment;
+        for (text_seg.sections.items) |_, sect_id| {
+            try self.createSectionOrdinal(.{
+                .seg = self.text_segment_cmd_index.?,
+                .sect = @intCast(u16, sect_id),
+            });
+        }
+        const data_const_seg = self.load_commands.items[self.data_const_segment_cmd_index.?].Segment;
+        for (data_const_seg.sections.items) |_, sect_id| {
+            try self.createSectionOrdinal(.{
+                .seg = self.data_const_segment_cmd_index.?,
+                .sect = @intCast(u16, sect_id),
+            });
+        }
+        const data_seg = self.load_commands.items[self.data_segment_cmd_index.?].Segment;
+        for (data_seg.sections.items) |_, sect_id| {
+            try self.createSectionOrdinal(.{
+                .seg = self.data_segment_cmd_index.?,
+                .sect = @intCast(u16, sect_id),
+            });
+        }
+    }
 }
 
 fn allocateTextSegment(self: *MachO) !void {
@@ -1732,7 +1778,7 @@ fn allocateTextBlocks(self: *MachO) !void {
         const sect = seg.sections.items[match.sect];
 
         var base_addr: u64 = sect.addr;
-        const n_sect = self.sectionId(match);
+        const n_sect = self.section_to_ordinal.get(match) orelse unreachable;
 
         log.debug("  within section {s},{s}", .{ commands.segmentName(sect), commands.sectionName(sect) });
         log.debug("    {}", .{sect});
@@ -2260,6 +2306,7 @@ fn resolveSymbols(self: *MachO) !void {
                 .sect = self.common_section_index.?,
             };
         };
+        try self.createSectionOrdinal(match);
 
         const size = sym.n_value;
         const code = try self.base.allocator.alloc(u8, size);
@@ -2272,7 +2319,7 @@ fn resolveSymbols(self: *MachO) !void {
         var nlist = macho.nlist_64{
             .n_strx = sym.n_strx,
             .n_type = macho.N_SECT,
-            .n_sect = self.sectionId(match),
+            .n_sect = self.section_to_ordinal.get(match) orelse unreachable,
             .n_desc = 0,
             .n_value = 0,
         };
@@ -2402,7 +2449,7 @@ fn resolveSymbols(self: *MachO) !void {
         var nlist = macho.nlist_64{
             .n_strx = undef.n_strx,
             .n_type = macho.N_SECT,
-            .n_sect = self.sectionId(match),
+            .n_sect = self.section_to_ordinal.get(match) orelse unreachable,
             .n_desc = 0,
             .n_value = 0,
         };
@@ -2498,6 +2545,10 @@ fn populateMetadata(self: *MachO) !void {
             .@"align" = alignment,
             .flags = macho.S_REGULAR | macho.S_ATTR_PURE_INSTRUCTIONS | macho.S_ATTR_SOME_INSTRUCTIONS,
         });
+        try self.createSectionOrdinal(.{
+            .seg = self.text_segment_cmd_index.?,
+            .sect = self.text_section_index.?,
+        });
     }
 
     if (self.stubs_section_index == null) {
@@ -2518,6 +2569,10 @@ fn populateMetadata(self: *MachO) !void {
             .flags = macho.S_SYMBOL_STUBS | macho.S_ATTR_PURE_INSTRUCTIONS | macho.S_ATTR_SOME_INSTRUCTIONS,
             .reserved2 = stub_size,
         });
+        try self.createSectionOrdinal(.{
+            .seg = self.text_segment_cmd_index.?,
+            .sect = self.stubs_section_index.?,
+        });
     }
 
     if (self.stub_helper_section_index == null) {
@@ -2538,6 +2593,10 @@ fn populateMetadata(self: *MachO) !void {
             .@"align" = alignment,
             .flags = macho.S_REGULAR | macho.S_ATTR_PURE_INSTRUCTIONS | macho.S_ATTR_SOME_INSTRUCTIONS,
         });
+        try self.createSectionOrdinal(.{
+            .seg = self.text_segment_cmd_index.?,
+            .sect = self.stub_helper_section_index.?,
+        });
     }
 
     if (self.data_const_segment_cmd_index == null) {
@@ -2557,6 +2616,10 @@ fn populateMetadata(self: *MachO) !void {
             .@"align" = 3, // 2^3 = @sizeOf(u64)
             .flags = macho.S_NON_LAZY_SYMBOL_POINTERS,
         });
+        try self.createSectionOrdinal(.{
+            .seg = self.data_const_segment_cmd_index.?,
+            .sect = self.got_section_index.?,
+        });
     }
 
     if (self.data_segment_cmd_index == null) {
@@ -2576,6 +2639,10 @@ fn populateMetadata(self: *MachO) !void {
             .@"align" = 3, // 2^3 = @sizeOf(u64)
             .flags = macho.S_LAZY_SYMBOL_POINTERS,
         });
+        try self.createSectionOrdinal(.{
+            .seg = self.data_segment_cmd_index.?,
+            .sect = self.la_symbol_ptr_section_index.?,
+        });
     }
 
     if (self.data_section_index == null) {
@@ -2584,6 +2651,10 @@ fn populateMetadata(self: *MachO) !void {
         try data_seg.addSection(self.base.allocator, "__data", .{
             .@"align" = 3, // 2^3 = @sizeOf(u64)
         });
+        try self.createSectionOrdinal(.{
+            .seg = self.data_segment_cmd_index.?,
+            .sect = self.data_section_index.?,
+        });
     }
 
     if (self.linkedit_segment_cmd_index == null) {
@@ -3290,6 +3361,8 @@ pub fn deinit(self: *MachO) void {
         ds.deinit(self.base.allocator);
     }
 
+    self.section_ordinals.deinit(self.base.allocator);
+    self.section_to_ordinal.deinit(self.base.allocator);
     self.pending_updates.deinit(self.base.allocator);
     self.got_entries.deinit(self.base.allocator);
     self.got_entries_map.deinit(self.base.allocator);
@@ -5816,37 +5889,6 @@ pub fn symbolIsTemp(sym: macho.nlist_64, sym_name: []const u8) bool {
     return mem.startsWith(u8, sym_name, "l") or mem.startsWith(u8, sym_name, "L");
 }
 
-pub fn sectionId(self: MachO, match: MatchingSection) u8 {
-    // TODO there might be a more generic way of doing this.
-    var section: u8 = 0;
-    for (self.load_commands.items) |cmd, cmd_id| {
-        if (cmd != .Segment) break;
-        if (cmd_id == match.seg) {
-            section += @intCast(u8, match.sect) + 1;
-            break;
-        }
-        section += @intCast(u8, cmd.Segment.sections.items.len);
-    }
-    return section;
-}
-
-pub fn unpackSectionId(self: MachO, section_id: u8) MatchingSection {
-    var match: MatchingSection = undefined;
-    var section: u8 = 0;
-    outer: for (self.load_commands.items) |cmd, cmd_id| {
-        assert(cmd == .Segment);
-        for (cmd.Segment.sections.items) |_, sect_id| {
-            section += 1;
-            if (section_id == section) {
-                match.seg = @intCast(u16, cmd_id);
-                match.sect = @intCast(u16, sect_id);
-                break :outer;
-            }
-        }
-    }
-    return match;
-}
-
 fn packDylibOrdinal(ordinal: u16) u16 {
     return ordinal * macho.N_SYMBOL_RESOLVER;
 }
@@ -5867,3 +5909,10 @@ pub fn findFirst(comptime T: type, haystack: []T, start: usize, predicate: anyty
     }
     return i;
 }
+
+fn createSectionOrdinal(self: *MachO, match: MatchingSection) !void {
+    if (self.section_to_ordinal.contains(match)) return;
+    const ordinal = @intCast(u8, self.section_ordinals.items.len);
+    try self.section_ordinals.append(self.base.allocator, match);
+    try self.section_to_ordinal.putNoClobber(self.base.allocator, match, ordinal);
+}
diff --git a/src/link/MachO/Object.zig b/src/link/MachO/Object.zig
index c5ff19b9be..846b87a65a 100644
--- a/src/link/MachO/Object.zig
+++ b/src/link/MachO/Object.zig
@@ -733,7 +733,7 @@ pub fn parseTextBlocks(self: *Object, macho_file: *MachO) !void {
                 try macho_file.locals.append(macho_file.base.allocator, .{
                     .n_strx = try macho_file.makeString(sym_name),
                     .n_type = macho.N_SECT,
-                    .n_sect = macho_file.sectionId(match),
+                    .n_sect = macho_file.section_to_ordinal.get(match) orelse unreachable,
                     .n_desc = 0,
                     .n_value = sect.addr,
                 });
@@ -779,7 +779,7 @@ pub fn parseTextBlocks(self: *Object, macho_file: *MachO) !void {
                 const nlist = nlist_with_index.nlist;
                 const local_sym_index = self.symbol_mapping.get(nlist_with_index.index) orelse unreachable;
                 const local = &macho_file.locals.items[local_sym_index];
-                local.n_sect = macho_file.sectionId(match);
+                local.n_sect = macho_file.section_to_ordinal.get(match) orelse unreachable;
 
                 const stab: ?TextBlock.Stab = if (self.debug_info) |di| blk: {
                     // TODO there has to be a better to handle this.
diff --git a/src/link/MachO/TextBlock.zig b/src/link/MachO/TextBlock.zig
index 9dc02d1f4d..17d1d82db8 100644
--- a/src/link/MachO/TextBlock.zig
+++ b/src/link/MachO/TextBlock.zig
@@ -620,7 +620,7 @@ fn initRelocFromObject(rel: macho.relocation_info, object: *Object, ctx: RelocCo
             try ctx.macho_file.locals.append(ctx.macho_file.base.allocator, .{
                 .n_strx = try ctx.macho_file.makeString(sym_name),
                 .n_type = macho.N_SECT,
-                .n_sect = ctx.macho_file.sectionId(match),
+                .n_sect = ctx.macho_file.section_to_ordinal.get(match) orelse unreachable,
                 .n_desc = 0,
                 .n_value = sect.addr,
             });
@@ -832,7 +832,7 @@ pub fn parseRelocsFromObject(
                 },
                 .local => {
                     const source_sym = ctx.macho_file.locals.items[self.local_sym_index];
-                    const match = ctx.macho_file.unpackSectionId(source_sym.n_sect);
+                    const match = ctx.macho_file.section_ordinals.items[source_sym.n_sect];
                     const seg = ctx.macho_file.load_commands.items[match.seg].Segment;
                     const sect = seg.sections.items[match.sect];
                     const sect_type = commands.sectionType(sect);
@@ -1096,7 +1096,7 @@ pub fn resolveRelocs(self: *TextBlock, macho_file: *MachO) !void {
                     const sym = macho_file.locals.items[rel.where_index];
                     const is_tlv = is_tlv: {
                         const source_sym = macho_file.locals.items[self.local_sym_index];
-                        const match = macho_file.unpackSectionId(source_sym.n_sect);
+                        const match = macho_file.section_ordinals.items[source_sym.n_sect];
                         const seg = macho_file.load_commands.items[match.seg].Segment;
                         const sect = seg.sections.items[match.sect];
                         break :is_tlv commands.sectionType(sect) == macho.S_THREAD_LOCAL_VARIABLES;