zld: track symbols defined within TextBlock

in case TextBlock represents an entire section with symbols
defined within.

2 files changed, 100 insertions, 48 deletions

diff --git a/src/link/MachO/Object.zig b/src/link/MachO/Object.zig
index 867821fa6f..a3b1ff4f79 100644
--- a/src/link/MachO/Object.zig
+++ b/src/link/MachO/Object.zig
@@ -430,17 +430,12 @@ const TextBlockParser = struct {
         const block = try self.allocator.create(TextBlock);
         errdefer self.allocator.destroy(block);
 
-        block.* = .{
-            .local_sym_index = senior_nlist.index,
-            .aliases = alias_only_indices,
-            .references = std.AutoArrayHashMap(u32, void).init(self.allocator),
-            .code = try self.allocator.dupe(u8, code),
-            .relocs = std.ArrayList(Relocation).init(self.allocator),
-            .rebases = std.ArrayList(u64).init(self.allocator),
-            .tlv_offsets = std.ArrayList(TextBlock.TlvOffset).init(self.allocator),
-            .size = size,
-            .alignment = self.section.@"align",
-        };
+        block.* = TextBlock.init(self.allocator);
+        block.local_sym_index = senior_nlist.index;
+        block.aliases = alias_only_indices;
+        block.code = try self.allocator.dupe(u8, code);
+        block.size = size;
+        block.alignment = self.section.@"align";
 
         const relocs = filterRelocs(self.relocs, start_addr, end_addr);
         if (relocs.len > 0) {
@@ -499,16 +494,17 @@ pub fn parseTextBlocks(self: *Object, zld: *Zld) !void {
         _ = try self.file.?.preadAll(raw_relocs, sect.reloff);
         const relocs = mem.bytesAsSlice(macho.relocation_info, raw_relocs);
 
+        // Symbols within this section only.
+        const filtered_nlists = NlistWithIndex.filterInSection(
+            sorted_nlists.items,
+            sect_id + 1,
+        );
+
         // Is there any padding between symbols within the section?
         const is_splittable = self.header.?.flags & macho.MH_SUBSECTIONS_VIA_SYMBOLS != 0;
 
         next: {
             if (is_splittable) blocks: {
-                const filtered_nlists = NlistWithIndex.filterInSection(
-                    sorted_nlists.items,
-                    sect_id + 1,
-                );
-
                 if (filtered_nlists.len == 0) break :blocks;
 
                 var parser = TextBlockParser{
@@ -528,7 +524,7 @@ pub fn parseTextBlocks(self: *Object, zld: *Zld) !void {
                     if (reg.file) |file| {
                         if (file != self) {
                             log.warn("deduping definition of {s} in {s}", .{ sym.name, self.name.? });
-                            block.deinit(self.allocator);
+                            block.deinit();
                             self.allocator.destroy(block);
                             continue;
                         }
@@ -583,21 +579,43 @@ pub fn parseTextBlocks(self: *Object, zld: *Zld) !void {
             const block = try self.allocator.create(TextBlock);
             errdefer self.allocator.destroy(block);
 
-            block.* = .{
-                .local_sym_index = local_sym_index,
-                .references = std.AutoArrayHashMap(u32, void).init(self.allocator),
-                .code = try self.allocator.dupe(u8, code),
-                .relocs = std.ArrayList(Relocation).init(self.allocator),
-                .rebases = std.ArrayList(u64).init(self.allocator),
-                .tlv_offsets = std.ArrayList(TextBlock.TlvOffset).init(self.allocator),
-                .size = sect.size,
-                .alignment = sect.@"align",
-            };
+            block.* = TextBlock.init(self.allocator);
+            block.local_sym_index = local_sym_index;
+            block.code = try self.allocator.dupe(u8, code);
+            block.size = sect.size;
+            block.alignment = sect.@"align";
 
             if (relocs.len > 0) {
                 try self.parseRelocs(zld, relocs, block, 0);
             }
 
+            // Since this is block gets a helper local temporary symbol that didn't exist
+            // in the object file which encompasses the entire section, we need traverse
+            // the filtered symbols and note which symbol is contained within so that
+            // we can properly allocate addresses down the line.
+            // While we're at it, we need to update segment,section mapping of each symbol too.
+            if (filtered_nlists.len > 0) {
+                var contained = std.ArrayList(TextBlock.SymbolAtOffset).init(self.allocator);
+                defer contained.deinit();
+                try contained.ensureTotalCapacity(filtered_nlists.len);
+
+                for (filtered_nlists) |nlist_with_index| {
+                    const sym = self.symbols.items[nlist_with_index.index];
+                    assert(sym.payload == .regular);
+                    const reg = &sym.payload.regular;
+
+                    reg.segment_id = match.seg;
+                    reg.section_id = match.sect;
+
+                    contained.appendAssumeCapacity(.{
+                        .local_sym_index = reg.local_sym_index,
+                        .offset = nlist_with_index.nlist.n_value - sect.addr,
+                    });
+                }
+
+                block.contained = contained.toOwnedSlice();
+            }
+
             // Update target section's metadata
             // TODO should we update segment's size here too?
             // How does it tie with incremental space allocs?
diff --git a/src/link/MachO/Zld.zig b/src/link/MachO/Zld.zig
index e137e16e9b..cdc0732c34 100644
--- a/src/link/MachO/Zld.zig
+++ b/src/link/MachO/Zld.zig
@@ -121,9 +121,11 @@ pub const Output = struct {
 };
 
 pub const TextBlock = struct {
+    allocator: *Allocator,
     local_sym_index: u32,
     aliases: ?[]u32 = null,
     references: std.AutoArrayHashMap(u32, void),
+    contained: ?[]SymbolAtOffset = null,
     code: []u8,
     relocs: std.ArrayList(Relocation),
     size: u64,
@@ -133,20 +135,42 @@ pub const TextBlock = struct {
     next: ?*TextBlock = null,
     prev: ?*TextBlock = null,
 
+    pub const SymbolAtOffset = struct {
+        local_sym_index: u32,
+        offset: u64,
+    };
+
     pub const TlvOffset = struct {
         local_sym_index: u32,
         offset: u64,
     };
 
-    pub fn deinit(block: *TextBlock, allocator: *Allocator) void {
-        if (block.aliases) |aliases| {
-            allocator.free(aliases);
+    pub fn init(allocator: *Allocator) TextBlock {
+        return .{
+            .allocator = allocator,
+            .local_sym_index = undefined,
+            .references = std.AutoArrayHashMap(u32, void).init(allocator),
+            .code = undefined,
+            .relocs = std.ArrayList(Relocation).init(allocator),
+            .size = undefined,
+            .alignment = undefined,
+            .rebases = std.ArrayList(u64).init(allocator),
+            .tlv_offsets = std.ArrayList(TextBlock.TlvOffset).init(allocator),
+        };
+    }
+
+    pub fn deinit(self: *TextBlock) void {
+        if (self.aliases) |aliases| {
+            self.allocator.free(aliases);
+        }
+        self.references.deinit();
+        if (self.contained) |contained| {
+            self.allocator.free(contained);
         }
-        block.relocs.deinit();
-        block.references.deinit();
-        block.rebases.deinit();
-        block.tlv_offsets.deinit();
-        allocator.free(block.code);
+        self.allocator.free(self.code);
+        self.relocs.deinit();
+        self.rebases.deinit();
+        self.tlv_offsets.deinit();
     }
 
     pub fn print_this(self: *const TextBlock, zld: *Zld) void {
@@ -164,6 +188,12 @@ pub const TextBlock = struct {
                 log.warn("    | {}: {}", .{ index, zld.locals.items[index] });
             }
         }
+        if (self.contained) |contained| {
+            log.warn("  | contained symbols:", .{});
+            for (contained) |sym_at_off| {
+                log.warn("    | {}: {}", .{ sym_at_off.offset, zld.locals.items[sym_at_off.local_sym_index] });
+            }
+        }
         log.warn("  | code.len = {}", .{self.code.len});
         if (self.relocs.items.len > 0) {
             log.warn("  | relocations:", .{});
@@ -1021,10 +1051,20 @@ fn allocateTextBlocks(self: *Zld) !void {
         var base_addr: u64 = sect.addr + sect.size;
 
         while (true) {
+            base_addr -= block.size;
+
             const sym = self.locals.items[block.local_sym_index];
             assert(sym.payload == .regular);
-            sym.payload.regular.address = base_addr - block.size;
-            base_addr -= block.size;
+            sym.payload.regular.address = base_addr;
+
+            // Update each symbol contained within the TextBlock
+            if (block.contained) |contained| {
+                for (contained) |sym_at_off| {
+                    const contained_sym = self.locals.items[sym_at_off.local_sym_index];
+                    assert(contained_sym.payload == .regular);
+                    contained_sym.payload.regular.address = base_addr + sym_at_off.offset;
+                }
+            }
 
             if (block.prev) |prev| {
                 block = prev;
@@ -1476,16 +1516,11 @@ fn resolveSymbols(self: *Zld) !void {
                 const block = try self.allocator.create(TextBlock);
                 errdefer self.allocator.destroy(block);
 
-                block.* = .{
-                    .local_sym_index = local_sym_index,
-                    .references = std.AutoArrayHashMap(u32, void).init(self.allocator),
-                    .code = code,
-                    .relocs = std.ArrayList(Relocation).init(self.allocator),
-                    .rebases = std.ArrayList(u64).init(self.allocator),
-                    .tlv_offsets = std.ArrayList(TextBlock.TlvOffset).init(self.allocator),
-                    .size = size,
-                    .alignment = alignment,
-                };
+                block.* = TextBlock.init(self.allocator);
+                block.local_sym_index = local_sym_index;
+                block.code = code;
+                block.size = size;
+                block.alignment = alignment;
 
                 if (self.blocks.getPtr(match)) |last| {
                     last.*.next = block;
@@ -1907,7 +1942,6 @@ fn addRpaths(self: *Zld, rpaths: []const []const u8) !void {
 
 fn flush(self: *Zld) !void {
     try self.writeStubHelperCommon();
-    try self.resolveRelocsAndWriteSections();
 
     if (self.common_section_index) |index| {
         const seg = &self.load_commands.items[self.data_segment_cmd_index.?].Segment;