authorgravatar for kubkon@jakubkonka.comJakub Konka <kubkon@jakubkonka.com> 2020-12-10 08:32:59+01:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2020-12-10 08:32:59+01:00
log23c1b7faee13f95bce6ba48051b9ac1a1fab9576
tree998dee2967d421e2368c61542b2cf3d153b1a5a6
parent8951f72fa39de019683b233f45a3b45d83b57904
parent44e2f210bb3fee78d339bc5c75822e0311600f70
signaturebadge-question-mark Signed by PGP key 4AEE18F83AFDEB23

Merge pull request #7368 from kubkon/macho-trie-cleanup

stage2: MachO export trie cleanup

3 files changed, 418 insertions(+), 222 deletions(-)

lib/std/macho.zig+9
......@@ -1333,6 +1333,15 @@ pub const N_WEAK_DEF: u16 = 0x80;
13331333/// This bit is only available in .o files (MH_OBJECT filetype)
13341334pub 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
13361345// Codesign consts and structs taken from:
13371346// https://opensource.apple.com/source/xnu/xnu-6153.81.5/osfmk/kern/cs_blobs.h.auto.html
13381347
src/link/MachO.zig+73-31
......@@ -301,7 +301,10 @@ pub fn flushModule(self: *MachO, comp: *Compilation) !void {
301301 const tracy = trace(@src());
302302 defer tracy.end();
303303
304 switch (self.base.options.output_mode) {
304 const output_mode = self.base.options.output_mode;
305 const target = self.base.options.target;
306
307 switch (output_mode) {
305308 .Exe => {
306309 if (self.entry_addr) |addr| {
307310 // Update LC_MAIN with entry offset.
......@@ -312,12 +315,15 @@ pub fn flushModule(self: *MachO, comp: *Compilation) !void {
312315 try self.writeExportTrie();
313316 try self.writeSymbolTable();
314317 try self.writeStringTable();
315 // Preallocate space for the code signature.
316 // We need to do this at this stage so that we have the load commands with proper values
317 // written out to the file.
318 // The most important here is to have the correct vm and filesize of the __LINKEDIT segment
319 // where the code signature goes into.
320 try self.writeCodeSignaturePadding();
318
319 if (target.cpu.arch == .aarch64) {
320 // Preallocate space for the code signature.
321 // We need to do this at this stage so that we have the load commands with proper values
322 // written out to the file.
323 // The most important here is to have the correct vm and filesize of the __LINKEDIT segment
324 // where the code signature goes into.
325 try self.writeCodeSignaturePadding();
326 }
321327 },
322328 .Obj => {},
323329 .Lib => return error.TODOImplementWritingLibFiles,
......@@ -339,9 +345,11 @@ pub fn flushModule(self: *MachO, comp: *Compilation) !void {
339345
340346 assert(!self.cmd_table_dirty);
341347
342 switch (self.base.options.output_mode) {
343 .Exe, .Lib => try self.writeCodeSignature(), // code signing always comes last
344 else => {},
348 if (target.cpu.arch == .aarch64) {
349 switch (output_mode) {
350 .Exe, .Lib => try self.writeCodeSignature(), // code signing always comes last
351 else => {},
352 }
345353 }
346354}
347355
......@@ -752,17 +760,15 @@ fn linkWithLLD(self: *MachO, comp: *Compilation) !void {
752760 const text_segment = self.load_commands.items[self.text_segment_cmd_index.?].Segment;
753761 const text_section = text_segment.sections.items[self.text_section_index.?];
754762 const after_last_cmd_offset = self.header.?.sizeofcmds + @sizeOf(macho.mach_header_64);
755 const needed_size = @sizeOf(macho.linkedit_data_command);
763 const needed_size = @sizeOf(macho.linkedit_data_command) * alloc_num / alloc_den;
764
756765 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 sections
758 // by the required offset, but this requires a little bit more thinking and bookkeeping.
759 // For now, return an error informing the user of the problem.
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});
766 std.log.err("Unable to extend padding between the end of load commands and start of __text section.", .{});
767 std.log.err("Re-run the linker with '-headerpad 0x{x}' option if available, or", .{needed_size});
768 std.log.err("fall back to the system linker by exporting 'ZIG_SYSTEM_LINKER_HACK=1'.", .{});
764769 return error.NotEnoughPadding;
765770 }
771
766772 const linkedit_segment = self.load_commands.items[self.linkedit_segment_cmd_index.?].Segment;
767773 // TODO This is clunky.
768774 self.linkedit_segment_next_offset = @intCast(u32, mem.alignForwardGeneric(u64, linkedit_segment.inner.fileoff + linkedit_segment.inner.filesize, @sizeOf(u64)));
......@@ -1799,38 +1805,40 @@ fn writeCodeSignature(self: *MachO) !void {
17991805fn writeExportTrie(self: *MachO) !void {
18001806 if (self.global_symbols.items.len == 0) return;
18011807
1802 var trie: Trie = .{};
1803 defer trie.deinit(self.base.allocator);
1808 var trie = Trie.init(self.base.allocator);
1809 defer trie.deinit();
18041810
18051811 const text_segment = self.load_commands.items[self.text_segment_cmd_index.?].Segment;
18061812 for (self.global_symbols.items) |symbol| {
18071813 // TODO figure out if we should put all global symbols into the export trie
18081814 const name = self.getString(symbol.n_strx);
18091815 assert(symbol.n_value >= text_segment.inner.vmaddr);
1810 try trie.put(self.base.allocator, .{
1816 try trie.put(.{
18111817 .name = name,
18121818 .vmaddr_offset = symbol.n_value - text_segment.inner.vmaddr,
1813 .export_flags = 0, // TODO workout creation of export flags
1819 .export_flags = macho.EXPORT_SYMBOL_FLAGS_KIND_REGULAR,
18141820 });
18151821 }
18161822
1817 var buffer: std.ArrayListUnmanaged(u8) = .{};
1818 defer buffer.deinit(self.base.allocator);
1819
1820 try trie.writeULEB128Mem(self.base.allocator, &buffer);
1823 try trie.finalize();
1824 var buffer = try self.base.allocator.alloc(u8, trie.size);
1825 defer self.base.allocator.free(buffer);
1826 var stream = std.io.fixedBufferStream(buffer);
1827 const nwritten = try trie.write(stream.writer());
1828 assert(nwritten == trie.size);
18211829
18221830 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)));
1831 const export_size = @intCast(u32, mem.alignForward(buffer.len, @sizeOf(u64)));
18241832 dyld_info.export_off = self.linkedit_segment_next_offset.?;
18251833 dyld_info.export_size = export_size;
18261834
18271835 log.debug("writing export trie from 0x{x} to 0x{x}\n", .{ dyld_info.export_off, dyld_info.export_off + export_size });
18281836
1829 if (export_size > buffer.items.len) {
1837 if (export_size > buffer.len) {
18301838 // Pad out to align(8).
18311839 try self.base.file.?.pwriteAll(&[_]u8{0}, dyld_info.export_off + export_size);
18321840 }
1833 try self.base.file.?.pwriteAll(buffer.items, dyld_info.export_off);
1841 try self.base.file.?.pwriteAll(buffer, dyld_info.export_off);
18341842
18351843 self.linkedit_segment_next_offset = dyld_info.export_off + dyld_info.export_size;
18361844 // Advance size of __LINKEDIT segment
......@@ -1917,7 +1925,9 @@ fn parseFromFile(self: *MachO, file: fs.File) !void {
19171925 switch (cmd.cmd()) {
19181926 macho.LC_SEGMENT_64 => {
19191927 const x = cmd.Segment;
1920 if (isSegmentOrSection(&x.inner.segname, "__LINKEDIT")) {
1928 if (isSegmentOrSection(&x.inner.segname, "__PAGEZERO")) {
1929 self.pagezero_segment_cmd_index = i;
1930 } else if (isSegmentOrSection(&x.inner.segname, "__LINKEDIT")) {
19211931 self.linkedit_segment_cmd_index = i;
19221932 } else if (isSegmentOrSection(&x.inner.segname, "__TEXT")) {
19231933 self.text_segment_cmd_index = i;
......@@ -1926,16 +1936,48 @@ fn parseFromFile(self: *MachO, file: fs.File) !void {
19261936 self.text_section_index = @intCast(u16, j);
19271937 }
19281938 }
1939 } else if (isSegmentOrSection(&x.inner.segname, "__DATA")) {
1940 self.data_segment_cmd_index = i;
19291941 }
19301942 },
1943 macho.LC_DYLD_INFO_ONLY => {
1944 self.dyld_info_cmd_index = i;
1945 },
19311946 macho.LC_SYMTAB => {
19321947 self.symtab_cmd_index = i;
19331948 },
1949 macho.LC_DYSYMTAB => {
1950 self.dysymtab_cmd_index = i;
1951 },
1952 macho.LC_LOAD_DYLINKER => {
1953 self.dylinker_cmd_index = i;
1954 },
1955 macho.LC_VERSION_MIN_MACOSX, macho.LC_VERSION_MIN_IPHONEOS, macho.LC_VERSION_MIN_WATCHOS, macho.LC_VERSION_MIN_TVOS => {
1956 self.version_min_cmd_index = i;
1957 },
1958 macho.LC_SOURCE_VERSION => {
1959 self.source_version_cmd_index = i;
1960 },
1961 macho.LC_MAIN => {
1962 self.main_cmd_index = i;
1963 },
1964 macho.LC_LOAD_DYLIB => {
1965 self.libsystem_cmd_index = i; // TODO This is incorrect, but we'll fixup later.
1966 },
1967 macho.LC_FUNCTION_STARTS => {
1968 self.function_starts_cmd_index = i;
1969 },
1970 macho.LC_DATA_IN_CODE => {
1971 self.data_in_code_cmd_index = i;
1972 },
19341973 macho.LC_CODE_SIGNATURE => {
19351974 self.code_signature_cmd_index = i;
19361975 },
19371976 // TODO populate more MachO fields
1938 else => {},
1977 else => {
1978 std.log.err("Unknown load command detected: 0x{x}.", .{cmd.cmd()});
1979 return error.UnknownLoadCommand;
1980 },
19391981 }
19401982 self.load_commands.appendAssumeCapacity(cmd);
19411983 }
src/link/MachO/Trie.zig+336-191
......@@ -34,156 +34,218 @@ const std = @import("std");
3434const mem = std.mem;
3535const leb = std.leb;
3636const log = std.log.scoped(.link);
37const macho = std.macho;
3738const testing = std.testing;
3839const assert = std.debug.assert;
3940const Allocator = mem.Allocator;
4041
41pub const Symbol = struct {
42 name: []const u8,
43 vmaddr_offset: u64,
44 export_flags: u64,
45};
46
47const Edge = struct {
48 from: *Node,
49 to: *Node,
50 label: []const u8,
42pub const Node = struct {
43 base: *Trie,
5144
52 fn deinit(self: *Edge, alloc: *Allocator) void {
53 self.to.deinit(alloc);
54 alloc.destroy(self.to);
55 self.from = undefined;
56 self.to = undefined;
57 }
58};
45 /// Terminal info associated with this node.
46 /// If this node is not a terminal node, info is null.
47 terminal_info: ?struct {
48 /// Export flags associated with this exported symbol.
49 export_flags: u64,
50 /// VM address offset wrt to the section this symbol is defined against.
51 vmaddr_offset: u64,
52 } = null,
5953
60const Node = struct {
61 /// Export flags associated with this exported symbol (if any).
62 export_flags: ?u64 = null,
63 /// VM address offset wrt to the section this symbol is defined against (if any).
64 vmaddr_offset: ?u64 = null,
6554 /// Offset of this node in the trie output byte stream.
6655 trie_offset: ?usize = null,
56
6757 /// List of all edges originating from this node.
6858 edges: std.ArrayListUnmanaged(Edge) = .{},
6959
70 fn deinit(self: *Node, alloc: *Allocator) void {
60 node_dirty: bool = true,
61
62 /// Edge connecting to nodes in the trie.
63 pub const Edge = struct {
64 from: *Node,
65 to: *Node,
66 label: []u8,
67
68 fn deinit(self: *Edge, allocator: *Allocator) void {
69 self.to.deinit(allocator);
70 allocator.destroy(self.to);
71 allocator.free(self.label);
72 self.from = undefined;
73 self.to = undefined;
74 self.label = undefined;
75 }
76 };
77
78 fn deinit(self: *Node, allocator: *Allocator) void {
7179 for (self.edges.items) |*edge| {
72 edge.deinit(alloc);
80 edge.deinit(allocator);
7381 }
74 self.edges.deinit(alloc);
82 self.edges.deinit(allocator);
7583 }
7684
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
8485 /// Inserts a new node starting from `self`.
85 fn put(self: *Node, alloc: *Allocator, label: []const u8, node_count: usize) !PutResult {
86 var curr_node_count = node_count;
86 fn put(self: *Node, allocator: *Allocator, label: []const u8) !*Node {
8787 // Check for match with edges from this node.
8888 for (self.edges.items) |*edge| {
89 const match = mem.indexOfDiff(u8, edge.label, label) orelse return PutResult{
90 .node = edge.to,
91 .node_count = curr_node_count,
92 };
89 const match = mem.indexOfDiff(u8, edge.label, label) orelse return edge.to;
9390 if (match == 0) continue;
94 if (match == edge.label.len) return edge.to.put(alloc, label[match..], curr_node_count);
91 if (match == edge.label.len) return edge.to.put(allocator, label[match..]);
9592
9693 // Found a match, need to splice up nodes.
9794 // From: A -> B
9895 // To: A -> C -> B
99 const mid = try alloc.create(Node);
100 mid.* = .{};
101 const to_label = edge.label;
96 const mid = try allocator.create(Node);
97 mid.* = .{ .base = self.base };
98 var to_label = try allocator.dupe(u8, edge.label[match..]);
99 allocator.free(edge.label);
102100 const to_node = edge.to;
103101 edge.to = mid;
104 edge.label = label[0..match];
105 curr_node_count += 1;
102 edge.label = try allocator.dupe(u8, label[0..match]);
103 self.base.node_count += 1;
106104
107 try mid.edges.append(alloc, .{
105 try mid.edges.append(allocator, .{
108106 .from = mid,
109107 .to = to_node,
110 .label = to_label[match..],
108 .label = to_label,
111109 });
112110
113 if (match == label.len) {
114 return PutResult{ .node = to_node, .node_count = curr_node_count };
115 } else {
116 return mid.put(alloc, label[match..], curr_node_count);
117 }
111 return if (match == label.len) to_node else mid.put(allocator, label[match..]);
118112 }
119113
120114 // Add a new node.
121 const node = try alloc.create(Node);
122 node.* = .{};
123 curr_node_count += 1;
115 const node = try allocator.create(Node);
116 node.* = .{ .base = self.base };
117 self.base.node_count += 1;
124118
125 try self.edges.append(alloc, .{
119 try self.edges.append(allocator, .{
126120 .from = self,
127121 .to = node,
128 .label = label,
122 .label = try allocator.dupe(u8, label),
129123 });
130124
131 return PutResult{ .node = node, .node_count = curr_node_count };
125 return node;
132126 }
133127
134 /// This method should only be called *after* updateOffset has been called!
135 /// In case this is not upheld, this method will panic.
136 fn writeULEB128Mem(self: Node, buffer: *std.ArrayListUnmanaged(u8)) !void {
137 assert(self.trie_offset != null); // You need to call updateOffset first.
138 if (self.vmaddr_offset) |offset| {
128 /// Recursively parses the node from the input byte stream.
129 fn read(self: *Node, allocator: *Allocator, reader: anytype) Trie.ReadError!usize {
130 self.node_dirty = true;
131 const trie_offset = try reader.context.getPos();
132 self.trie_offset = trie_offset;
133
134 var nread: usize = 0;
135
136 const node_size = try leb.readULEB128(u64, reader);
137 if (node_size > 0) {
138 const export_flags = try leb.readULEB128(u64, reader);
139 // TODO Parse special flags.
140 assert(export_flags & macho.EXPORT_SYMBOL_FLAGS_REEXPORT == 0 and
141 export_flags & macho.EXPORT_SYMBOL_FLAGS_STUB_AND_RESOLVER == 0);
142
143 const vmaddr_offset = try leb.readULEB128(u64, reader);
144
145 self.terminal_info = .{
146 .export_flags = export_flags,
147 .vmaddr_offset = vmaddr_offset,
148 };
149 }
150
151 const nedges = try reader.readByte();
152 self.base.node_count += nedges;
153
154 nread += (try reader.context.getPos()) - trie_offset;
155
156 var i: usize = 0;
157 while (i < nedges) : (i += 1) {
158 const edge_start_pos = try reader.context.getPos();
159
160 const label = blk: {
161 var label_buf = std.ArrayList(u8).init(allocator);
162 while (true) {
163 const next = try reader.readByte();
164 if (next == @as(u8, 0))
165 break;
166 try label_buf.append(next);
167 }
168 break :blk label_buf.toOwnedSlice();
169 };
170
171 const seek_to = try leb.readULEB128(u64, reader);
172 const return_pos = try reader.context.getPos();
173
174 nread += return_pos - edge_start_pos;
175 try reader.context.seekTo(seek_to);
176
177 const node = try allocator.create(Node);
178 node.* = .{ .base = self.base };
179
180 nread += try node.read(allocator, reader);
181 try self.edges.append(allocator, .{
182 .from = self,
183 .to = node,
184 .label = label,
185 });
186 try reader.context.seekTo(return_pos);
187 }
188
189 return nread;
190 }
191
192 /// Writes this node to a byte stream.
193 /// The children of this node *are* not written to the byte stream
194 /// recursively. To write all nodes to a byte stream in sequence,
195 /// iterate over `Trie.ordered_nodes` and call this method on each node.
196 /// This is one of the requirements of the MachO.
197 /// Panics if `finalize` was not called before calling this method.
198 fn write(self: Node, writer: anytype) !void {
199 assert(!self.node_dirty);
200 if (self.terminal_info) |info| {
139201 // Terminal node info: encode export flags and vmaddr offset of this symbol.
140202 var info_buf_len: usize = 0;
141203 var info_buf: [@sizeOf(u64) * 2]u8 = undefined;
142204 var info_stream = std.io.fixedBufferStream(&info_buf);
143 try leb.writeULEB128(info_stream.writer(), self.export_flags.?);
144 try leb.writeULEB128(info_stream.writer(), offset);
205 // TODO Implement for special flags.
206 assert(info.export_flags & macho.EXPORT_SYMBOL_FLAGS_REEXPORT == 0 and
207 info.export_flags & macho.EXPORT_SYMBOL_FLAGS_STUB_AND_RESOLVER == 0);
208 try leb.writeULEB128(info_stream.writer(), info.export_flags);
209 try leb.writeULEB128(info_stream.writer(), info.vmaddr_offset);
145210
146211 // Encode the size of the terminal node info.
147212 var size_buf: [@sizeOf(u64)]u8 = undefined;
148213 var size_stream = std.io.fixedBufferStream(&size_buf);
149214 try leb.writeULEB128(size_stream.writer(), info_stream.pos);
150215
151 // Now, write them to the output buffer.
152 buffer.appendSliceAssumeCapacity(size_buf[0..size_stream.pos]);
153 buffer.appendSliceAssumeCapacity(info_buf[0..info_stream.pos]);
216 // Now, write them to the output stream.
217 try writer.writeAll(size_buf[0..size_stream.pos]);
218 try writer.writeAll(info_buf[0..info_stream.pos]);
154219 } else {
155220 // Non-terminal node is delimited by 0 byte.
156 buffer.appendAssumeCapacity(0);
221 try writer.writeByte(0);
157222 }
158223 // Write number of edges (max legal number of edges is 256).
159 buffer.appendAssumeCapacity(@intCast(u8, self.edges.items.len));
224 try writer.writeByte(@intCast(u8, self.edges.items.len));
160225
161226 for (self.edges.items) |edge| {
162 // Write edges labels.
163 buffer.appendSliceAssumeCapacity(edge.label);
164 buffer.appendAssumeCapacity(0);
165
166 var buf: [@sizeOf(u64)]u8 = undefined;
167 var buf_stream = std.io.fixedBufferStream(&buf);
168 try leb.writeULEB128(buf_stream.writer(), edge.to.trie_offset.?);
169 buffer.appendSliceAssumeCapacity(buf[0..buf_stream.pos]);
227 // Write edge label and offset to next node in trie.
228 try writer.writeAll(edge.label);
229 try writer.writeByte(0);
230 try leb.writeULEB128(writer, edge.to.trie_offset.?);
170231 }
171232 }
172233
173 const UpdateResult = struct {
234 const FinalizeResult = struct {
174235 /// Current size of this node in bytes.
175236 node_size: usize,
237
176238 /// True if the trie offset of this node in the output byte stream
177239 /// would need updating; false otherwise.
178240 updated: bool,
179241 };
180242
181243 /// Updates offset of this node in the output byte stream.
182 fn updateOffset(self: *Node, offset: usize) UpdateResult {
244 fn finalize(self: *Node, offset_in_trie: usize) FinalizeResult {
183245 var node_size: usize = 0;
184 if (self.vmaddr_offset) |vmaddr| {
185 node_size += sizeULEB128Mem(self.export_flags.?);
186 node_size += sizeULEB128Mem(vmaddr);
246 if (self.terminal_info) |info| {
247 node_size += sizeULEB128Mem(info.export_flags);
248 node_size += sizeULEB128Mem(info.vmaddr_offset);
187249 node_size += sizeULEB128Mem(node_size);
188250 } else {
189251 node_size += 1; // 0x0 for non-terminal nodes
......@@ -196,8 +258,9 @@ const Node = struct {
196258 }
197259
198260 const trie_offset = self.trie_offset orelse 0;
199 const updated = offset != trie_offset;
200 self.trie_offset = offset;
261 const updated = offset_in_trie != trie_offset;
262 self.trie_offset = offset_in_trie;
263 self.node_dirty = false;
201264
202265 return .{ .node_size = node_size, .updated = updated };
203266 }
......@@ -215,70 +278,146 @@ const Node = struct {
215278 }
216279};
217280
218/// Count of nodes in the trie.
219/// The count is updated at every `put` call.
220/// The trie always consists of at least a root node, hence
221/// the count always starts at 1.
222node_count: usize = 1,
223281/// The root node of the trie.
224root: Node = .{},
282root: ?*Node = null,
283
284allocator: *Allocator,
285
286/// If you want to access nodes ordered in DFS fashion,
287/// you should call `finalize` first since the nodes
288/// in this container are not guaranteed to not be stale
289/// if more insertions took place after the last `finalize`
290/// call.
291ordered_nodes: std.ArrayListUnmanaged(*Node) = .{},
292
293/// The size of the trie in bytes.
294/// This value may be outdated if there were additional
295/// insertions performed after `finalize` was called.
296/// Call `finalize` before accessing this value to ensure
297/// it is up-to-date.
298size: usize = 0,
299
300/// Number of nodes currently in the trie.
301node_count: usize = 0,
302
303trie_dirty: bool = true,
304
305pub fn init(allocator: *Allocator) Trie {
306 return .{ .allocator = allocator };
307}
308
309/// Export symbol that is to be placed in the trie.
310pub const ExportSymbol = struct {
311 /// Name of the symbol.
312 name: []const u8,
313
314 /// Offset of this symbol's virtual memory address from the beginning
315 /// of the __TEXT segment.
316 vmaddr_offset: u64,
317
318 /// Export flags of this exported symbol.
319 export_flags: u64,
320};
225321
226322/// Insert a symbol into the trie, updating the prefixes in the process.
227323/// This operation may change the layout of the trie by splicing edges in
228324/// certain circumstances.
229pub fn put(self: *Trie, alloc: *Allocator, symbol: Symbol) !void {
230 const res = try self.root.put(alloc, symbol.name, 0);
231 self.node_count += res.node_count;
232 res.node.vmaddr_offset = symbol.vmaddr_offset;
233 res.node.export_flags = symbol.export_flags;
325pub fn put(self: *Trie, symbol: ExportSymbol) !void {
326 try self.createRoot();
327 const node = try self.root.?.put(self.allocator, symbol.name);
328 node.terminal_info = .{
329 .vmaddr_offset = symbol.vmaddr_offset,
330 .export_flags = symbol.export_flags,
331 };
332 self.trie_dirty = true;
234333}
235334
236/// Write the trie to a buffer ULEB128 encoded.
237pub fn writeULEB128Mem(self: *Trie, alloc: *Allocator, buffer: *std.ArrayListUnmanaged(u8)) !void {
238 var ordered_nodes: std.ArrayListUnmanaged(*Node) = .{};
239 defer ordered_nodes.deinit(alloc);
335/// Finalizes this trie for writing to a byte stream.
336/// This step performs multiple passes through the trie ensuring
337/// there are no gaps after every `Node` is ULEB128 encoded.
338/// Call this method before trying to `write` the trie to a byte stream.
339pub fn finalize(self: *Trie) !void {
340 if (!self.trie_dirty) return;
341
342 self.ordered_nodes.shrinkRetainingCapacity(0);
343 try self.ordered_nodes.ensureCapacity(self.allocator, self.node_count);
344
345 comptime const Fifo = std.fifo.LinearFifo(*Node, .{ .Static = std.math.maxInt(u8) });
346 var fifo = Fifo.init();
347 try fifo.writeItem(self.root.?);
240348
241 try ordered_nodes.ensureCapacity(alloc, self.node_count);
242 walkInOrder(&self.root, &ordered_nodes);
349 while (fifo.readItem()) |next| {
350 for (next.edges.items) |*edge| {
351 try fifo.writeItem(edge.to);
352 }
353 self.ordered_nodes.appendAssumeCapacity(next);
354 }
243355
244 var offset: usize = 0;
245356 var more: bool = true;
246357 while (more) {
247 offset = 0;
358 self.size = 0;
248359 more = false;
249 for (ordered_nodes.items) |node| {
250 const res = node.updateOffset(offset);
251 offset += res.node_size;
360 for (self.ordered_nodes.items) |node| {
361 const res = node.finalize(self.size);
362 self.size += res.node_size;
252363 if (res.updated) more = true;
253364 }
254365 }
255366
256 try buffer.ensureCapacity(alloc, buffer.items.len + offset);
257 for (ordered_nodes.items) |node| {
258 try node.writeULEB128Mem(buffer);
367 self.trie_dirty = false;
368}
369
370const ReadError = error{
371 OutOfMemory,
372 EndOfStream,
373 Overflow,
374};
375
376/// Parse the trie from a byte stream.
377pub fn read(self: *Trie, reader: anytype) ReadError!usize {
378 try self.createRoot();
379 return self.root.?.read(self.allocator, reader);
380}
381
382/// Write the trie to a byte stream.
383/// Caller owns the memory and needs to free it.
384/// Panics if the trie was not finalized using `finalize`
385/// before calling this method.
386pub fn write(self: Trie, writer: anytype) !usize {
387 assert(!self.trie_dirty);
388 var counting_writer = std.io.countingWriter(writer);
389 for (self.ordered_nodes.items) |node| {
390 try node.write(counting_writer.writer());
259391 }
392 return counting_writer.bytes_written;
260393}
261394
262/// Walks the trie in DFS order gathering all nodes into a linear stream of nodes.
263fn walkInOrder(node: *Node, list: *std.ArrayListUnmanaged(*Node)) void {
264 list.appendAssumeCapacity(node);
265 for (node.edges.items) |*edge| {
266 walkInOrder(edge.to, list);
395pub fn deinit(self: *Trie) void {
396 if (self.root) |root| {
397 root.deinit(self.allocator);
398 self.allocator.destroy(root);
267399 }
400 self.ordered_nodes.deinit(self.allocator);
268401}
269402
270pub fn deinit(self: *Trie, alloc: *Allocator) void {
271 self.root.deinit(alloc);
403fn createRoot(self: *Trie) !void {
404 if (self.root == null) {
405 const root = try self.allocator.create(Node);
406 root.* = .{ .base = self };
407 self.root = root;
408 self.node_count += 1;
409 }
272410}
273411
274412test "Trie node count" {
275413 var gpa = testing.allocator;
276 var trie: Trie = .{};
277 defer trie.deinit(gpa);
414 var trie = Trie.init(gpa);
415 defer trie.deinit();
278416
279 testing.expectEqual(trie.node_count, 1);
417 testing.expectEqual(trie.node_count, 0);
418 testing.expect(trie.root == null);
280419
281 try trie.put(gpa, .{
420 try trie.put(.{
282421 .name = "_main",
283422 .vmaddr_offset = 0,
284423 .export_flags = 0,
......@@ -286,14 +425,14 @@ test "Trie node count" {
286425 testing.expectEqual(trie.node_count, 2);
287426
288427 // Inserting the same node shouldn't update the trie.
289 try trie.put(gpa, .{
428 try trie.put(.{
290429 .name = "_main",
291430 .vmaddr_offset = 0,
292431 .export_flags = 0,
293432 });
294433 testing.expectEqual(trie.node_count, 2);
295434
296 try trie.put(gpa, .{
435 try trie.put(.{
297436 .name = "__mh_execute_header",
298437 .vmaddr_offset = 0x1000,
299438 .export_flags = 0,
......@@ -301,13 +440,13 @@ test "Trie node count" {
301440 testing.expectEqual(trie.node_count, 4);
302441
303442 // Inserting the same node shouldn't update the trie.
304 try trie.put(gpa, .{
443 try trie.put(.{
305444 .name = "__mh_execute_header",
306445 .vmaddr_offset = 0x1000,
307446 .export_flags = 0,
308447 });
309448 testing.expectEqual(trie.node_count, 4);
310 try trie.put(gpa, .{
449 try trie.put(.{
311450 .name = "_main",
312451 .vmaddr_offset = 0,
313452 .export_flags = 0,
......@@ -317,31 +456,28 @@ test "Trie node count" {
317456
318457test "Trie basic" {
319458 var gpa = testing.allocator;
320 var trie: Trie = .{};
321 defer trie.deinit(gpa);
322
323 // root
324 testing.expect(trie.root.edges.items.len == 0);
459 var trie = Trie.init(gpa);
460 defer trie.deinit();
325461
326462 // root --- _st ---> node
327 try trie.put(gpa, .{
463 try trie.put(.{
328464 .name = "_st",
329465 .vmaddr_offset = 0,
330466 .export_flags = 0,
331467 });
332 testing.expect(trie.root.edges.items.len == 1);
333 testing.expect(mem.eql(u8, trie.root.edges.items[0].label, "_st"));
468 testing.expect(trie.root.?.edges.items.len == 1);
469 testing.expect(mem.eql(u8, trie.root.?.edges.items[0].label, "_st"));
334470
335471 {
336472 // root --- _st ---> node --- art ---> node
337 try trie.put(gpa, .{
473 try trie.put(.{
338474 .name = "_start",
339475 .vmaddr_offset = 0,
340476 .export_flags = 0,
341477 });
342 testing.expect(trie.root.edges.items.len == 1);
478 testing.expect(trie.root.?.edges.items.len == 1);
343479
344 const nextEdge = &trie.root.edges.items[0];
480 const nextEdge = &trie.root.?.edges.items[0];
345481 testing.expect(mem.eql(u8, nextEdge.label, "_st"));
346482 testing.expect(nextEdge.to.edges.items.len == 1);
347483 testing.expect(mem.eql(u8, nextEdge.to.edges.items[0].label, "art"));
......@@ -350,14 +486,14 @@ test "Trie basic" {
350486 // root --- _ ---> node --- st ---> node --- art ---> node
351487 // |
352488 // | --- main ---> node
353 try trie.put(gpa, .{
489 try trie.put(.{
354490 .name = "_main",
355491 .vmaddr_offset = 0,
356492 .export_flags = 0,
357493 });
358 testing.expect(trie.root.edges.items.len == 1);
494 testing.expect(trie.root.?.edges.items.len == 1);
359495
360 const nextEdge = &trie.root.edges.items[0];
496 const nextEdge = &trie.root.?.edges.items[0];
361497 testing.expect(mem.eql(u8, nextEdge.label, "_"));
362498 testing.expect(nextEdge.to.edges.items.len == 2);
363499 testing.expect(mem.eql(u8, nextEdge.to.edges.items[0].label, "st"));
......@@ -368,72 +504,81 @@ test "Trie basic" {
368504 }
369505}
370506
371test "Trie.writeULEB128Mem" {
507test "write Trie to a byte stream" {
372508 var gpa = testing.allocator;
373 var trie: Trie = .{};
374 defer trie.deinit(gpa);
509 var trie = Trie.init(gpa);
510 defer trie.deinit();
375511
376 try trie.put(gpa, .{
512 try trie.put(.{
377513 .name = "__mh_execute_header",
378514 .vmaddr_offset = 0,
379515 .export_flags = 0,
380516 });
381 try trie.put(gpa, .{
517 try trie.put(.{
382518 .name = "_main",
383519 .vmaddr_offset = 0x1000,
384520 .export_flags = 0,
385521 });
386522
387 var buffer: std.ArrayListUnmanaged(u8) = .{};
388 defer buffer.deinit(gpa);
389
390 try trie.writeULEB128Mem(gpa, &buffer);
523 try trie.finalize();
524 try trie.finalize(); // Finalizing mulitple times is a nop subsequently unless we add new nodes.
391525
392526 const exp_buffer = [_]u8{
393 0x0,
394 0x1,
395 0x5f,
396 0x0,
397 0x5,
398 0x0,
399 0x2,
400 0x5f,
401 0x6d,
402 0x68,
403 0x5f,
404 0x65,
405 0x78,
406 0x65,
407 0x63,
408 0x75,
409 0x74,
410 0x65,
411 0x5f,
412 0x68,
413 0x65,
414 0x61,
415 0x64,
416 0x65,
417 0x72,
418 0x0,
419 0x21,
420 0x6d,
421 0x61,
422 0x69,
423 0x6e,
424 0x0,
425 0x25,
426 0x2,
427 0x0,
428 0x0,
429 0x0,
430 0x3,
431 0x0,
432 0x80,
433 0x20,
434 0x0,
527 0x0, 0x1, // node root
528 0x5f, 0x0, 0x5, // edge '_'
529 0x0, 0x2, // non-terminal node
530 0x5f, 0x6d, 0x68, 0x5f, 0x65, 0x78, 0x65, 0x63, 0x75, 0x74, // edge '_mh_execute_header'
531 0x65, 0x5f, 0x68, 0x65, 0x61, 0x64, 0x65, 0x72, 0x0, 0x21, // edge '_mh_execute_header'
532 0x6d, 0x61, 0x69, 0x6e, 0x0, 0x25, // edge 'main'
533 0x2, 0x0, 0x0, 0x0, // terminal node
534 0x3, 0x0, 0x80, 0x20, 0x0, // terminal node
435535 };
436536
437 testing.expect(buffer.items.len == exp_buffer.len);
438 testing.expect(mem.eql(u8, buffer.items, exp_buffer[0..]));
537 var buffer = try gpa.alloc(u8, trie.size);
538 defer gpa.free(buffer);
539 var stream = std.io.fixedBufferStream(buffer);
540 {
541 const nwritten = try trie.write(stream.writer());
542 testing.expect(nwritten == trie.size);
543 testing.expect(mem.eql(u8, buffer, exp_buffer[0..]));
544 }
545 {
546 // Writing finalized trie again should yield the same result.
547 try stream.seekTo(0);
548 const nwritten = try trie.write(stream.writer());
549 testing.expect(nwritten == trie.size);
550 testing.expect(mem.eql(u8, buffer, exp_buffer[0..]));
551 }
552}
553
554test "parse Trie from byte stream" {
555 var gpa = testing.allocator;
556
557 const in_buffer = [_]u8{
558 0x0, 0x1, // node root
559 0x5f, 0x0, 0x5, // edge '_'
560 0x0, 0x2, // non-terminal node
561 0x5f, 0x6d, 0x68, 0x5f, 0x65, 0x78, 0x65, 0x63, 0x75, 0x74, // edge '_mh_execute_header'
562 0x65, 0x5f, 0x68, 0x65, 0x61, 0x64, 0x65, 0x72, 0x0, 0x21, // edge '_mh_execute_header'
563 0x6d, 0x61, 0x69, 0x6e, 0x0, 0x25, // edge 'main'
564 0x2, 0x0, 0x0, 0x0, // terminal node
565 0x3, 0x0, 0x80, 0x20, 0x0, // terminal node
566 };
567
568 var in_stream = std.io.fixedBufferStream(in_buffer[0..]);
569 var trie = Trie.init(gpa);
570 defer trie.deinit();
571 const nread = try trie.read(in_stream.reader());
572
573 testing.expect(nread == in_buffer.len);
574
575 try trie.finalize();
576
577 var out_buffer = try gpa.alloc(u8, trie.size);
578 defer gpa.free(out_buffer);
579 var out_stream = std.io.fixedBufferStream(out_buffer);
580 const nwritten = try trie.write(out_stream.writer());
581
582 testing.expect(nwritten == trie.size);
583 testing.expect(mem.eql(u8, in_buffer[0..], out_buffer));
439584}