link/macho: introduce Atom extras for out-of-band data

6 files changed, 157 insertions, 41 deletions

diff --git a/src/link/MachO.zig b/src/link/MachO.zig
index 01e03a10ba..7e20c8e97d 100644
--- a/src/link/MachO.zig
+++ b/src/link/MachO.zig
@@ -65,6 +65,7 @@ entry_index: ?Symbol.Index = null,
 
 /// List of atoms that are either synthetic or map directly to the Zig source program.
 atoms: std.ArrayListUnmanaged(Atom) = .{},
+atoms_extra: std.ArrayListUnmanaged(u32) = .{},
 thunks: std.ArrayListUnmanaged(Thunk) = .{},
 unwind_records: std.ArrayListUnmanaged(UnwindInfo.Record) = .{},
 
@@ -246,6 +247,7 @@ pub fn createEmpty(
     try self.files.append(gpa, .null);
     // Atom at index 0 is reserved as null atom
     try self.atoms.append(gpa, .{});
+    try self.atoms_extra.append(gpa, 0);
     // Append empty string to string tables
     try self.strings.buffer.append(gpa, 0);
     try self.strtab.append(gpa, 0);
@@ -350,6 +352,7 @@ pub fn deinit(self: *MachO) void {
     self.unwind_info.deinit(gpa);
 
     self.atoms.deinit(gpa);
+    self.atoms_extra.deinit(gpa);
     for (self.thunks.items) |*thunk| {
         thunk.deinit(gpa);
     }
@@ -3865,6 +3868,50 @@ pub fn getAtom(self: *MachO, index: Atom.Index) ?*Atom {
     return &self.atoms.items[index];
 }
 
+pub fn addAtomExtra(self: *MachO, extra: Atom.Extra) !u32 {
+    const fields = @typeInfo(Atom.Extra).Struct.fields;
+    try self.atoms_extra.ensureUnusedCapacity(self.base.comp.gpa, fields.len);
+    return self.addAtomExtraAssumeCapacity(extra);
+}
+
+pub fn addAtomExtraAssumeCapacity(self: *MachO, extra: Atom.Extra) u32 {
+    const index = @as(u32, @intCast(self.atoms_extra.items.len));
+    const fields = @typeInfo(Atom.Extra).Struct.fields;
+    inline for (fields) |field| {
+        self.atoms_extra.appendAssumeCapacity(switch (field.type) {
+            u32 => @field(extra, field.name),
+            else => @compileError("bad field type"),
+        });
+    }
+    return index;
+}
+
+pub fn getAtomExtra(self: *MachO, index: u32) ?Atom.Extra {
+    if (index == 0) return null;
+    const fields = @typeInfo(Atom.Extra).Struct.fields;
+    var i: usize = index;
+    var result: Atom.Extra = undefined;
+    inline for (fields) |field| {
+        @field(result, field.name) = switch (field.type) {
+            u32 => self.atoms_extra.items[i],
+            else => @compileError("bad field type"),
+        };
+        i += 1;
+    }
+    return result;
+}
+
+pub fn setAtomExtra(self: *MachO, index: u32, extra: Atom.Extra) void {
+    assert(index > 0);
+    const fields = @typeInfo(Atom.Extra).Struct.fields;
+    inline for (fields, 0..) |field, i| {
+        self.atoms_extra.items[index + i] = switch (field.type) {
+            u32 => @field(extra, field.name),
+            else => @compileError("bad field type"),
+        };
+    }
+}
+
 pub fn addSymbol(self: *MachO) !Symbol.Index {
     const index = @as(Symbol.Index, @intCast(self.symbols.items.len));
     const symbol = try self.symbols.addOne(self.base.comp.gpa);
diff --git a/src/link/MachO/Atom.zig b/src/link/MachO/Atom.zig
index 86edceda05..b39157e588 100644
--- a/src/link/MachO/Atom.zig
+++ b/src/link/MachO/Atom.zig
@@ -23,18 +23,9 @@ out_n_sect: u8 = 0,
 /// off + size <= parent section size.
 off: u64 = 0,
 
-/// Relocations of this atom.
-relocs: Loc = .{},
-
 /// Index of this atom in the linker's atoms table.
 atom_index: Index = 0,
 
-/// Index of the thunk for this atom.
-thunk_index: Thunk.Index = 0,
-
-/// Unwind records associated with this atom.
-unwind_records: Loc = .{},
-
 flags: Flags = .{},
 
 /// Points to the previous and next neighbors, based on the `text_offset`.
@@ -42,6 +33,8 @@ flags: Flags = .{},
 prev_index: Index = 0,
 next_index: Index = 0,
 
+extra: u32 = 0,
+
 pub fn getName(self: Atom, macho_file: *MachO) [:0]const u8 {
     return switch (self.getFile(macho_file)) {
         .dylib => unreachable,
@@ -67,7 +60,7 @@ pub fn getData(self: Atom, macho_file: *MachO, buffer: []u8) !void {
 pub fn getRelocs(self: Atom, macho_file: *MachO) []const Relocation {
     return switch (self.getFile(macho_file)) {
         .dylib => unreachable,
-        inline else => |x| x.getAtomRelocs(self),
+        inline else => |x| x.getAtomRelocs(self, macho_file),
     };
 }
 
@@ -95,10 +88,11 @@ pub fn getPriority(self: Atom, macho_file: *MachO) u64 {
 }
 
 pub fn getUnwindRecords(self: Atom, macho_file: *MachO) []const UnwindInfo.Record.Index {
+    if (!self.flags.unwind) return &[0]UnwindInfo.Record.Index{};
+    const extra = self.getExtra(macho_file).?;
     return switch (self.getFile(macho_file)) {
-        .dylib => unreachable,
-        .zig_object, .internal => &[0]UnwindInfo.Record.Index{},
-        .object => |x| x.unwind_records.items[self.unwind_records.pos..][0..self.unwind_records.len],
+        .dylib, .zig_object, .internal => unreachable,
+        .object => |x| x.unwind_records.items[extra.unwind_index..][0..extra.unwind_count],
     };
 }
 
@@ -114,7 +108,38 @@ pub fn markUnwindRecordsDead(self: Atom, macho_file: *MachO) void {
 }
 
 pub fn getThunk(self: Atom, macho_file: *MachO) *Thunk {
-    return macho_file.getThunk(self.thunk_index);
+    assert(self.flags.thunk);
+    const extra = self.getExtra(macho_file).?;
+    return macho_file.getThunk(extra.thunk);
+}
+
+const AddExtraOpts = struct {
+    thunk: ?u32 = null,
+    rel_index: ?u32 = null,
+    rel_count: ?u32 = null,
+    unwind_index: ?u32 = null,
+    unwind_count: ?u32 = null,
+};
+
+pub fn addExtra(atom: *Atom, opts: AddExtraOpts, macho_file: *MachO) !void {
+    if (atom.getExtra(macho_file) == null) {
+        atom.extra = try macho_file.addAtomExtra(.{});
+    }
+    var extra = atom.getExtra(macho_file).?;
+    inline for (@typeInfo(@TypeOf(opts)).Struct.fields) |field| {
+        if (@field(opts, field.name)) |x| {
+            @field(extra, field.name) = x;
+        }
+    }
+    atom.setExtra(extra, macho_file);
+}
+
+pub inline fn getExtra(atom: Atom, macho_file: *MachO) ?Extra {
+    return macho_file.getAtomExtra(atom.extra);
+}
+
+pub inline fn setExtra(atom: Atom, extra: Extra, macho_file: *MachO) void {
+    macho_file.setAtomExtra(atom.extra, extra);
 }
 
 pub fn initOutputSection(sect: macho.section_64, macho_file: *MachO) !u8 {
@@ -403,14 +428,20 @@ pub fn addReloc(self: *Atom, macho_file: *MachO, reloc: Relocation) !void {
     const gpa = macho_file.base.comp.gpa;
     const file = self.getFile(macho_file);
     assert(file == .zig_object);
-    const rels = &file.zig_object.relocs.items[self.relocs.pos];
+    assert(self.flags.relocs);
+    var extra = self.getExtra(macho_file).?;
+    const rels = &file.zig_object.relocs.items[extra.rel_index];
     try rels.append(gpa, reloc);
-    self.relocs.len += 1;
+    extra.rel_count += 1;
+    self.setExtra(extra, macho_file);
 }
 
 pub fn freeRelocs(self: *Atom, macho_file: *MachO) void {
-    self.getFile(macho_file).zig_object.freeAtomRelocs(self.*);
-    self.relocs.len = 0;
+    if (!self.flags.relocs) return;
+    self.getFile(macho_file).zig_object.freeAtomRelocs(self.*, macho_file);
+    var extra = self.getExtra(macho_file).?;
+    extra.rel_count = 0;
+    self.setExtra(extra, macho_file);
 }
 
 pub fn scanRelocs(self: Atom, macho_file: *MachO) !void {
@@ -1117,19 +1148,21 @@ fn format2(
     _ = unused_fmt_string;
     const atom = ctx.atom;
     const macho_file = ctx.macho_file;
-    try writer.print("atom({d}) : {s} : @{x} : sect({d}) : align({x}) : size({x}) : nreloc({d}) : thunk({d})", .{
+    try writer.print("atom({d}) : {s} : @{x} : sect({d}) : align({x}) : size({x}) : nreloc({d})", .{
         atom.atom_index,                atom.getName(macho_file), atom.getAddress(macho_file),
         atom.out_n_sect,                atom.alignment,           atom.size,
-        atom.getRelocs(macho_file).len, atom.thunk_index,
+        atom.getRelocs(macho_file).len,
     });
+    if (atom.flags.thunk) try writer.print(" : thunk({d})", .{atom.getExtra(macho_file).?.thunk});
     if (!atom.flags.alive) try writer.writeAll(" : [*]");
-    if (atom.unwind_records.len > 0) {
+    if (atom.flags.unwind) {
         try writer.writeAll(" : unwind{ ");
-        for (atom.getUnwindRecords(macho_file), atom.unwind_records.pos..) |index, i| {
+        const extra = atom.getExtra(macho_file).?;
+        for (atom.getUnwindRecords(macho_file), extra.unwind_index..) |index, i| {
             const rec = macho_file.getUnwindRecord(index);
             try writer.print("{d}", .{index});
             if (!rec.alive) try writer.writeAll("([*])");
-            if (i < atom.unwind_records.pos + atom.unwind_records.len - 1) try writer.writeAll(", ");
+            if (i < extra.unwind_index + extra.unwind_count - 1) try writer.writeAll(", ");
         }
         try writer.writeAll(" }");
     }
@@ -1143,11 +1176,32 @@ pub const Flags = packed struct {
 
     /// Specifies if the atom has been visited during garbage collection.
     visited: bool = false,
+
+    /// Whether this atom has a range extension thunk.
+    thunk: bool = false,
+
+    /// Whether this atom has any relocations.
+    relocs: bool = false,
+
+    /// Whether this atom has any unwind records.
+    unwind: bool = false,
 };
 
-pub const Loc = struct {
-    pos: u32 = 0,
-    len: u32 = 0,
+pub const Extra = struct {
+    /// Index of the range extension thunk of this atom.
+    thunk: u32 = 0,
+
+    /// Start index of relocations belonging to this atom.
+    rel_index: u32 = 0,
+
+    /// Count of relocations belonging to this atom.
+    rel_count: u32 = 0,
+
+    /// Start index of relocations belonging to this atom.
+    unwind_index: u32 = 0,
+
+    /// Count of relocations belonging to this atom.
+    unwind_count: u32 = 0,
 };
 
 pub const Alignment = @import("../../InternPool.zig").Alignment;
diff --git a/src/link/MachO/InternalObject.zig b/src/link/MachO/InternalObject.zig
index db8a8fd939..9f42eca114 100644
--- a/src/link/MachO/InternalObject.zig
+++ b/src/link/MachO/InternalObject.zig
@@ -115,7 +115,8 @@ fn addObjcSelrefsSection(
             .has_subtractor = false,
         },
     });
-    atom.relocs = .{ .pos = 0, .len = 1 };
+    try atom.addExtra(.{ .rel_index = 0, .rel_count = 1 }, macho_file);
+    atom.flags.relocs = true;
     self.num_rebase_relocs += 1;
 
     return atom_index;
@@ -183,9 +184,11 @@ pub fn getAtomData(self: *const InternalObject, atom: Atom, buffer: []u8) !void
     @memcpy(buffer, data[off..][0..size]);
 }
 
-pub fn getAtomRelocs(self: *const InternalObject, atom: Atom) []const Relocation {
+pub fn getAtomRelocs(self: *const InternalObject, atom: Atom, macho_file: *MachO) []const Relocation {
+    if (!atom.flags.relocs) return &[0]Relocation{};
+    const extra = atom.getExtra(macho_file).?;
     const relocs = self.sections.items(.relocs)[atom.n_sect];
-    return relocs.items[atom.relocs.pos..][0..atom.relocs.len];
+    return relocs.items[extra.rel_index..][0..extra.rel_count];
 }
 
 fn addString(self: *InternalObject, allocator: Allocator, name: [:0]const u8) error{OutOfMemory}!u32 {
diff --git a/src/link/MachO/Object.zig b/src/link/MachO/Object.zig
index 5d2f537a4f..a49000ce30 100644
--- a/src/link/MachO/Object.zig
+++ b/src/link/MachO/Object.zig
@@ -690,11 +690,13 @@ fn initRelocs(self: *Object, macho_file: *MachO) !void {
             if (!atom.flags.alive) continue;
             if (next_reloc >= relocs.items.len) break;
             const end_addr = atom.off + atom.size;
-            atom.relocs.pos = next_reloc;
+            const rel_index = next_reloc;
 
             while (next_reloc < relocs.items.len and relocs.items[next_reloc].offset < end_addr) : (next_reloc += 1) {}
 
-            atom.relocs.len = next_reloc - atom.relocs.pos;
+            const rel_count = next_reloc - rel_index;
+            try atom.addExtra(.{ .rel_index = rel_index, .rel_count = rel_count }, macho_file);
+            atom.flags.relocs = true;
         }
     }
 }
@@ -1004,7 +1006,8 @@ fn parseUnwindRecords(self: *Object, macho_file: *MachO) !void {
         {}
 
         const atom = rec.getAtom(macho_file);
-        atom.unwind_records = .{ .pos = start, .len = next_cu - start };
+        try atom.addExtra(.{ .unwind_index = start, .unwind_count = next_cu - start }, macho_file);
+        atom.flags.unwind = true;
     }
 }
 
@@ -1724,9 +1727,11 @@ pub fn getAtomData(self: *const Object, macho_file: *MachO, atom: Atom, buffer:
     if (amt != buffer.len) return error.InputOutput;
 }
 
-pub fn getAtomRelocs(self: *const Object, atom: Atom) []const Relocation {
+pub fn getAtomRelocs(self: *const Object, atom: Atom, macho_file: *MachO) []const Relocation {
+    if (!atom.flags.relocs) return &[0]Relocation{};
+    const extra = atom.getExtra(macho_file).?;
     const relocs = self.sections.items(.relocs)[atom.n_sect];
-    return relocs.items[atom.relocs.pos..][0..atom.relocs.len];
+    return relocs.items[extra.rel_index..][0..extra.rel_count];
 }
 
 fn addString(self: *Object, allocator: Allocator, name: [:0]const u8) error{OutOfMemory}!u32 {
diff --git a/src/link/MachO/ZigObject.zig b/src/link/MachO/ZigObject.zig
index bb788fcacc..1e0297f5ec 100644
--- a/src/link/MachO/ZigObject.zig
+++ b/src/link/MachO/ZigObject.zig
@@ -163,7 +163,8 @@ pub fn addAtom(self: *ZigObject, macho_file: *MachO) !Symbol.Index {
     const relocs_index = @as(u32, @intCast(self.relocs.items.len));
     const relocs = try self.relocs.addOne(gpa);
     relocs.* = .{};
-    atom.relocs = .{ .pos = relocs_index, .len = 0 };
+    try atom.addExtra(.{ .rel_index = relocs_index, .rel_count = 0 }, macho_file);
+    atom.flags.relocs = true;
 
     return symbol_index;
 }
@@ -190,13 +191,18 @@ pub fn getAtomData(self: ZigObject, macho_file: *MachO, atom: Atom, buffer: []u8
     }
 }
 
-pub fn getAtomRelocs(self: *ZigObject, atom: Atom) []const Relocation {
-    const relocs = self.relocs.items[atom.relocs.pos];
-    return relocs.items[0..atom.relocs.len];
+pub fn getAtomRelocs(self: *ZigObject, atom: Atom, macho_file: *MachO) []const Relocation {
+    if (!atom.flags.relocs) return &[0]Relocation{};
+    const extra = atom.getExtra(macho_file).?;
+    const relocs = self.relocs.items[extra.rel_index];
+    return relocs.items[0..extra.rel_count];
 }
 
-pub fn freeAtomRelocs(self: *ZigObject, atom: Atom) void {
-    self.relocs.items[atom.relocs.pos].clearRetainingCapacity();
+pub fn freeAtomRelocs(self: *ZigObject, atom: Atom, macho_file: *MachO) void {
+    if (atom.flags.relocs) {
+        const extra = atom.getExtra(macho_file).?;
+        self.relocs.items[extra.rel_index].clearRetainingCapacity();
+    }
 }
 
 pub fn resolveSymbols(self: *ZigObject, macho_file: *MachO) void {
diff --git a/src/link/MachO/thunks.zig b/src/link/MachO/thunks.zig
index f66a990b1d..710dac6d42 100644
--- a/src/link/MachO/thunks.zig
+++ b/src/link/MachO/thunks.zig
@@ -43,7 +43,8 @@ pub fn createThunks(sect_id: u8, macho_file: *MachO) !void {
                 if (isReachable(atom, rel, macho_file)) continue;
                 try thunk.symbols.put(gpa, rel.target, {});
             }
-            atom.thunk_index = thunk_index;
+            try atom.addExtra(.{ .thunk = thunk_index }, macho_file);
+            atom.flags.thunk = true;
         }
 
         thunk.value = try advance(header, thunk.size(), .@"4");