1 files changed, 13 insertions, 22 deletions

diff --git a/src/link/MachO/Trie.zig b/src/link/MachO/Trie.zig
index 0016ff329c..e445ea4fef 100644
--- a/src/link/MachO/Trie.zig
+++ b/src/link/MachO/Trie.zig
@@ -241,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 += sizeULEB128Mem(info.export_flags);
-            node_size += sizeULEB128Mem(info.vmaddr_offset);
-            node_size += sizeULEB128Mem(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
         }
@@ -254,27 +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 + sizeULEB128Mem(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 };
-    }
-
-    /// Calculates number of bytes in ULEB128 encoding of value.
-    fn sizeULEB128Mem(value: u64) usize {
-        var res: usize = 0;
-        var v = value;
-        while (true) {
-            v = v >> 7;
-            res += 1;
-            if (v == 0) break;
-        }
-        return res;
+        return FinalizeResult{ .node_size = node_size, .updated = updated };
     }
 };
 
@@ -358,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;
         }
@@ -380,9 +373,7 @@ pub fn read(self: *Trie, reader: anytype) ReadError!usize {
 }
 
 /// Write the trie to a byte stream.
-/// Caller owns the memory and needs to free it.
-/// Panics if the trie was not finalized using `finalize`
-/// before calling this method.
+/// Panics if the trie was not finalized using `finalize` before calling this method.
 pub fn write(self: Trie, writer: anytype) !usize {
     assert(!self.trie_dirty);
     var counting_writer = std.io.countingWriter(writer);