macho: tweak section and segment sorting logic

kubkon · Jan 23, 2024 · 422d3f5 · 422d3f5
1 parent 059af7e
commit 422d3f5
Showing 1 changed file with 54 additions and 34 deletions.
diff --git a/src/MachO.zig b/src/MachO.zig
@@ -1536,50 +1536,71 @@ fn getSegmentProt(segname: []const u8) macho.vm_prot_t {
 
 fn getSegmentRank(segname: []const u8) u4 {
     if (mem.eql(u8, segname, "__PAGEZERO")) return 0x0;
-    if (mem.eql(u8, segname, "__TEXT")) return 0x1;
-    if (mem.eql(u8, segname, "__DATA_CONST")) return 0x2;
-    if (mem.eql(u8, segname, "__DATA")) return 0x3;
     if (mem.eql(u8, segname, "__LINKEDIT")) return 0x5;
+    if (mem.startsWith(u8, segname, "__TEXT")) return 0x1;
+    if (mem.startsWith(u8, segname, "__DATA_CONST")) return 0x2;
+    if (mem.startsWith(u8, segname, "__DATA")) return 0x3;
     return 0x4;
 }
 
-fn getSectionRank(self: *MachO, sect_index: u8) u8 {
-    const header = self.sections.items(.header)[sect_index];
-    const segment_rank = getSegmentRank(header.segName());
-    const section_rank: u4 = blk: {
-        if (header.isCode()) {
-            if (mem.eql(u8, "__text", header.sectName())) break :blk 0x0;
-            if (header.type() == macho.S_SYMBOL_STUBS) break :blk 0x1;
-            break :blk 0x2;
-        }
-        switch (header.type()) {
-            macho.S_NON_LAZY_SYMBOL_POINTERS,
-            macho.S_LAZY_SYMBOL_POINTERS,
-            => break :blk 0x0,
-
-            macho.S_MOD_INIT_FUNC_POINTERS => break :blk 0x1,
-            macho.S_MOD_TERM_FUNC_POINTERS => break :blk 0x2,
-            macho.S_ZEROFILL => break :blk 0xf,
-            macho.S_THREAD_LOCAL_REGULAR => break :blk 0xd,
-            macho.S_THREAD_LOCAL_ZEROFILL => break :blk 0xe,
-
-            else => {
-                if (mem.eql(u8, "__unwind_info", header.sectName())) break :blk 0xe;
-                if (mem.eql(u8, "__compact_unwind", header.sectName())) break :blk 0xe;
-                if (mem.eql(u8, "__eh_frame", header.sectName())) break :blk 0xf;
-                break :blk 0x3;
-            },
+fn segmentLessThan(ctx: void, lhs: []const u8, rhs: []const u8) bool {
+    _ = ctx;
+    const lhs_rank = getSegmentRank(lhs);
+    const rhs_rank = getSegmentRank(rhs);
+    if (lhs_rank == rhs_rank) {
+        return mem.order(u8, lhs, rhs) == .lt;
+    }
+    return lhs_rank < rhs_rank;
+}
+
+fn getSectionRank(section: macho.section_64) u8 {
+    if (section.isCode()) {
+        if (mem.eql(u8, "__text", section.sectName())) return 0x0;
+        if (section.type() == macho.S_SYMBOL_STUBS) return 0x1;
+        return 0x2;
+    }
+    switch (section.type()) {
+        macho.S_NON_LAZY_SYMBOL_POINTERS,
+        macho.S_LAZY_SYMBOL_POINTERS,
+        => return 0x0,
+
+        macho.S_MOD_INIT_FUNC_POINTERS => return 0x1,
+        macho.S_MOD_TERM_FUNC_POINTERS => return 0x2,
+        macho.S_ZEROFILL => return 0xf,
+        macho.S_THREAD_LOCAL_REGULAR => return 0xd,
+        macho.S_THREAD_LOCAL_ZEROFILL => return 0xe,
+
+        else => {
+            if (mem.eql(u8, "__unwind_info", section.sectName())) return 0xe;
+            if (mem.eql(u8, "__compact_unwind", section.sectName())) return 0xe;
+            if (mem.eql(u8, "__eh_frame", section.sectName())) return 0xf;
+            return 0x3;
+        },
+    }
+}
+
+fn sectionLessThan(ctx: void, lhs: macho.section_64, rhs: macho.section_64) bool {
+    if (mem.eql(u8, lhs.segName(), rhs.segName())) {
+        const lhs_rank = getSectionRank(lhs);
+        const rhs_rank = getSectionRank(rhs);
+        if (lhs_rank == rhs_rank) {
+            return mem.order(u8, lhs.sectName(), rhs.sectName()) == .lt;
         }
-    };
-    return (@as(u8, @intCast(segment_rank)) << 4) + section_rank;
+        return lhs_rank < rhs_rank;
+    }
+    return segmentLessThan(ctx, lhs.segName(), rhs.segName());
 }
 
 pub fn sortSections(self: *MachO) !void {
     const Entry = struct {
         index: u8,
 
         pub fn lessThan(macho_file: *MachO, lhs: @This(), rhs: @This()) bool {
-            return macho_file.getSectionRank(lhs.index) < macho_file.getSectionRank(rhs.index);
+            return sectionLessThan(
+                {},
+                macho_file.sections.items(.header)[lhs.index],
+                macho_file.sections.items(.header)[rhs.index],
+            );
         }
     };
 
@@ -1839,8 +1860,7 @@ fn initSegments(self: *MachO) !void {
 
     const sortFn = struct {
         fn sortFn(ctx: void, lhs: macho.segment_command_64, rhs: macho.segment_command_64) bool {
-            _ = ctx;
-            return getSegmentRank(lhs.segName()) < getSegmentRank(rhs.segName());
+            return segmentLessThan(ctx, lhs.segName(), rhs.segName());
         }
     }.sortFn;