authorgravatar for kubkon@jakubkonka.comJakub Konka <kubkon@jakubkonka.com> 2020-10-08 18:10:32+02:00
committergravatar for kubkon@jakubkonka.comJakub Konka <kubkon@jakubkonka.com> 2020-10-08 19:53:23+02:00
logba41e599bfaff2c614c4edfbe5c7ffe94b437486
tree4b3cd8d8f9bb3ac1bb52addf00e1b3c581ade710
parent5f86505cf79a0ce75e1a02602ae0e9c845024982
signature Commit is signed but in an unrecognized format.

Clean up writing the trie into ULEB128 byte stream

Prealloc as much as possible to improve alloc performance. Signed-off-by: Jakub Konka <kubkon@jakubkonka.com>

2 files changed, 112 insertions(+), 36 deletions(-)

src/link/MachO.zig+1-3
......@@ -1403,9 +1403,7 @@ fn writeAllUndefSymbols(self: *MachO) !void {
14031403fn writeExportTrie(self: *MachO) !void {
14041404 if (self.global_symbols.items.len == 0) return; // No exports, nothing to do.
14051405
1406 var trie: Trie = .{
1407 .root = .{},
1408 };
1406 var trie: Trie = .{};
14091407 defer trie.deinit(self.base.allocator);
14101408
14111409 const text_segment = self.load_commands.items[self.text_segment_cmd_index.?].Segment;
src/link/MachO/Trie.zig+111-33
......@@ -35,6 +35,7 @@ const mem = std.mem;
3535const leb = std.debug.leb;
3636const log = std.log.scoped(.link);
3737const testing = std.testing;
38const assert = std.debug.assert;
3839const Allocator = mem.Allocator;
3940
4041pub const Symbol = struct {
......@@ -57,9 +58,13 @@ const Edge = struct {
5758};
5859
5960const Node = struct {
61 /// Export flags associated with this exported symbol (if any).
6062 export_flags: ?u64 = null,
63 /// VM address offset wrt to the section this symbol is defined against (if any).
6164 vmaddr_offset: ?u64 = null,
62 trie_offset: usize = 0,
65 /// Offset of this node in the trie output byte stream.
66 trie_offset: ?usize = null,
67 /// List of all edges originating from this node.
6368 edges: std.ArrayListUnmanaged(Edge) = .{},
6469
6570 fn deinit(self: *Node, alloc: *Allocator) void {
......@@ -69,12 +74,24 @@ const Node = struct {
6974 self.edges.deinit(alloc);
7075 }
7176
72 fn put(self: *Node, alloc: *Allocator, label: []const u8) !*Node {
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`.
85 fn put(self: *Node, alloc: *Allocator, label: []const u8, node_count: usize) !PutResult {
86 var curr_node_count = node_count;
7387 // Check for match with edges from this node.
7488 for (self.edges.items) |*edge| {
75 const match = mem.indexOfDiff(u8, edge.label, label) orelse return edge.to;
89 const match = mem.indexOfDiff(u8, edge.label, label) orelse return PutResult{
90 .node = edge.to,
91 .node_count = curr_node_count,
92 };
7693 if (match == 0) continue;
77 if (match == edge.label.len) return edge.to.put(alloc, label[match..]);
94 if (match == edge.label.len) return edge.to.put(alloc, label[match..], curr_node_count);
7895
7996 // Found a match, need to splice up nodes.
8097 // From: A -> B
......@@ -85,6 +102,7 @@ const Node = struct {
85102 const to_node = edge.to;
86103 edge.to = mid;
87104 edge.label = label[0..match];
105 curr_node_count += 1;
88106
89107 try mid.edges.append(alloc, .{
90108 .from = mid,
......@@ -93,15 +111,16 @@ const Node = struct {
93111 });
94112
95113 if (match == label.len) {
96 return to_node;
114 return PutResult{ .node = to_node, .node_count = curr_node_count };
97115 } else {
98 return mid.put(alloc, label[match..]);
116 return mid.put(alloc, label[match..], curr_node_count);
99117 }
100118 }
101119
102 // Add a new edge.
120 // Add a new node.
103121 const node = try alloc.create(Node);
104122 node.* = .{};
123 curr_node_count += 1;
105124
106125 try self.edges.append(alloc, .{
107126 .from = self,
......@@ -109,10 +128,13 @@ const Node = struct {
109128 .label = label,
110129 });
111130
112 return node;
131 return PutResult{ .node = node, .node_count = curr_node_count };
113132 }
114133
115 fn writeULEB128Mem(self: Node, alloc: *Allocator, buffer: *std.ArrayListUnmanaged(u8)) !void {
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.
116138 if (self.vmaddr_offset) |offset| {
117139 // Terminal node info: encode export flags and vmaddr offset of this symbol.
118140 var info_buf_len: usize = 0;
......@@ -125,33 +147,35 @@ const Node = struct {
125147 const size_buf_len = try leb.writeULEB128Mem(size_buf[0..], info_buf_len);
126148
127149 // Now, write them to the output buffer.
128 try buffer.ensureCapacity(alloc, buffer.items.len + info_buf_len + size_buf_len);
129150 buffer.appendSliceAssumeCapacity(size_buf[0..size_buf_len]);
130151 buffer.appendSliceAssumeCapacity(info_buf[0..info_buf_len]);
131152 } else {
132153 // Non-terminal node is delimited by 0 byte.
133 try buffer.append(alloc, 0);
154 buffer.appendAssumeCapacity(0);
134155 }
135156 // Write number of edges (max legal number of edges is 256).
136 try buffer.append(alloc, @intCast(u8, self.edges.items.len));
157 buffer.appendAssumeCapacity(@intCast(u8, self.edges.items.len));
137158
138159 for (self.edges.items) |edge| {
139160 // Write edges labels.
140 try buffer.ensureCapacity(alloc, buffer.items.len + edge.label.len + 1); // +1 to account for null-byte
141161 buffer.appendSliceAssumeCapacity(edge.label);
142162 buffer.appendAssumeCapacity(0);
143163
144164 var buf: [@sizeOf(u64)]u8 = undefined;
145 const buf_len = try leb.writeULEB128Mem(buf[0..], edge.to.trie_offset);
146 try buffer.appendSlice(alloc, buf[0..buf_len]);
165 const buf_len = try leb.writeULEB128Mem(buf[0..], edge.to.trie_offset.?);
166 buffer.appendSliceAssumeCapacity(buf[0..buf_len]);
147167 }
148168 }
149169
150170 const UpdateResult = struct {
171 /// Current size of this node in bytes.
151172 node_size: usize,
173 /// True if the trie offset of this node in the output byte stream
174 /// would need updating; false otherwise.
152175 updated: bool,
153176 };
154177
178 /// Updates offset of this node in the output byte stream.
155179 fn updateOffset(self: *Node, offset: usize) UpdateResult {
156180 var node_size: usize = 0;
157181 if (self.vmaddr_offset) |vmaddr| {
......@@ -164,15 +188,18 @@ const Node = struct {
164188 node_size += 1; // 1 byte for edge count
165189
166190 for (self.edges.items) |edge| {
167 node_size += edge.label.len + 1 + sizeULEB128Mem(edge.to.trie_offset);
191 const next_node_offset = edge.to.trie_offset orelse 0;
192 node_size += edge.label.len + 1 + sizeULEB128Mem(next_node_offset);
168193 }
169194
170 const updated = offset != self.trie_offset;
195 const trie_offset = self.trie_offset orelse 0;
196 const updated = offset != trie_offset;
171197 self.trie_offset = offset;
172198
173199 return .{ .node_size = node_size, .updated = updated };
174200 }
175201
202 /// Calculates number of bytes in ULEB128 encoding of value.
176203 fn sizeULEB128Mem(value: u64) usize {
177204 var res: usize = 0;
178205 var v = value;
......@@ -185,15 +212,22 @@ const Node = struct {
185212 }
186213};
187214
188root: Node,
215/// Count of nodes in the trie.
216/// The count is updated at every `put` call.
217/// The trie always consists of at least a root node, hence
218/// the count always starts at 1.
219node_count: usize = 1,
220/// The root node of the trie.
221root: Node = .{},
189222
190223/// Insert a symbol into the trie, updating the prefixes in the process.
191224/// This operation may change the layout of the trie by splicing edges in
192225/// certain circumstances.
193226pub fn put(self: *Trie, alloc: *Allocator, symbol: Symbol) !void {
194 const node = try self.root.put(alloc, symbol.name);
195 node.vmaddr_offset = symbol.vmaddr_offset;
196 node.export_flags = symbol.export_flags;
227 const res = try self.root.put(alloc, symbol.name, 0);
228 self.node_count += res.node_count;
229 res.node.vmaddr_offset = symbol.vmaddr_offset;
230 res.node.export_flags = symbol.export_flags;
197231}
198232
199233/// Write the trie to a buffer ULEB128 encoded.
......@@ -201,11 +235,13 @@ pub fn writeULEB128Mem(self: *Trie, alloc: *Allocator, buffer: *std.ArrayListUnm
201235 var ordered_nodes: std.ArrayListUnmanaged(*Node) = .{};
202236 defer ordered_nodes.deinit(alloc);
203237
204 try walkInOrder(&self.root, alloc, &ordered_nodes);
238 try ordered_nodes.ensureCapacity(alloc, self.node_count);
239 walkInOrder(&self.root, &ordered_nodes);
205240
241 var offset: usize = 0;
206242 var more: bool = true;
207243 while (more) {
208 var offset: usize = 0;
244 offset = 0;
209245 more = false;
210246 for (ordered_nodes.items) |node| {
211247 const res = node.updateOffset(offset);
......@@ -214,15 +250,17 @@ pub fn writeULEB128Mem(self: *Trie, alloc: *Allocator, buffer: *std.ArrayListUnm
214250 }
215251 }
216252
253 try buffer.ensureCapacity(alloc, buffer.items.len + offset);
217254 for (ordered_nodes.items) |node| {
218 try node.writeULEB128Mem(alloc, buffer);
255 try node.writeULEB128Mem(buffer);
219256 }
220257}
221258
222fn walkInOrder(node: *Node, alloc: *Allocator, list: *std.ArrayListUnmanaged(*Node)) error{OutOfMemory}!void {
223 try list.append(alloc, node);
259/// Walks the trie in DFS order gathering all nodes into a linear stream of nodes.
260fn walkInOrder(node: *Node, list: *std.ArrayListUnmanaged(*Node)) void {
261 list.appendAssumeCapacity(node);
224262 for (node.edges.items) |*edge| {
225 try walkInOrder(edge.to, alloc, list);
263 walkInOrder(edge.to, list);
226264 }
227265}
228266
......@@ -230,11 +268,53 @@ pub fn deinit(self: *Trie, alloc: *Allocator) void {
230268 self.root.deinit(alloc);
231269}
232270
271test "Trie node count" {
272 var gpa = testing.allocator;
273 var trie: Trie = .{};
274 defer trie.deinit(gpa);
275
276 testing.expectEqual(trie.node_count, 1);
277
278 try trie.put(gpa, .{
279 .name = "_main",
280 .vmaddr_offset = 0,
281 .export_flags = 0,
282 });
283 testing.expectEqual(trie.node_count, 2);
284
285 // Inserting the same node shouldn't update the trie.
286 try trie.put(gpa, .{
287 .name = "_main",
288 .vmaddr_offset = 0,
289 .export_flags = 0,
290 });
291 testing.expectEqual(trie.node_count, 2);
292
293 try trie.put(gpa, .{
294 .name = "__mh_execute_header",
295 .vmaddr_offset = 0x1000,
296 .export_flags = 0,
297 });
298 testing.expectEqual(trie.node_count, 4);
299
300 // Inserting the same node shouldn't update the trie.
301 try trie.put(gpa, .{
302 .name = "__mh_execute_header",
303 .vmaddr_offset = 0x1000,
304 .export_flags = 0,
305 });
306 testing.expectEqual(trie.node_count, 4);
307 try trie.put(gpa, .{
308 .name = "_main",
309 .vmaddr_offset = 0,
310 .export_flags = 0,
311 });
312 testing.expectEqual(trie.node_count, 4);
313}
314
233315test "Trie basic" {
234316 var gpa = testing.allocator;
235 var trie: Trie = .{
236 .root = .{},
237 };
317 var trie: Trie = .{};
238318 defer trie.deinit(gpa);
239319
240320 // root
......@@ -287,9 +367,7 @@ test "Trie basic" {
287367
288368test "Trie.writeULEB128Mem" {
289369 var gpa = testing.allocator;
290 var trie: Trie = .{
291 .root = .{},
292 };
370 var trie: Trie = .{};
293371 defer trie.deinit(gpa);
294372
295373 try trie.put(gpa, .{