authorgravatar for kubkon@jakubkonka.comJakub Konka <kubkon@jakubkonka.com> 2020-12-08 16:52:50+01:00
committergravatar for kubkon@jakubkonka.comJakub Konka <kubkon@jakubkonka.com> 2020-12-09 20:36:58+01:00
log4c3e6c5bff967388dddc7ec352017c7b712d9f06
tree994837c7efe2d145bdba169985016c7288aa9fff
parentae3fd86dcc6f39601aa26ce1a6dc51a0db5cbf30

macho: cleanup export trie generation and parsing

Now, ExportTrie is becoming usable for larger linking contexts such as linking in multiple object files, or relinking dylibs, etc.

3 files changed, 263 insertions(+), 115 deletions(-)

lib/std/macho.zig+9
...@@ -1333,6 +1333,15 @@ pub const N_WEAK_DEF: u16 = 0x80;...@@ -1333,6 +1333,15 @@ pub const N_WEAK_DEF: u16 = 0x80;
1333/// This bit is only available in .o files (MH_OBJECT filetype)1333/// This bit is only available in .o files (MH_OBJECT filetype)
1334pub const N_SYMBOL_RESOLVER: u16 = 0x100;1334pub const N_SYMBOL_RESOLVER: u16 = 0x100;
13351335
1336// The following are used on the flags byte of a terminal node // in the export information.
1337pub const EXPORT_SYMBOL_FLAGS_KIND_MASK: u8 = 0x03;
1338pub const EXPORT_SYMBOL_FLAGS_KIND_REGULAR: u8 = 0x00;
1339pub const EXPORT_SYMBOL_FLAGS_KIND_THREAD_LOCAL: u8 = 0x01;
1340pub const EXPORT_SYMBOL_FLAGS_KIND_ABSOLUTE: u8 = 0x02;
1341pub const EXPORT_SYMBOL_FLAGS_KIND_WEAK_DEFINITION: u8 = 0x04;
1342pub const EXPORT_SYMBOL_FLAGS_REEXPORT: u8 = 0x08;
1343pub const EXPORT_SYMBOL_FLAGS_STUB_AND_RESOLVER: u8 = 0x10;
1344
1336// Codesign consts and structs taken from:1345// Codesign consts and structs taken from:
1337// https://opensource.apple.com/source/xnu/xnu-6153.81.5/osfmk/kern/cs_blobs.h.auto.html1346// https://opensource.apple.com/source/xnu/xnu-6153.81.5/osfmk/kern/cs_blobs.h.auto.html
13381347
src/link/MachO.zig+47-19
...@@ -754,13 +754,9 @@ fn linkWithLLD(self: *MachO, comp: *Compilation) !void {...@@ -754,13 +754,9 @@ fn linkWithLLD(self: *MachO, comp: *Compilation) !void {
754 const after_last_cmd_offset = self.header.?.sizeofcmds + @sizeOf(macho.mach_header_64);754 const after_last_cmd_offset = self.header.?.sizeofcmds + @sizeOf(macho.mach_header_64);
755 const needed_size = @sizeOf(macho.linkedit_data_command);755 const needed_size = @sizeOf(macho.linkedit_data_command);
756 if (needed_size + after_last_cmd_offset > text_section.offset) {756 if (needed_size + after_last_cmd_offset > text_section.offset) {
757 // TODO We are in the position to be able to increase the padding by moving all sections757 std.log.err("Unable to extend padding between load commands and start of __text section.", .{});
758 // by the required offset, but this requires a little bit more thinking and bookkeeping.758 std.log.err("Re-run the linker with '-headerpad 0x{x}' option if available, or", .{needed_size * alloc_num / alloc_den});
759 // For now, return an error informing the user of the problem.759 std.log.err("fall back to the system linker.", .{});
760 log.err("Not enough padding between load commands and start of __text section:\n", .{});
761 log.err("Offset after last load command: 0x{x}\n", .{after_last_cmd_offset});
762 log.err("Beginning of __text section: 0x{x}\n", .{text_section.offset});
763 log.err("Needed size: 0x{x}\n", .{needed_size});
764 return error.NotEnoughPadding;760 return error.NotEnoughPadding;
765 }761 }
766 const linkedit_segment = self.load_commands.items[self.linkedit_segment_cmd_index.?].Segment;762 const linkedit_segment = self.load_commands.items[self.linkedit_segment_cmd_index.?].Segment;
...@@ -1799,38 +1795,36 @@ fn writeCodeSignature(self: *MachO) !void {...@@ -1799,38 +1795,36 @@ fn writeCodeSignature(self: *MachO) !void {
1799fn writeExportTrie(self: *MachO) !void {1795fn writeExportTrie(self: *MachO) !void {
1800 if (self.global_symbols.items.len == 0) return;1796 if (self.global_symbols.items.len == 0) return;
18011797
1802 var trie: Trie = .{};1798 var trie = Trie.init(self.base.allocator);
1803 defer trie.deinit(self.base.allocator);1799 defer trie.deinit();
18041800
1805 const text_segment = self.load_commands.items[self.text_segment_cmd_index.?].Segment;1801 const text_segment = self.load_commands.items[self.text_segment_cmd_index.?].Segment;
1806 for (self.global_symbols.items) |symbol| {1802 for (self.global_symbols.items) |symbol| {
1807 // TODO figure out if we should put all global symbols into the export trie1803 // TODO figure out if we should put all global symbols into the export trie
1808 const name = self.getString(symbol.n_strx);1804 const name = self.getString(symbol.n_strx);
1809 assert(symbol.n_value >= text_segment.inner.vmaddr);1805 assert(symbol.n_value >= text_segment.inner.vmaddr);
1810 try trie.put(self.base.allocator, .{1806 try trie.put(.{
1811 .name = name,1807 .name = name,
1812 .vmaddr_offset = symbol.n_value - text_segment.inner.vmaddr,1808 .vmaddr_offset = symbol.n_value - text_segment.inner.vmaddr,
1813 .export_flags = 0, // TODO workout creation of export flags1809 .export_flags = 0, // TODO workout creation of export flags
1814 });1810 });
1815 }1811 }
18161812
1817 var buffer: std.ArrayListUnmanaged(u8) = .{};1813 var buffer = try trie.writeULEB128Mem();
1818 defer buffer.deinit(self.base.allocator);1814 defer self.base.allocator.free(buffer);
1819
1820 try trie.writeULEB128Mem(self.base.allocator, &buffer);
18211815
1822 const dyld_info = &self.load_commands.items[self.dyld_info_cmd_index.?].DyldInfoOnly;1816 const dyld_info = &self.load_commands.items[self.dyld_info_cmd_index.?].DyldInfoOnly;
1823 const export_size = @intCast(u32, mem.alignForward(buffer.items.len, @sizeOf(u64)));1817 const export_size = @intCast(u32, mem.alignForward(buffer.len, @sizeOf(u64)));
1824 dyld_info.export_off = self.linkedit_segment_next_offset.?;1818 dyld_info.export_off = self.linkedit_segment_next_offset.?;
1825 dyld_info.export_size = export_size;1819 dyld_info.export_size = export_size;
18261820
1827 log.debug("writing export trie from 0x{x} to 0x{x}\n", .{ dyld_info.export_off, dyld_info.export_off + export_size });1821 log.debug("writing export trie from 0x{x} to 0x{x}\n", .{ dyld_info.export_off, dyld_info.export_off + export_size });
18281822
1829 if (export_size > buffer.items.len) {1823 if (export_size > buffer.len) {
1830 // Pad out to align(8).1824 // Pad out to align(8).
1831 try self.base.file.?.pwriteAll(&[_]u8{0}, dyld_info.export_off + export_size);1825 try self.base.file.?.pwriteAll(&[_]u8{0}, dyld_info.export_off + export_size);
1832 }1826 }
1833 try self.base.file.?.pwriteAll(buffer.items, dyld_info.export_off);1827 try self.base.file.?.pwriteAll(buffer, dyld_info.export_off);
18341828
1835 self.linkedit_segment_next_offset = dyld_info.export_off + dyld_info.export_size;1829 self.linkedit_segment_next_offset = dyld_info.export_off + dyld_info.export_size;
1836 // Advance size of __LINKEDIT segment1830 // Advance size of __LINKEDIT segment
...@@ -1917,7 +1911,9 @@ fn parseFromFile(self: *MachO, file: fs.File) !void {...@@ -1917,7 +1911,9 @@ fn parseFromFile(self: *MachO, file: fs.File) !void {
1917 switch (cmd.cmd()) {1911 switch (cmd.cmd()) {
1918 macho.LC_SEGMENT_64 => {1912 macho.LC_SEGMENT_64 => {
1919 const x = cmd.Segment;1913 const x = cmd.Segment;
1920 if (isSegmentOrSection(&x.inner.segname, "__LINKEDIT")) {1914 if (isSegmentOrSection(&x.inner.segname, "__PAGEZERO")) {
1915 self.pagezero_segment_cmd_index = i;
1916 } else if (isSegmentOrSection(&x.inner.segname, "__LINKEDIT")) {
1921 self.linkedit_segment_cmd_index = i;1917 self.linkedit_segment_cmd_index = i;
1922 } else if (isSegmentOrSection(&x.inner.segname, "__TEXT")) {1918 } else if (isSegmentOrSection(&x.inner.segname, "__TEXT")) {
1923 self.text_segment_cmd_index = i;1919 self.text_segment_cmd_index = i;
...@@ -1926,16 +1922,48 @@ fn parseFromFile(self: *MachO, file: fs.File) !void {...@@ -1926,16 +1922,48 @@ fn parseFromFile(self: *MachO, file: fs.File) !void {
1926 self.text_section_index = @intCast(u16, j);1922 self.text_section_index = @intCast(u16, j);
1927 }1923 }
1928 }1924 }
1925 } else if (isSegmentOrSection(&x.inner.segname, "__DATA")) {
1926 self.data_segment_cmd_index = i;
1929 }1927 }
1930 },1928 },
1929 macho.LC_DYLD_INFO_ONLY => {
1930 self.dyld_info_cmd_index = i;
1931 },
1931 macho.LC_SYMTAB => {1932 macho.LC_SYMTAB => {
1932 self.symtab_cmd_index = i;1933 self.symtab_cmd_index = i;
1933 },1934 },
1935 macho.LC_DYSYMTAB => {
1936 self.dysymtab_cmd_index = i;
1937 },
1938 macho.LC_LOAD_DYLINKER => {
1939 self.dylinker_cmd_index = i;
1940 },
1941 macho.LC_VERSION_MIN_MACOSX, macho.LC_VERSION_MIN_IPHONEOS, macho.LC_VERSION_MIN_WATCHOS, macho.LC_VERSION_MIN_TVOS => {
1942 self.version_min_cmd_index = i;
1943 },
1944 macho.LC_SOURCE_VERSION => {
1945 self.source_version_cmd_index = i;
1946 },
1947 macho.LC_MAIN => {
1948 self.main_cmd_index = i;
1949 },
1950 macho.LC_LOAD_DYLIB => {
1951 self.libsystem_cmd_index = i; // TODO This is incorrect, but we'll fixup later.
1952 },
1953 macho.LC_FUNCTION_STARTS => {
1954 self.function_starts_cmd_index = i;
1955 },
1956 macho.LC_DATA_IN_CODE => {
1957 self.data_in_code_cmd_index = i;
1958 },
1934 macho.LC_CODE_SIGNATURE => {1959 macho.LC_CODE_SIGNATURE => {
1935 self.code_signature_cmd_index = i;1960 self.code_signature_cmd_index = i;
1936 },1961 },
1937 // TODO populate more MachO fields1962 // TODO populate more MachO fields
1938 else => {},1963 else => {
1964 std.log.err("Unknown load command detected: 0x{x}.", .{cmd.cmd()});
1965 return error.UnknownLoadCommand;
1966 },
1939 }1967 }
1940 self.load_commands.appendAssumeCapacity(cmd);1968 self.load_commands.appendAssumeCapacity(cmd);
1941 }1969 }
src/link/MachO/Trie.zig+207-96
...@@ -44,20 +44,23 @@ pub const Symbol = struct {...@@ -44,20 +44,23 @@ pub const Symbol = struct {
44 export_flags: u64,44 export_flags: u64,
45};45};
4646
47const Edge = struct {47pub const Edge = struct {
48 from: *Node,48 from: *Node,
49 to: *Node,49 to: *Node,
50 label: []const u8,50 label: []u8,
5151
52 fn deinit(self: *Edge, alloc: *Allocator) void {52 fn deinit(self: *Edge, allocator: *Allocator) void {
53 self.to.deinit(alloc);53 self.to.deinit();
54 alloc.destroy(self.to);54 allocator.destroy(self.to);
55 allocator.free(self.label);
55 self.from = undefined;56 self.from = undefined;
56 self.to = undefined;57 self.to = undefined;
58 self.label = undefined;
57 }59 }
58};60};
5961
60const Node = struct {62pub const Node = struct {
63 base: *Trie,
61 /// Export flags associated with this exported symbol (if any).64 /// Export flags associated with this exported symbol (if any).
62 export_flags: ?u64 = null,65 export_flags: ?u64 = null,
63 /// VM address offset wrt to the section this symbol is defined against (if any).66 /// VM address offset wrt to the section this symbol is defined against (if any).
...@@ -67,73 +70,97 @@ const Node = struct {...@@ -67,73 +70,97 @@ const Node = struct {
67 /// List of all edges originating from this node.70 /// List of all edges originating from this node.
68 edges: std.ArrayListUnmanaged(Edge) = .{},71 edges: std.ArrayListUnmanaged(Edge) = .{},
6972
70 fn deinit(self: *Node, alloc: *Allocator) void {73 fn deinit(self: *Node) void {
71 for (self.edges.items) |*edge| {74 for (self.edges.items) |*edge| {
72 edge.deinit(alloc);75 edge.deinit(self.base.allocator);
73 }76 }
74 self.edges.deinit(alloc);77 self.edges.deinit(self.base.allocator);
75 }78 }
7679
77 const PutResult = struct {
78 /// Node reached at this stage of `put` op.
79 node: *Node,
80 /// Count of newly inserted nodes at this stage of `put` op.
81 node_count: usize,
82 };
83
84 /// Inserts a new node starting from `self`.80 /// Inserts a new node starting from `self`.
85 fn put(self: *Node, alloc: *Allocator, label: []const u8, node_count: usize) !PutResult {81 fn put(self: *Node, label: []const u8) !*Node {
86 var curr_node_count = node_count;
87 // Check for match with edges from this node.82 // Check for match with edges from this node.
88 for (self.edges.items) |*edge| {83 for (self.edges.items) |*edge| {
89 const match = mem.indexOfDiff(u8, edge.label, label) orelse return PutResult{84 const match = mem.indexOfDiff(u8, edge.label, label) orelse return edge.to;
90 .node = edge.to,
91 .node_count = curr_node_count,
92 };
93 if (match == 0) continue;85 if (match == 0) continue;
94 if (match == edge.label.len) return edge.to.put(alloc, label[match..], curr_node_count);86 if (match == edge.label.len) return edge.to.put(label[match..]);
9587
96 // Found a match, need to splice up nodes.88 // Found a match, need to splice up nodes.
97 // From: A -> B89 // From: A -> B
98 // To: A -> C -> B90 // To: A -> C -> B
99 const mid = try alloc.create(Node);91 const mid = try self.base.allocator.create(Node);
100 mid.* = .{};92 mid.* = .{ .base = self.base };
101 const to_label = edge.label;93 var to_label = try self.base.allocator.dupe(u8, edge.label[match..]);
94 self.base.allocator.free(edge.label);
102 const to_node = edge.to;95 const to_node = edge.to;
103 edge.to = mid;96 edge.to = mid;
104 edge.label = label[0..match];97 edge.label = try self.base.allocator.dupe(u8, label[0..match]);
105 curr_node_count += 1;98 self.base.node_count += 1;
10699
107 try mid.edges.append(alloc, .{100 try mid.edges.append(self.base.allocator, .{
108 .from = mid,101 .from = mid,
109 .to = to_node,102 .to = to_node,
110 .label = to_label[match..],103 .label = to_label,
111 });104 });
112105
113 if (match == label.len) {106 return if (match == label.len) to_node else mid.put(label[match..]);
114 return PutResult{ .node = to_node, .node_count = curr_node_count };
115 } else {
116 return mid.put(alloc, label[match..], curr_node_count);
117 }
118 }107 }
119108
120 // Add a new node.109 // Add a new node.
121 const node = try alloc.create(Node);110 const node = try self.base.allocator.create(Node);
122 node.* = .{};111 node.* = .{ .base = self.base };
123 curr_node_count += 1;112 self.base.node_count += 1;
124113
125 try self.edges.append(alloc, .{114 try self.edges.append(self.base.allocator, .{
126 .from = self,115 .from = self,
127 .to = node,116 .to = node,
128 .label = label,117 .label = try self.base.allocator.dupe(u8, label),
129 });118 });
130119
131 return PutResult{ .node = node, .node_count = curr_node_count };120 return node;
121 }
122
123 fn fromByteStream(self: *Node, stream: anytype) Trie.FromByteStreamError!void {
124 self.trie_offset = try stream.getPos();
125 var reader = stream.reader();
126 const node_size = try leb.readULEB128(u64, reader);
127 if (node_size > 0) {
128 self.export_flags = try leb.readULEB128(u64, reader);
129 // TODO Parse flags.
130 self.vmaddr_offset = try leb.readULEB128(u64, reader);
131 }
132 const nedges = try reader.readByte();
133 self.base.node_count += nedges;
134 var i: usize = 0;
135 while (i < nedges) : (i += 1) {
136 var label = blk: {
137 var label_buf = std.ArrayList(u8).init(self.base.allocator);
138 while (true) {
139 const next = try reader.readByte();
140 if (next == @as(u8, 0))
141 break;
142 try label_buf.append(next);
143 }
144 break :blk label_buf.toOwnedSlice();
145 };
146 const seek_to = try leb.readULEB128(u64, reader);
147 const cur_pos = try stream.getPos();
148 try stream.seekTo(seek_to);
149 var node = try self.base.allocator.create(Node);
150 node.* = .{ .base = self.base };
151 try node.fromByteStream(stream);
152 try self.edges.append(self.base.allocator, .{
153 .from = self,
154 .to = node,
155 .label = label,
156 });
157 try stream.seekTo(cur_pos);
158 }
132 }159 }
133160
134 /// This method should only be called *after* updateOffset has been called!161 /// This method should only be called *after* updateOffset has been called!
135 /// In case this is not upheld, this method will panic.162 /// In case this is not upheld, this method will panic.
136 fn writeULEB128Mem(self: Node, buffer: *std.ArrayListUnmanaged(u8)) !void {163 fn writeULEB128Mem(self: Node, buffer: *std.ArrayList(u8)) !void {
137 assert(self.trie_offset != null); // You need to call updateOffset first.164 assert(self.trie_offset != null); // You need to call updateOffset first.
138 if (self.vmaddr_offset) |offset| {165 if (self.vmaddr_offset) |offset| {
139 // Terminal node info: encode export flags and vmaddr offset of this symbol.166 // Terminal node info: encode export flags and vmaddr offset of this symbol.
...@@ -221,64 +248,95 @@ const Node = struct {...@@ -221,64 +248,95 @@ const Node = struct {
221/// the count always starts at 1.248/// the count always starts at 1.
222node_count: usize = 1,249node_count: usize = 1,
223/// The root node of the trie.250/// The root node of the trie.
224root: Node = .{},251root: ?Node = null,
252allocator: *Allocator,
253
254pub fn init(allocator: *Allocator) Trie {
255 return .{ .allocator = allocator };
256}
225257
226/// Insert a symbol into the trie, updating the prefixes in the process.258/// Insert a symbol into the trie, updating the prefixes in the process.
227/// This operation may change the layout of the trie by splicing edges in259/// This operation may change the layout of the trie by splicing edges in
228/// certain circumstances.260/// certain circumstances.
229pub fn put(self: *Trie, alloc: *Allocator, symbol: Symbol) !void {261pub fn put(self: *Trie, symbol: Symbol) !void {
230 const res = try self.root.put(alloc, symbol.name, 0);262 if (self.root == null) {
231 self.node_count += res.node_count;263 self.root = .{ .base = self };
232 res.node.vmaddr_offset = symbol.vmaddr_offset;264 }
233 res.node.export_flags = symbol.export_flags;265 const node = try self.root.?.put(symbol.name);
266 node.vmaddr_offset = symbol.vmaddr_offset;
267 node.export_flags = symbol.export_flags;
234}268}
235269
236/// Write the trie to a buffer ULEB128 encoded.270const FromByteStreamError = error{
237pub fn writeULEB128Mem(self: *Trie, alloc: *Allocator, buffer: *std.ArrayListUnmanaged(u8)) !void {271 OutOfMemory,
238 var ordered_nodes: std.ArrayListUnmanaged(*Node) = .{};272 EndOfStream,
239 defer ordered_nodes.deinit(alloc);273 Overflow,
274};
240275
241 try ordered_nodes.ensureCapacity(alloc, self.node_count);276/// Parse the trie from a byte stream.
242 walkInOrder(&self.root, &ordered_nodes);277pub fn fromByteStream(self: *Trie, stream: anytype) FromByteStreamError!void {
278 if (self.root == null) {
279 self.root = .{ .base = self };
280 }
281 return self.root.?.fromByteStream(stream);
282}
283
284/// Write the trie to a buffer ULEB128 encoded.
285/// Caller owns the memory and needs to free it.
286pub fn writeULEB128Mem(self: *Trie) ![]u8 {
287 var ordered_nodes = try self.nodes();
288 defer self.allocator.free(ordered_nodes);
243289
244 var offset: usize = 0;290 var offset: usize = 0;
245 var more: bool = true;291 var more: bool = true;
246 while (more) {292 while (more) {
247 offset = 0;293 offset = 0;
248 more = false;294 more = false;
249 for (ordered_nodes.items) |node| {295 for (ordered_nodes) |node| {
250 const res = node.updateOffset(offset);296 const res = node.updateOffset(offset);
251 offset += res.node_size;297 offset += res.node_size;
252 if (res.updated) more = true;298 if (res.updated) more = true;
253 }299 }
254 }300 }
255301
256 try buffer.ensureCapacity(alloc, buffer.items.len + offset);302 var buffer = std.ArrayList(u8).init(self.allocator);
257 for (ordered_nodes.items) |node| {303 try buffer.ensureCapacity(offset);
258 try node.writeULEB128Mem(buffer);304 for (ordered_nodes) |node| {
305 try node.writeULEB128Mem(&buffer);
259 }306 }
307 return buffer.toOwnedSlice();
260}308}
261309
262/// Walks the trie in DFS order gathering all nodes into a linear stream of nodes.310pub fn nodes(self: *Trie) ![]*Node {
263fn walkInOrder(node: *Node, list: *std.ArrayListUnmanaged(*Node)) void {311 var ordered_nodes = std.ArrayList(*Node).init(self.allocator);
264 list.appendAssumeCapacity(node);312 try ordered_nodes.ensureCapacity(self.node_count);
265 for (node.edges.items) |*edge| {313
266 walkInOrder(edge.to, list);314 comptime const Fifo = std.fifo.LinearFifo(*Node, .{ .Static = std.math.maxInt(u8) });
315 var fifo = Fifo.init();
316 try fifo.writeItem(&self.root.?);
317
318 while (fifo.readItem()) |next| {
319 for (next.edges.items) |*edge| {
320 try fifo.writeItem(edge.to);
321 }
322 ordered_nodes.appendAssumeCapacity(next);
267 }323 }
324
325 return ordered_nodes.toOwnedSlice();
268}326}
269327
270pub fn deinit(self: *Trie, alloc: *Allocator) void {328pub fn deinit(self: *Trie) void {
271 self.root.deinit(alloc);329 self.root.?.deinit();
272}330}
273331
274test "Trie node count" {332test "Trie node count" {
275 var gpa = testing.allocator;333 var gpa = testing.allocator;
276 var trie: Trie = .{};334 var trie = Trie.init(gpa);
277 defer trie.deinit(gpa);335 defer trie.deinit();
278336
279 testing.expectEqual(trie.node_count, 1);337 testing.expectEqual(trie.node_count, 1);
280338
281 try trie.put(gpa, .{339 try trie.put(.{
282 .name = "_main",340 .name = "_main",
283 .vmaddr_offset = 0,341 .vmaddr_offset = 0,
284 .export_flags = 0,342 .export_flags = 0,
...@@ -286,14 +344,14 @@ test "Trie node count" {...@@ -286,14 +344,14 @@ test "Trie node count" {
286 testing.expectEqual(trie.node_count, 2);344 testing.expectEqual(trie.node_count, 2);
287345
288 // Inserting the same node shouldn't update the trie.346 // Inserting the same node shouldn't update the trie.
289 try trie.put(gpa, .{347 try trie.put(.{
290 .name = "_main",348 .name = "_main",
291 .vmaddr_offset = 0,349 .vmaddr_offset = 0,
292 .export_flags = 0,350 .export_flags = 0,
293 });351 });
294 testing.expectEqual(trie.node_count, 2);352 testing.expectEqual(trie.node_count, 2);
295353
296 try trie.put(gpa, .{354 try trie.put(.{
297 .name = "__mh_execute_header",355 .name = "__mh_execute_header",
298 .vmaddr_offset = 0x1000,356 .vmaddr_offset = 0x1000,
299 .export_flags = 0,357 .export_flags = 0,
...@@ -301,13 +359,13 @@ test "Trie node count" {...@@ -301,13 +359,13 @@ test "Trie node count" {
301 testing.expectEqual(trie.node_count, 4);359 testing.expectEqual(trie.node_count, 4);
302360
303 // Inserting the same node shouldn't update the trie.361 // Inserting the same node shouldn't update the trie.
304 try trie.put(gpa, .{362 try trie.put(.{
305 .name = "__mh_execute_header",363 .name = "__mh_execute_header",
306 .vmaddr_offset = 0x1000,364 .vmaddr_offset = 0x1000,
307 .export_flags = 0,365 .export_flags = 0,
308 });366 });
309 testing.expectEqual(trie.node_count, 4);367 testing.expectEqual(trie.node_count, 4);
310 try trie.put(gpa, .{368 try trie.put(.{
311 .name = "_main",369 .name = "_main",
312 .vmaddr_offset = 0,370 .vmaddr_offset = 0,
313 .export_flags = 0,371 .export_flags = 0,
...@@ -317,31 +375,28 @@ test "Trie node count" {...@@ -317,31 +375,28 @@ test "Trie node count" {
317375
318test "Trie basic" {376test "Trie basic" {
319 var gpa = testing.allocator;377 var gpa = testing.allocator;
320 var trie: Trie = .{};378 var trie = Trie.init(gpa);
321 defer trie.deinit(gpa);379 defer trie.deinit();
322
323 // root
324 testing.expect(trie.root.edges.items.len == 0);
325380
326 // root --- _st ---> node381 // root --- _st ---> node
327 try trie.put(gpa, .{382 try trie.put(.{
328 .name = "_st",383 .name = "_st",
329 .vmaddr_offset = 0,384 .vmaddr_offset = 0,
330 .export_flags = 0,385 .export_flags = 0,
331 });386 });
332 testing.expect(trie.root.edges.items.len == 1);387 testing.expect(trie.root.?.edges.items.len == 1);
333 testing.expect(mem.eql(u8, trie.root.edges.items[0].label, "_st"));388 testing.expect(mem.eql(u8, trie.root.?.edges.items[0].label, "_st"));
334389
335 {390 {
336 // root --- _st ---> node --- art ---> node391 // root --- _st ---> node --- art ---> node
337 try trie.put(gpa, .{392 try trie.put(.{
338 .name = "_start",393 .name = "_start",
339 .vmaddr_offset = 0,394 .vmaddr_offset = 0,
340 .export_flags = 0,395 .export_flags = 0,
341 });396 });
342 testing.expect(trie.root.edges.items.len == 1);397 testing.expect(trie.root.?.edges.items.len == 1);
343398
344 const nextEdge = &trie.root.edges.items[0];399 const nextEdge = &trie.root.?.edges.items[0];
345 testing.expect(mem.eql(u8, nextEdge.label, "_st"));400 testing.expect(mem.eql(u8, nextEdge.label, "_st"));
346 testing.expect(nextEdge.to.edges.items.len == 1);401 testing.expect(nextEdge.to.edges.items.len == 1);
347 testing.expect(mem.eql(u8, nextEdge.to.edges.items[0].label, "art"));402 testing.expect(mem.eql(u8, nextEdge.to.edges.items[0].label, "art"));
...@@ -350,14 +405,14 @@ test "Trie basic" {...@@ -350,14 +405,14 @@ test "Trie basic" {
350 // root --- _ ---> node --- st ---> node --- art ---> node405 // root --- _ ---> node --- st ---> node --- art ---> node
351 // |406 // |
352 // | --- main ---> node407 // | --- main ---> node
353 try trie.put(gpa, .{408 try trie.put(.{
354 .name = "_main",409 .name = "_main",
355 .vmaddr_offset = 0,410 .vmaddr_offset = 0,
356 .export_flags = 0,411 .export_flags = 0,
357 });412 });
358 testing.expect(trie.root.edges.items.len == 1);413 testing.expect(trie.root.?.edges.items.len == 1);
359414
360 const nextEdge = &trie.root.edges.items[0];415 const nextEdge = &trie.root.?.edges.items[0];
361 testing.expect(mem.eql(u8, nextEdge.label, "_"));416 testing.expect(mem.eql(u8, nextEdge.label, "_"));
362 testing.expect(nextEdge.to.edges.items.len == 2);417 testing.expect(nextEdge.to.edges.items.len == 2);
363 testing.expect(mem.eql(u8, nextEdge.to.edges.items[0].label, "st"));418 testing.expect(mem.eql(u8, nextEdge.to.edges.items[0].label, "st"));
...@@ -370,24 +425,22 @@ test "Trie basic" {...@@ -370,24 +425,22 @@ test "Trie basic" {
370425
371test "Trie.writeULEB128Mem" {426test "Trie.writeULEB128Mem" {
372 var gpa = testing.allocator;427 var gpa = testing.allocator;
373 var trie: Trie = .{};428 var trie = Trie.init(gpa);
374 defer trie.deinit(gpa);429 defer trie.deinit();
375430
376 try trie.put(gpa, .{431 try trie.put(.{
377 .name = "__mh_execute_header",432 .name = "__mh_execute_header",
378 .vmaddr_offset = 0,433 .vmaddr_offset = 0,
379 .export_flags = 0,434 .export_flags = 0,
380 });435 });
381 try trie.put(gpa, .{436 try trie.put(.{
382 .name = "_main",437 .name = "_main",
383 .vmaddr_offset = 0x1000,438 .vmaddr_offset = 0x1000,
384 .export_flags = 0,439 .export_flags = 0,
385 });440 });
386441
387 var buffer: std.ArrayListUnmanaged(u8) = .{};442 var buffer = try trie.writeULEB128Mem();
388 defer buffer.deinit(gpa);443 defer gpa.free(buffer);
389
390 try trie.writeULEB128Mem(gpa, &buffer);
391444
392 const exp_buffer = [_]u8{445 const exp_buffer = [_]u8{
393 0x0,446 0x0,
...@@ -434,6 +487,64 @@ test "Trie.writeULEB128Mem" {...@@ -434,6 +487,64 @@ test "Trie.writeULEB128Mem" {
434 0x0,487 0x0,
435 };488 };
436489
437 testing.expect(buffer.items.len == exp_buffer.len);490 testing.expect(buffer.len == exp_buffer.len);
438 testing.expect(mem.eql(u8, buffer.items, exp_buffer[0..]));491 testing.expect(mem.eql(u8, buffer, exp_buffer[0..]));
492}
493
494test "parse Trie from byte stream" {
495 var gpa = testing.allocator;
496
497 const in_buffer = [_]u8{
498 0x0,
499 0x1,
500 0x5f,
501 0x0,
502 0x5,
503 0x0,
504 0x2,
505 0x5f,
506 0x6d,
507 0x68,
508 0x5f,
509 0x65,
510 0x78,
511 0x65,
512 0x63,
513 0x75,
514 0x74,
515 0x65,
516 0x5f,
517 0x68,
518 0x65,
519 0x61,
520 0x64,
521 0x65,
522 0x72,
523 0x0,
524 0x21,
525 0x6d,
526 0x61,
527 0x69,
528 0x6e,
529 0x0,
530 0x25,
531 0x2,
532 0x0,
533 0x0,
534 0x0,
535 0x3,
536 0x0,
537 0x80,
538 0x20,
539 0x0,
540 };
541 var stream = std.io.fixedBufferStream(in_buffer[0..]);
542 var trie = Trie.init(gpa);
543 defer trie.deinit();
544 try trie.fromByteStream(&stream);
545
546 var out_buffer = try trie.writeULEB128Mem();
547 defer gpa.free(out_buffer);
548
549 testing.expect(mem.eql(u8, in_buffer[0..], out_buffer));
439}550}