macho: refactor calculating LEB128 sizes

3 files changed, 60 insertions, 30 deletions

diff --git a/src/link/MachO.zig b/src/link/MachO.zig
index e71b89e0b6..a80e272935 100644
--- a/src/link/MachO.zig
+++ b/src/link/MachO.zig
@@ -826,7 +826,7 @@ fn linkWithLLD(self: *MachO, comp: *Compilation) !void {
                 // Write update dyld info
                 const dyld_info = self.load_commands.items[self.dyld_info_cmd_index.?].DyldInfoOnly;
                 {
-                    const size = self.binding_info_table.calcSize();
+                    const size = try self.binding_info_table.calcSize();
                     assert(dyld_info.bind_size >= size);
 
                     var buffer = try self.base.allocator.alloc(u8, size);
@@ -838,7 +838,7 @@ fn linkWithLLD(self: *MachO, comp: *Compilation) !void {
                     try self.base.file.?.pwriteAll(buffer, dyld_info.bind_off);
                 }
                 {
-                    const size = self.lazy_binding_info_table.calcSize();
+                    const size = try self.lazy_binding_info_table.calcSize();
                     assert(dyld_info.lazy_bind_size >= size);
 
                     var buffer = try self.base.allocator.alloc(u8, size);
@@ -2204,15 +2204,3 @@ fn parseLazyBindingInfoTable(self: *MachO) !void {
     var stream = std.io.fixedBufferStream(buffer);
     try self.lazy_binding_info_table.read(stream.reader(), self.base.allocator);
 }
-
-/// Calculates number of bytes in LEB128 encoding of value.
-pub fn sizeLEB128(value: anytype) usize {
-    var res: usize = 0;
-    var v = value;
-    while (true) {
-        v = v >> 7;
-        res += 1;
-        if (v == 0) break;
-    }
-    return res;
-}
diff --git a/src/link/MachO/Trie.zig b/src/link/MachO/Trie.zig
index 690588cd53..e445ea4fef 100644
--- a/src/link/MachO/Trie.zig
+++ b/src/link/MachO/Trie.zig
@@ -38,7 +38,6 @@ const macho = std.macho;
 const testing = std.testing;
 const assert = std.debug.assert;
 const Allocator = mem.Allocator;
-const sizeLEB128 = @import("../MachO.zig").sizeLEB128;
 
 pub const Node = struct {
     base: *Trie,
@@ -242,12 +241,15 @@ pub const Node = struct {
     };
 
     /// Updates offset of this node in the output byte stream.
-    fn finalize(self: *Node, offset_in_trie: usize) FinalizeResult {
+    fn finalize(self: *Node, offset_in_trie: usize) !FinalizeResult {
+        var stream = std.io.countingWriter(std.io.null_writer);
+        var writer = stream.writer();
+
         var node_size: usize = 0;
         if (self.terminal_info) |info| {
-            node_size += sizeLEB128(info.export_flags);
-            node_size += sizeLEB128(info.vmaddr_offset);
-            node_size += sizeLEB128(node_size);
+            try leb.writeULEB128(writer, info.export_flags);
+            try leb.writeULEB128(writer, info.vmaddr_offset);
+            try leb.writeULEB128(writer, stream.bytes_written);
         } else {
             node_size += 1; // 0x0 for non-terminal nodes
         }
@@ -255,15 +257,17 @@ pub const Node = struct {
 
         for (self.edges.items) |edge| {
             const next_node_offset = edge.to.trie_offset orelse 0;
-            node_size += edge.label.len + 1 + sizeLEB128(next_node_offset);
+            node_size += edge.label.len + 1;
+            try leb.writeULEB128(writer, next_node_offset);
         }
 
         const trie_offset = self.trie_offset orelse 0;
         const updated = offset_in_trie != trie_offset;
         self.trie_offset = offset_in_trie;
         self.node_dirty = false;
+        node_size += stream.bytes_written;
 
-        return .{ .node_size = node_size, .updated = updated };
+        return FinalizeResult{ .node_size = node_size, .updated = updated };
     }
 };
 
@@ -347,7 +351,7 @@ pub fn finalize(self: *Trie) !void {
         self.size = 0;
         more = false;
         for (self.ordered_nodes.items) |node| {
-            const res = node.finalize(self.size);
+            const res = try node.finalize(self.size);
             self.size += res.node_size;
             if (res.updated) more = true;
         }
diff --git a/src/link/MachO/imports.zig b/src/link/MachO/imports.zig
index 6c03649b84..988c72c151 100644
--- a/src/link/MachO/imports.zig
+++ b/src/link/MachO/imports.zig
@@ -5,7 +5,6 @@ const mem = std.mem;
 
 const assert = std.debug.assert;
 const Allocator = mem.Allocator;
-const sizeLEB128 = @import("../MachO.zig").sizeLEB128;
 
 /// Table of binding info entries used to tell the dyld which
 /// symbols to bind at loading time.
@@ -27,6 +26,9 @@ pub const BindingInfoTable = struct {
 
         /// Offset of this symbol wrt to the segment id encoded in `segment`.
         offset: i64,
+
+        /// Addend value (if any).
+        addend: ?i64 = null,
     };
 
     pub fn deinit(self: *BindingInfoTable, allocator: *Allocator) void {
@@ -91,6 +93,9 @@ pub const BindingInfoTable = struct {
                 macho.BIND_OPCODE_SET_TYPE_IMM => {
                     self.binding_type = imm;
                 },
+                macho.BIND_OPCODE_SET_ADDEND_SLEB => {
+                    symbol.addend = try leb.readILEB128(i64, reader);
+                },
                 else => {
                     std.log.warn("unhandled BIND_OPCODE_: 0x{x}", .{opcode});
                 },
@@ -121,6 +126,11 @@ pub const BindingInfoTable = struct {
             try writer.writeByte(macho.BIND_OPCODE_SET_SEGMENT_AND_OFFSET_ULEB | @truncate(u4, symbol.segment));
             try leb.writeILEB128(writer, symbol.offset);
 
+            if (symbol.addend) |addend| {
+                try writer.writeByte(macho.BIND_OPCODE_SET_ADDEND_SLEB);
+                try leb.writeILEB128(writer, addend);
+            }
+
             try writer.writeByte(macho.BIND_OPCODE_DO_BIND);
         }
 
@@ -128,10 +138,13 @@ pub const BindingInfoTable = struct {
     }
 
     /// Calculate size in bytes of this binding info table.
-    pub fn calcSize(self: *BindingInfoTable) usize {
+    pub fn calcSize(self: *BindingInfoTable) !usize {
+        var stream = std.io.countingWriter(std.io.null_writer);
+        var writer = stream.writer();
         var size: usize = 1;
+
         if (self.dylib_ordinal > 15) {
-            size += sizeLEB128(self.dylib_ordinal);
+            try leb.writeULEB128(writer, @bitCast(u64, self.dylib_ordinal));
         }
 
         size += 1;
@@ -144,12 +157,17 @@ pub const BindingInfoTable = struct {
             }
 
             size += 1;
-            size += sizeLEB128(symbol.offset);
+            try leb.writeILEB128(writer, symbol.offset);
+
+            if (symbol.addend) |addend| {
+                size += 1;
+                try leb.writeILEB128(writer, addend);
+            }
 
             size += 1;
         }
 
-        size += 1;
+        size += 1 + stream.bytes_written;
         return size;
     }
 };
@@ -173,6 +191,9 @@ pub const LazyBindingInfoTable = struct {
 
         /// Id of the segment where to bind this symbol to.
         segment: u8,
+
+        /// Addend value (if any).
+        addend: ?i64 = null,
     };
 
     pub fn deinit(self: *LazyBindingInfoTable, allocator: *Allocator) void {
@@ -232,6 +253,9 @@ pub const LazyBindingInfoTable = struct {
                 macho.BIND_OPCODE_SET_DYLIB_ORDINAL_ULEB => {
                     symbol.dylib_ordinal = try leb.readILEB128(i64, reader);
                 },
+                macho.BIND_OPCODE_SET_ADDEND_SLEB => {
+                    symbol.addend = try leb.readILEB128(i64, reader);
+                },
                 else => {
                     std.log.warn("unhandled BIND_OPCODE_: 0x{x}", .{opcode});
                 },
@@ -246,6 +270,11 @@ pub const LazyBindingInfoTable = struct {
             try writer.writeByte(macho.BIND_OPCODE_SET_SEGMENT_AND_OFFSET_ULEB | @truncate(u4, symbol.segment));
             try leb.writeILEB128(writer, symbol.offset);
 
+            if (symbol.addend) |addend| {
+                try writer.writeByte(macho.BIND_OPCODE_SET_ADDEND_SLEB);
+                try leb.writeILEB128(writer, addend);
+            }
+
             if (symbol.dylib_ordinal > 15) {
                 try writer.writeByte(macho.BIND_OPCODE_SET_DYLIB_ORDINAL_ULEB);
                 try leb.writeULEB128(writer, @bitCast(u64, symbol.dylib_ordinal));
@@ -267,15 +296,23 @@ pub const LazyBindingInfoTable = struct {
     }
 
     /// Calculate size in bytes of this binding info table.
-    pub fn calcSize(self: *LazyBindingInfoTable) usize {
+    pub fn calcSize(self: *LazyBindingInfoTable) !usize {
+        var stream = std.io.countingWriter(std.io.null_writer);
+        var writer = stream.writer();
         var size: usize = 0;
 
         for (self.symbols.items) |symbol| {
             size += 1;
-            size += sizeLEB128(symbol.offset);
+            try leb.writeILEB128(writer, symbol.offset);
+
+            if (symbol.addend) |addend| {
+                size += 1;
+                try leb.writeILEB128(writer, addend);
+            }
+
             size += 1;
             if (symbol.dylib_ordinal > 15) {
-                size += sizeLEB128(symbol.dylib_ordinal);
+                try leb.writeULEB128(writer, @bitCast(u64, symbol.dylib_ordinal));
             }
             if (symbol.name) |name| {
                 size += 1;
@@ -285,6 +322,7 @@ pub const LazyBindingInfoTable = struct {
             size += 2;
         }
 
+        size += stream.bytes_written;
         return size;
     }
 };