macho: do not preempt segment headers; do it when commiting to file

This way, tracking segment-to-section mapping becomes a lot easier
since it's effectively just start index plus number of sections
defined within the segment. If a section becomes empty however
care needs to be taken to remove the header upon committing to the
final binary.

3 files changed, 35 insertions, 13 deletions

diff --git a/src/link/MachO.zig b/src/link/MachO.zig
index 9d3ca34b9e..352e3bedf8 100644
--- a/src/link/MachO.zig
+++ b/src/link/MachO.zig
@@ -4888,13 +4888,26 @@ fn getSegmentAllocBase(self: MachO, indices: []const ?u8) struct { vmaddr: u64,
 
 fn writeSegmentHeaders(self: *MachO, ncmds: *u32, writer: anytype) !void {
     for (self.segments.items) |seg, i| {
-        if (seg.nsects == 0 and
-            (mem.eql(u8, seg.segName(), "__DATA_CONST") or
-            mem.eql(u8, seg.segName(), "__DATA"))) continue;
-        try writer.writeStruct(seg);
-
         const indexes = self.getSectionIndexes(@intCast(u8, i));
+        var out_seg = seg;
+        out_seg.cmdsize = @sizeOf(macho.segment_command_64);
+        out_seg.nsects = 0;
+
+        // Update section headers count; any section with size of 0 is excluded
+        // since it doesn't have any data in the final binary file.
+        for (self.sections.items(.header)[indexes.start..indexes.end]) |header| {
+            if (header.size == 0) continue;
+            out_seg.cmdsize += @sizeOf(macho.section_64);
+            out_seg.nsects += 1;
+        }
+
+        if (out_seg.nsects == 0 and
+            (mem.eql(u8, out_seg.segName(), "__DATA_CONST") or
+            mem.eql(u8, out_seg.segName(), "__DATA"))) continue;
+
+        try writer.writeStruct(out_seg);
         for (self.sections.items(.header)[indexes.start..indexes.end]) |header| {
+            if (header.size == 0) continue;
             try writer.writeStruct(header);
         }
 
diff --git a/src/link/MachO/DebugSymbols.zig b/src/link/MachO/DebugSymbols.zig
index 65d3319293..3bfe334302 100644
--- a/src/link/MachO/DebugSymbols.zig
+++ b/src/link/MachO/DebugSymbols.zig
@@ -367,16 +367,28 @@ fn writeSegmentHeaders(self: *DebugSymbols, ncmds: *u32, writer: anytype) !void
     // Write segment/section headers from the binary file first.
     const end = self.base.linkedit_segment_cmd_index.?;
     for (self.base.segments.items[0..end]) |seg, i| {
-        if (seg.nsects == 0 and
-            (mem.eql(u8, seg.segName(), "__DATA_CONST") or
-            mem.eql(u8, seg.segName(), "__DATA"))) continue;
+        const indexes = self.base.getSectionIndexes(@intCast(u8, i));
         var out_seg = seg;
         out_seg.fileoff = 0;
         out_seg.filesize = 0;
-        try writer.writeStruct(out_seg);
+        out_seg.cmdsize = @sizeOf(macho.segment_command_64);
+        out_seg.nsects = 0;
 
-        const indexes = self.base.getSectionIndexes(@intCast(u8, i));
+        // Update section headers count; any section with size of 0 is excluded
+        // since it doesn't have any data in the final binary file.
+        for (self.base.sections.items(.header)[indexes.start..indexes.end]) |header| {
+            if (header.size == 0) continue;
+            out_seg.cmdsize += @sizeOf(macho.section_64);
+            out_seg.nsects += 1;
+        }
+
+        if (out_seg.nsects == 0 and
+            (mem.eql(u8, out_seg.segName(), "__DATA_CONST") or
+            mem.eql(u8, out_seg.segName(), "__DATA"))) continue;
+
+        try writer.writeStruct(out_seg);
         for (self.base.sections.items(.header)[indexes.start..indexes.end]) |header| {
+            if (header.size == 0) continue;
             var out_header = header;
             out_header.offset = 0;
             try writer.writeStruct(out_header);
diff --git a/src/link/MachO/dead_strip.zig b/src/link/MachO/dead_strip.zig
index bf65b96049..eb2be6e5fe 100644
--- a/src/link/MachO/dead_strip.zig
+++ b/src/link/MachO/dead_strip.zig
@@ -43,9 +43,6 @@ fn removeAtomFromSection(atom: *Atom, match: u8, macho_file: *MachO) void {
             // The section will be GCed in the next step.
             section.last_atom = null;
             section.header.size = 0;
-            const segment = &macho_file.segments.items[section.segment_index];
-            segment.cmdsize -= @sizeOf(macho.section_64);
-            segment.nsects -= 1;
         }
     }