link/macho: apply fixes to deduping logic

* test non-ObjC literal deduping logic

5 files changed, 105 insertions, 42 deletions

diff --git a/src/link/MachO.zig b/src/link/MachO.zig
index 5dadf8a60c..326f714796 100644
--- a/src/link/MachO.zig
+++ b/src/link/MachO.zig
@@ -1500,14 +1500,25 @@ pub fn dedupLiterals(self: *MachO) !void {
     const gpa = self.base.comp.gpa;
     var lp: LiteralPool = .{};
     defer lp.deinit(gpa);
+
     if (self.getZigObject()) |zo| {
-        try zo.dedupLiterals(&lp, self);
+        try zo.resolveLiterals(&lp, self);
     }
     for (self.objects.items) |index| {
-        try self.getFile(index).?.object.dedupLiterals(&lp, self);
+        try self.getFile(index).?.object.resolveLiterals(&lp, self);
     }
     if (self.getInternalObject()) |object| {
-        try object.dedupLiterals(&lp, self);
+        try object.resolveLiterals(&lp, self);
+    }
+
+    if (self.getZigObject()) |zo| {
+        zo.dedupLiterals(lp, self);
+    }
+    for (self.objects.items) |index| {
+        self.getFile(index).?.object.dedupLiterals(lp, self);
+    }
+    if (self.getInternalObject()) |object| {
+        object.dedupLiterals(lp, self);
     }
 }
 
@@ -4415,8 +4426,14 @@ pub const LiteralPool = struct {
         lp.data.deinit(allocator);
     }
 
+    pub fn getAtom(lp: LiteralPool, index: Index, macho_file: *MachO) *Atom {
+        assert(index < lp.values.items.len);
+        return macho_file.getAtom(lp.values.items[index]).?;
+    }
+
     const InsertResult = struct {
         found_existing: bool,
+        index: Index,
         atom: *Atom.Index,
     };
 
@@ -4434,6 +4451,7 @@ pub const LiteralPool = struct {
         }
         return .{
             .found_existing = gop.found_existing,
+            .index = @intCast(gop.index),
             .atom = &lp.values.items[gop.index],
         };
     }
@@ -4472,6 +4490,8 @@ pub const LiteralPool = struct {
             return key.hash(ctx.lp);
         }
     };
+
+    pub const Index = u32;
 };
 
 const HotUpdateState = struct {
diff --git a/src/link/MachO/Atom.zig b/src/link/MachO/Atom.zig
index e37412dd53..e31619162e 100644
--- a/src/link/MachO/Atom.zig
+++ b/src/link/MachO/Atom.zig
@@ -113,12 +113,18 @@ pub fn getThunk(self: Atom, macho_file: *MachO) *Thunk {
     return macho_file.getThunk(extra.thunk);
 }
 
+pub fn getLiteralPoolIndex(self: Atom, macho_file: *MachO) ?MachO.LiteralPool.Index {
+    if (!self.flags.literal_pool) return null;
+    return self.getExtra(macho_file).?.literal_index;
+}
+
 const AddExtraOpts = struct {
     thunk: ?u32 = null,
     rel_index: ?u32 = null,
     rel_count: ?u32 = null,
     unwind_index: ?u32 = null,
     unwind_count: ?u32 = null,
+    literal_index: ?u32 = null,
 };
 
 pub fn addExtra(atom: *Atom, opts: AddExtraOpts, macho_file: *MachO) !void {
@@ -1177,7 +1183,7 @@ pub const Flags = packed struct {
     /// Specifies whether this atom is alive or has been garbage collected.
     alive: bool = true,
 
-    /// Specifies if the atom has been visited during garbage collection.
+    /// Specifies if this atom has been visited during garbage collection.
     visited: bool = false,
 
     /// Whether this atom has a range extension thunk.
@@ -1188,6 +1194,9 @@ pub const Flags = packed struct {
 
     /// Whether this atom has any unwind records.
     unwind: bool = false,
+
+    /// Whether this atom has LiteralPool entry.
+    literal_pool: bool = false,
 };
 
 pub const Extra = struct {
@@ -1205,6 +1214,9 @@ pub const Extra = struct {
 
     /// Count of relocations belonging to this atom.
     unwind_count: u32 = 0,
+
+    /// Index into LiteralPool entry for this atom.
+    literal_index: u32 = 0,
 };
 
 pub const Alignment = @import("../../InternPool.zig").Alignment;
diff --git a/src/link/MachO/InternalObject.zig b/src/link/MachO/InternalObject.zig
index f25508f037..cea32ca233 100644
--- a/src/link/MachO/InternalObject.zig
+++ b/src/link/MachO/InternalObject.zig
@@ -109,12 +109,9 @@ fn addObjcSelrefsSection(self: *InternalObject, methname_atom_index: Atom.Index,
     return atom_index;
 }
 
-pub fn dedupLiterals(self: InternalObject, lp: *MachO.LiteralPool, macho_file: *MachO) !void {
+pub fn resolveLiterals(self: InternalObject, lp: *MachO.LiteralPool, macho_file: *MachO) !void {
     const gpa = macho_file.base.comp.gpa;
 
-    var killed_atoms = std.AutoHashMap(Atom.Index, Atom.Index).init(gpa);
-    defer killed_atoms.deinit();
-
     var buffer = std.ArrayList(u8).init(gpa);
     defer buffer.deinit();
 
@@ -126,10 +123,9 @@ pub fn dedupLiterals(self: InternalObject, lp: *MachO.LiteralPool, macho_file: *
             const res = try lp.insert(gpa, header.type(), data);
             if (!res.found_existing) {
                 res.atom.* = atom_index;
-                continue;
             }
-            atom.flags.alive = false;
-            try killed_atoms.putNoClobber(atom_index, res.atom.*);
+            atom.flags.literal_pool = true;
+            try atom.addExtra(.{ .literal_index = res.index }, macho_file);
         } else if (Object.isPtrLiteral(header)) {
             const atom = macho_file.getAtom(atom_index).?;
             const relocs = atom.getRelocs(macho_file);
@@ -145,32 +141,45 @@ pub fn dedupLiterals(self: InternalObject, lp: *MachO.LiteralPool, macho_file: *
             buffer.clearRetainingCapacity();
             if (!res.found_existing) {
                 res.atom.* = atom_index;
-                continue;
             }
-            atom.flags.alive = false;
-            try killed_atoms.putNoClobber(atom_index, res.atom.*);
+            atom.flags.literal_pool = true;
+            try atom.addExtra(.{ .literal_index = res.index }, macho_file);
         }
     }
+}
 
+pub fn dedupLiterals(self: InternalObject, lp: MachO.LiteralPool, macho_file: *MachO) void {
     for (self.atoms.items) |atom_index| {
-        if (killed_atoms.get(atom_index)) |_| continue;
         const atom = macho_file.getAtom(atom_index) orelse continue;
         if (!atom.flags.alive) continue;
         if (!atom.flags.relocs) continue;
 
         const relocs = blk: {
             const extra = atom.getExtra(macho_file).?;
-            const relocs = slice.items(.relocs)[atom.n_sect].items;
+            const relocs = self.sections.items(.relocs)[atom.n_sect].items;
             break :blk relocs[extra.rel_index..][0..extra.rel_count];
         };
         for (relocs) |*rel| switch (rel.tag) {
-            .local => if (killed_atoms.get(rel.target)) |new_target| {
-                rel.target = new_target;
+            .local => {
+                const target = macho_file.getAtom(rel.target).?;
+                if (target.getLiteralPoolIndex(macho_file)) |lp_index| {
+                    const lp_atom = lp.getAtom(lp_index, macho_file);
+                    if (target.atom_index != lp_atom.atom_index) {
+                        target.flags.alive = false;
+                        rel.target = lp_atom.atom_index;
+                    }
+                }
             },
             .@"extern" => {
-                const target = rel.getTargetSymbol(macho_file);
-                if (killed_atoms.get(target.atom)) |new_atom| {
-                    target.atom = new_atom;
+                const target_sym = rel.getTargetSymbol(macho_file);
+                if (target_sym.getAtom(macho_file)) |target_atom| {
+                    if (target_atom.getLiteralPoolIndex(macho_file)) |lp_index| {
+                        const lp_atom = lp.getAtom(lp_index, macho_file);
+                        if (target_atom.atom_index != lp_atom.atom_index) {
+                            target_atom.flags.alive = false;
+                            target_sym.atom = lp_atom.atom_index;
+                        }
+                    }
                 }
             },
         };
@@ -179,9 +188,15 @@ pub fn dedupLiterals(self: InternalObject, lp: *MachO.LiteralPool, macho_file: *
     for (self.symbols.items) |sym_index| {
         const sym = macho_file.getSymbol(sym_index);
         if (!sym.flags.objc_stubs) continue;
-        const extra = sym.getExtra(macho_file).?;
-        if (killed_atoms.get(extra.objc_selrefs)) |new_atom| {
-            try sym.addExtra(.{ .objc_selrefs = new_atom }, macho_file);
+        var extra = sym.getExtra(macho_file).?;
+        const atom = macho_file.getAtom(extra.objc_selrefs).?;
+        if (atom.getLiteralPoolIndex(macho_file)) |lp_index| {
+            const lp_atom = lp.getAtom(lp_index, macho_file);
+            if (atom.atom_index != lp_atom.atom_index) {
+                atom.flags.alive = false;
+                extra.objc_selrefs = lp_atom.atom_index;
+                sym.setExtra(extra, macho_file);
+            }
         }
     }
 }
diff --git a/src/link/MachO/Object.zig b/src/link/MachO/Object.zig
index d32ace7a3e..42f5d8d021 100644
--- a/src/link/MachO/Object.zig
+++ b/src/link/MachO/Object.zig
@@ -527,12 +527,9 @@ fn initPointerLiterals(self: *Object, macho_file: *MachO) !void {
     }
 }
 
-pub fn dedupLiterals(self: Object, lp: *MachO.LiteralPool, macho_file: *MachO) !void {
+pub fn resolveLiterals(self: Object, lp: *MachO.LiteralPool, macho_file: *MachO) !void {
     const gpa = macho_file.base.comp.gpa;
 
-    var killed_atoms = std.AutoHashMap(Atom.Index, Atom.Index).init(gpa);
-    defer killed_atoms.deinit();
-
     var buffer = std.ArrayList(u8).init(gpa);
     defer buffer.deinit();
 
@@ -548,10 +545,9 @@ pub fn dedupLiterals(self: Object, lp: *MachO.LiteralPool, macho_file: *MachO) !
                 const res = try lp.insert(gpa, header.type(), atom_data);
                 if (!res.found_existing) {
                     res.atom.* = sub.atom;
-                    continue;
                 }
-                atom.flags.alive = false;
-                try killed_atoms.putNoClobber(sub.atom, res.atom.*);
+                atom.flags.literal_pool = true;
+                try atom.addExtra(.{ .literal_index = res.index }, macho_file);
             }
         } else if (isPtrLiteral(header)) {
             for (subs.items) |sub| {
@@ -572,33 +568,46 @@ pub fn dedupLiterals(self: Object, lp: *MachO.LiteralPool, macho_file: *MachO) !
                 buffer.clearRetainingCapacity();
                 if (!res.found_existing) {
                     res.atom.* = sub.atom;
-                    continue;
                 }
-                atom.flags.alive = false;
-                try killed_atoms.putNoClobber(sub.atom, res.atom.*);
+                atom.flags.literal_pool = true;
+                try atom.addExtra(.{ .literal_index = res.index }, macho_file);
             }
         }
     }
+}
 
+pub fn dedupLiterals(self: Object, lp: MachO.LiteralPool, macho_file: *MachO) void {
     for (self.atoms.items) |atom_index| {
-        if (killed_atoms.get(atom_index)) |_| continue;
         const atom = macho_file.getAtom(atom_index) orelse continue;
         if (!atom.flags.alive) continue;
         if (!atom.flags.relocs) continue;
 
         const relocs = blk: {
             const extra = atom.getExtra(macho_file).?;
-            const relocs = slice.items(.relocs)[atom.n_sect].items;
+            const relocs = self.sections.items(.relocs)[atom.n_sect].items;
             break :blk relocs[extra.rel_index..][0..extra.rel_count];
         };
         for (relocs) |*rel| switch (rel.tag) {
-            .local => if (killed_atoms.get(rel.target)) |new_target| {
-                rel.target = new_target;
+            .local => {
+                const target = macho_file.getAtom(rel.target).?;
+                if (target.getLiteralPoolIndex(macho_file)) |lp_index| {
+                    const lp_atom = lp.getAtom(lp_index, macho_file);
+                    if (target.atom_index != lp_atom.atom_index) {
+                        target.flags.alive = false;
+                        rel.target = lp_atom.atom_index;
+                    }
+                }
             },
             .@"extern" => {
-                const target = rel.getTargetSymbol(macho_file);
-                if (killed_atoms.get(target.atom)) |new_atom| {
-                    target.atom = new_atom;
+                const target_sym = rel.getTargetSymbol(macho_file);
+                if (target_sym.getAtom(macho_file)) |target_atom| {
+                    if (target_atom.getLiteralPoolIndex(macho_file)) |lp_index| {
+                        const lp_atom = lp.getAtom(lp_index, macho_file);
+                        if (target_atom.atom_index != lp_atom.atom_index) {
+                            target_atom.flags.alive = false;
+                            target_sym.atom = lp_atom.atom_index;
+                        }
+                    }
                 }
             },
         };
diff --git a/src/link/MachO/ZigObject.zig b/src/link/MachO/ZigObject.zig
index 338840e521..6168d40d96 100644
--- a/src/link/MachO/ZigObject.zig
+++ b/src/link/MachO/ZigObject.zig
@@ -314,7 +314,14 @@ pub fn checkDuplicates(self: *ZigObject, dupes: anytype, macho_file: *MachO) !vo
     }
 }
 
-pub fn dedupLiterals(self: *ZigObject, lp: *MachO.LiteralPool, macho_file: *MachO) !void {
+pub fn resolveLiterals(self: *ZigObject, lp: *MachO.LiteralPool, macho_file: *MachO) !void {
+    _ = self;
+    _ = lp;
+    _ = macho_file;
+    // TODO
+}
+
+pub fn dedupLiterals(self: *ZigObject, lp: MachO.LiteralPool, macho_file: *MachO) void {
     _ = self;
     _ = lp;
     _ = macho_file;