From 801c37209889883ad2d7b9cfb232fd8d9e0ed66c Mon Sep 17 00:00:00 2001
From: Jakub Konka <kubkon@jakubkonka.com>
Date: Sun, 28 Jul 2024 11:36:51 +0200
Subject: [PATCH] elf: init output merge sections in a separate step

---
 src/link/Elf.zig               | 10 ++++++++--
 src/link/Elf/Object.zig        | 17 ++++++++++++++---
 src/link/Elf/merge_section.zig |  7 +++----
 3 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/src/link/Elf.zig b/src/link/Elf.zig
index 4da9ffeb33..030c0499cc 100644
--- a/src/link/Elf.zig
+++ b/src/link/Elf.zig
@@ -3308,7 +3308,7 @@ pub fn resolveMergeSections(self: *Elf) !void {
     for (self.objects.items) |index| {
         const file_ptr = self.file(index).?;
         if (!file_ptr.isAlive()) continue;
-        file_ptr.object.initMergeSections(self) catch |err| switch (err) {
+        file_ptr.object.initInputMergeSections(self) catch |err| switch (err) {
             error.MalformedObject => has_errors = true,
             else => |e| return e,
         };
@@ -3316,6 +3316,12 @@ pub fn resolveMergeSections(self: *Elf) !void {
 
     if (has_errors) return error.FlushFailure;
 
+    for (self.objects.items) |index| {
+        const file_ptr = self.file(index).?;
+        if (!file_ptr.isAlive()) continue;
+        try file_ptr.object.initOutputMergeSections(self);
+    }
+
     for (self.objects.items) |index| {
         const file_ptr = self.file(index).?;
         if (!file_ptr.isAlive()) continue;
@@ -3330,7 +3336,7 @@ pub fn resolveMergeSections(self: *Elf) !void {
 
 pub fn finalizeMergeSections(self: *Elf) !void {
     for (self.merge_sections.items) |*msec| {
-        try msec.finalize(self);
+        try msec.finalize(self.base.comp.gpa);
     }
 }
 
diff --git a/src/link/Elf/Object.zig b/src/link/Elf/Object.zig
index 40c1896282..4944ef904e 100644
--- a/src/link/Elf/Object.zig
+++ b/src/link/Elf/Object.zig
@@ -691,7 +691,7 @@ pub fn checkDuplicates(self: *Object, dupes: anytype, elf_file: *Elf) error{OutO
     }
 }
 
-pub fn initMergeSections(self: *Object, elf_file: *Elf) !void {
+pub fn initInputMergeSections(self: *Object, elf_file: *Elf) !void {
     const gpa = elf_file.base.comp.gpa;
 
     try self.input_merge_sections.ensureUnusedCapacity(gpa, self.shdrs.items.len);
@@ -709,8 +709,6 @@ pub fn initMergeSections(self: *Object, elf_file: *Elf) !void {
         const imsec_idx = try self.addInputMergeSection(gpa);
         const imsec = self.inputMergeSection(imsec_idx).?;
         self.input_merge_sections_indexes.items[shndx] = imsec_idx;
-
-        imsec.merge_section_index = try elf_file.getOrCreateMergeSection(atom_ptr.name(elf_file), shdr.sh_flags, shdr.sh_type);
         imsec.atom_index = atom_index;
 
         const data = try self.codeDecompressAlloc(elf_file, atom_index);
@@ -769,6 +767,19 @@ pub fn initMergeSections(self: *Object, elf_file: *Elf) !void {
     }
 }
 
+pub fn initOutputMergeSections(self: *Object, elf_file: *Elf) !void {
+    for (self.input_merge_sections_indexes.items) |index| {
+        const imsec = self.inputMergeSection(index) orelse continue;
+        const atom_ptr = self.atom(imsec.atom_index).?;
+        const shdr = atom_ptr.inputShdr(elf_file);
+        imsec.merge_section_index = try elf_file.getOrCreateMergeSection(
+            atom_ptr.name(elf_file),
+            shdr.sh_flags,
+            shdr.sh_type,
+        );
+    }
+}
+
 pub fn resolveMergeSubsections(self: *Object, elf_file: *Elf) !void {
     const gpa = elf_file.base.comp.gpa;
 
diff --git a/src/link/Elf/merge_section.zig b/src/link/Elf/merge_section.zig
index 3dabe8eda3..b4b670062c 100644
--- a/src/link/Elf/merge_section.zig
+++ b/src/link/Elf/merge_section.zig
@@ -60,9 +60,8 @@ pub const MergeSection = struct {
 
     /// Finalizes the merge section and clears hash table.
     /// Sorts all owned subsections.
-    pub fn finalize(msec: *MergeSection, elf_file: *Elf) !void {
-        const gpa = elf_file.base.comp.gpa;
-        try msec.finalized_subsections.ensureTotalCapacityPrecise(gpa, msec.subsections.items.len);
+    pub fn finalize(msec: *MergeSection, allocator: Allocator) !void {
+        try msec.finalized_subsections.ensureTotalCapacityPrecise(allocator, msec.subsections.items.len);
 
         var it = msec.table.iterator();
         while (it.next()) |entry| {
@@ -70,7 +69,7 @@ pub const MergeSection = struct {
             if (!msub.alive) continue;
             msec.finalized_subsections.appendAssumeCapacity(entry.value_ptr.*);
         }
-        msec.table.clearAndFree(gpa);
+        msec.table.clearAndFree(allocator);
 
         const sortFn = struct {
             pub fn sortFn(ctx: *MergeSection, lhs: MergeSubsection.Index, rhs: MergeSubsection.Index) bool {