| ... | @@ -34,156 +34,218 @@ const std = @import("std"); | ... | @@ -34,156 +34,218 @@ const std = @import("std"); |
| 34 | const mem = std.mem; | 34 | const mem = std.mem; |
| 35 | const leb = std.leb; | 35 | const leb = std.leb; |
| 36 | const log = std.log.scoped(.link); | 36 | const log = std.log.scoped(.link); |
| | 37 | const macho = std.macho; |
| 37 | const testing = std.testing; | 38 | const testing = std.testing; |
| 38 | const assert = std.debug.assert; | 39 | const assert = std.debug.assert; |
| 39 | const Allocator = mem.Allocator; | 40 | const Allocator = mem.Allocator; |
| 40 | | 41 | |
| 41 | pub const Symbol = struct { | 42 | pub const Node = struct { |
| 42 | name: []const u8, | 43 | base: *Trie, |
| 43 | vmaddr_offset: u64, | | |
| 44 | export_flags: u64, | | |
| 45 | }; | | |
| 46 | | | |
| 47 | const Edge = struct { | | |
| 48 | from: *Node, | | |
| 49 | to: *Node, | | |
| 50 | label: []const u8, | | |
| 51 | | 44 | |
| 52 | fn deinit(self: *Edge, alloc: *Allocator) void { | 45 | /// Terminal info associated with this node. |
| 53 | self.to.deinit(alloc); | 46 | /// If this node is not a terminal node, info is null. |
| 54 | alloc.destroy(self.to); | 47 | terminal_info: ?struct { |
| 55 | self.from = undefined; | 48 | /// Export flags associated with this exported symbol. |
| 56 | self.to = undefined; | 49 | export_flags: u64, |
| 57 | } | 50 | /// VM address offset wrt to the section this symbol is defined against. |
| 58 | }; | 51 | vmaddr_offset: u64, |
| | 52 | } = null, |
| 59 | | 53 | |
| 60 | const 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, | | |
| 65 | /// Offset of this node in the trie output byte stream. | 54 | /// Offset of this node in the trie output byte stream. |
| 66 | trie_offset: ?usize = null, | 55 | trie_offset: ?usize = null, |
| | 56 | |
| 67 | /// List of all edges originating from this node. | 57 | /// List of all edges originating from this node. |
| 68 | edges: std.ArrayListUnmanaged(Edge) = .{}, | 58 | edges: std.ArrayListUnmanaged(Edge) = .{}, |
| 69 | | 59 | |
| 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 { |
| 71 | for (self.edges.items) |*edge| { | 79 | for (self.edges.items) |*edge| { |
| 72 | edge.deinit(alloc); | 80 | edge.deinit(allocator); |
| 73 | } | 81 | } |
| 74 | self.edges.deinit(alloc); | 82 | self.edges.deinit(allocator); |
| 75 | } | 83 | } |
| 76 | | 84 | |
| 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 | /// Inserts a new node starting from `self`. |
| 85 | fn put(self: *Node, alloc: *Allocator, label: []const u8, node_count: usize) !PutResult { | 86 | fn put(self: *Node, allocator: *Allocator, label: []const u8) !*Node { |
| 86 | var curr_node_count = node_count; | | |
| 87 | // Check for match with edges from this node. | 87 | // Check for match with edges from this node. |
| 88 | for (self.edges.items) |*edge| { | 88 | for (self.edges.items) |*edge| { |
| 89 | const match = mem.indexOfDiff(u8, edge.label, label) orelse return PutResult{ | 89 | 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; | 90 | 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..]); |
| 95 | | 92 | |
| 96 | // Found a match, need to splice up nodes. | 93 | // Found a match, need to splice up nodes. |
| 97 | // From: A -> B | 94 | // From: A -> B |
| 98 | // To: A -> C -> B | 95 | // To: A -> C -> B |
| 99 | const mid = try alloc.create(Node); | 96 | const mid = try allocator.create(Node); |
| 100 | mid.* = .{}; | 97 | mid.* = .{ .base = self.base }; |
| 101 | const to_label = edge.label; | 98 | var to_label = try allocator.dupe(u8, edge.label[match..]); |
| | 99 | allocator.free(edge.label); |
| 102 | const to_node = edge.to; | 100 | const to_node = edge.to; |
| 103 | edge.to = mid; | 101 | edge.to = mid; |
| 104 | edge.label = label[0..match]; | 102 | edge.label = try allocator.dupe(u8, label[0..match]); |
| 105 | curr_node_count += 1; | 103 | self.base.node_count += 1; |
| 106 | | 104 | |
| 107 | try mid.edges.append(alloc, .{ | 105 | try mid.edges.append(allocator, .{ |
| 108 | .from = mid, | 106 | .from = mid, |
| 109 | .to = to_node, | 107 | .to = to_node, |
| 110 | .label = to_label[match..], | 108 | .label = to_label, |
| 111 | }); | 109 | }); |
| 112 | | 110 | |
| 113 | if (match == label.len) { | 111 | return if (match == label.len) to_node else mid.put(allocator, 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 | } | 112 | } |
| 119 | | 113 | |
| 120 | // Add a new node. | 114 | // Add a new node. |
| 121 | const node = try alloc.create(Node); | 115 | const node = try allocator.create(Node); |
| 122 | node.* = .{}; | 116 | node.* = .{ .base = self.base }; |
| 123 | curr_node_count += 1; | 117 | self.base.node_count += 1; |
| 124 | | 118 | |
| 125 | try self.edges.append(alloc, .{ | 119 | try self.edges.append(allocator, .{ |
| 126 | .from = self, | 120 | .from = self, |
| 127 | .to = node, | 121 | .to = node, |
| 128 | .label = label, | 122 | .label = try allocator.dupe(u8, label), |
| 129 | }); | 123 | }); |
| 130 | | 124 | |
| 131 | return PutResult{ .node = node, .node_count = curr_node_count }; | 125 | return node; |
| 132 | } | 126 | } |
| 133 | | 127 | |
| 134 | /// This method should only be called *after* updateOffset has been called! | 128 | /// Recursively parses the node from the input byte stream. |
| 135 | /// In case this is not upheld, this method will panic. | 129 | fn read(self: *Node, allocator: *Allocator, reader: anytype) Trie.ReadError!usize { |
| 136 | fn writeULEB128Mem(self: Node, buffer: *std.ArrayListUnmanaged(u8)) !void { | 130 | self.node_dirty = true; |
| 137 | assert(self.trie_offset != null); // You need to call updateOffset first. | 131 | const trie_offset = try reader.context.getPos(); |
| 138 | if (self.vmaddr_offset) |offset| { | 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| { |
| 139 | // Terminal node info: encode export flags and vmaddr offset of this symbol. | 201 | // Terminal node info: encode export flags and vmaddr offset of this symbol. |
| 140 | var info_buf_len: usize = 0; | 202 | var info_buf_len: usize = 0; |
| 141 | var info_buf: [@sizeOf(u64) * 2]u8 = undefined; | 203 | var info_buf: [@sizeOf(u64) * 2]u8 = undefined; |
| 142 | var info_stream = std.io.fixedBufferStream(&info_buf); | 204 | var info_stream = std.io.fixedBufferStream(&info_buf); |
| 143 | try leb.writeULEB128(info_stream.writer(), self.export_flags.?); | 205 | // TODO Implement for special flags. |
| 144 | try leb.writeULEB128(info_stream.writer(), offset); | 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); |
| 145 | | 210 | |
| 146 | // Encode the size of the terminal node info. | 211 | // Encode the size of the terminal node info. |
| 147 | var size_buf: [@sizeOf(u64)]u8 = undefined; | 212 | var size_buf: [@sizeOf(u64)]u8 = undefined; |
| 148 | var size_stream = std.io.fixedBufferStream(&size_buf); | 213 | var size_stream = std.io.fixedBufferStream(&size_buf); |
| 149 | try leb.writeULEB128(size_stream.writer(), info_stream.pos); | 214 | try leb.writeULEB128(size_stream.writer(), info_stream.pos); |
| 150 | | 215 | |
| 151 | // Now, write them to the output buffer. | 216 | // Now, write them to the output stream. |
| 152 | buffer.appendSliceAssumeCapacity(size_buf[0..size_stream.pos]); | 217 | try writer.writeAll(size_buf[0..size_stream.pos]); |
| 153 | buffer.appendSliceAssumeCapacity(info_buf[0..info_stream.pos]); | 218 | try writer.writeAll(info_buf[0..info_stream.pos]); |
| 154 | } else { | 219 | } else { |
| 155 | // Non-terminal node is delimited by 0 byte. | 220 | // Non-terminal node is delimited by 0 byte. |
| 156 | buffer.appendAssumeCapacity(0); | 221 | try writer.writeByte(0); |
| 157 | } | 222 | } |
| 158 | // Write number of edges (max legal number of edges is 256). | 223 | // 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)); |
| 160 | | 225 | |
| 161 | for (self.edges.items) |edge| { | 226 | for (self.edges.items) |edge| { |
| 162 | // Write edges labels. | 227 | // Write edge label and offset to next node in trie. |
| 163 | buffer.appendSliceAssumeCapacity(edge.label); | 228 | try writer.writeAll(edge.label); |
| 164 | buffer.appendAssumeCapacity(0); | 229 | try writer.writeByte(0); |
| 165 | | 230 | try leb.writeULEB128(writer, edge.to.trie_offset.?); |
| 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]); | | |
| 170 | } | 231 | } |
| 171 | } | 232 | } |
| 172 | | 233 | |
| 173 | const UpdateResult = struct { | 234 | const FinalizeResult = struct { |
| 174 | /// Current size of this node in bytes. | 235 | /// Current size of this node in bytes. |
| 175 | node_size: usize, | 236 | node_size: usize, |
| | 237 | |
| 176 | /// True if the trie offset of this node in the output byte stream | 238 | /// True if the trie offset of this node in the output byte stream |
| 177 | /// would need updating; false otherwise. | 239 | /// would need updating; false otherwise. |
| 178 | updated: bool, | 240 | updated: bool, |
| 179 | }; | 241 | }; |
| 180 | | 242 | |
| 181 | /// Updates offset of this node in the output byte stream. | 243 | /// 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 { |
| 183 | var node_size: usize = 0; | 245 | var node_size: usize = 0; |
| 184 | if (self.vmaddr_offset) |vmaddr| { | 246 | if (self.terminal_info) |info| { |
| 185 | node_size += sizeULEB128Mem(self.export_flags.?); | 247 | node_size += sizeULEB128Mem(info.export_flags); |
| 186 | node_size += sizeULEB128Mem(vmaddr); | 248 | node_size += sizeULEB128Mem(info.vmaddr_offset); |
| 187 | node_size += sizeULEB128Mem(node_size); | 249 | node_size += sizeULEB128Mem(node_size); |
| 188 | } else { | 250 | } else { |
| 189 | node_size += 1; // 0x0 for non-terminal nodes | 251 | node_size += 1; // 0x0 for non-terminal nodes |
| ... | @@ -196,8 +258,9 @@ const Node = struct { | ... | @@ -196,8 +258,9 @@ const Node = struct { |
| 196 | } | 258 | } |
| 197 | | 259 | |
| 198 | const trie_offset = self.trie_offset orelse 0; | 260 | const trie_offset = self.trie_offset orelse 0; |
| 199 | const updated = offset != trie_offset; | 261 | const updated = offset_in_trie != trie_offset; |
| 200 | self.trie_offset = offset; | 262 | self.trie_offset = offset_in_trie; |
| | 263 | self.node_dirty = false; |
| 201 | | 264 | |
| 202 | return .{ .node_size = node_size, .updated = updated }; | 265 | return .{ .node_size = node_size, .updated = updated }; |
| 203 | } | 266 | } |
| ... | @@ -215,70 +278,146 @@ const Node = struct { | ... | @@ -215,70 +278,146 @@ const Node = struct { |
| 215 | } | 278 | } |
| 216 | }; | 279 | }; |
| 217 | | 280 | |
| 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. | | |
| 222 | node_count: usize = 1, | | |
| 223 | /// The root node of the trie. | 281 | /// The root node of the trie. |
| 224 | root: Node = .{}, | 282 | root: ?*Node = null, |
| | 283 | |
| | 284 | allocator: *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. |
| | 291 | ordered_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. |
| | 298 | size: usize = 0, |
| | 299 | |
| | 300 | /// Number of nodes currently in the trie. |
| | 301 | node_count: usize = 0, |
| | 302 | |
| | 303 | trie_dirty: bool = true, |
| | 304 | |
| | 305 | pub fn init(allocator: *Allocator) Trie { |
| | 306 | return .{ .allocator = allocator }; |
| | 307 | } |
| | 308 | |
| | 309 | /// Export symbol that is to be placed in the trie. |
| | 310 | pub 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 | }; |
| 225 | | 321 | |
| 226 | /// Insert a symbol into the trie, updating the prefixes in the process. | 322 | /// 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 in | 323 | /// This operation may change the layout of the trie by splicing edges in |
| 228 | /// certain circumstances. | 324 | /// certain circumstances. |
| 229 | pub fn put(self: *Trie, alloc: *Allocator, symbol: Symbol) !void { | 325 | pub fn put(self: *Trie, symbol: ExportSymbol) !void { |
| 230 | const res = try self.root.put(alloc, symbol.name, 0); | 326 | try self.createRoot(); |
| 231 | self.node_count += res.node_count; | 327 | const node = try self.root.?.put(self.allocator, symbol.name); |
| 232 | res.node.vmaddr_offset = symbol.vmaddr_offset; | 328 | node.terminal_info = .{ |
| 233 | res.node.export_flags = symbol.export_flags; | 329 | .vmaddr_offset = symbol.vmaddr_offset, |
| | 330 | .export_flags = symbol.export_flags, |
| | 331 | }; |
| | 332 | self.trie_dirty = true; |
| 234 | } | 333 | } |
| 235 | | 334 | |
| 236 | /// Write the trie to a buffer ULEB128 encoded. | 335 | /// Finalizes this trie for writing to a byte stream. |
| 237 | pub fn writeULEB128Mem(self: *Trie, alloc: *Allocator, buffer: *std.ArrayListUnmanaged(u8)) !void { | 336 | /// This step performs multiple passes through the trie ensuring |
| 238 | var ordered_nodes: std.ArrayListUnmanaged(*Node) = .{}; | 337 | /// there are no gaps after every `Node` is ULEB128 encoded. |
| 239 | defer ordered_nodes.deinit(alloc); | 338 | /// Call this method before trying to `write` the trie to a byte stream. |
| | 339 | pub 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.?); |
| 240 | | 348 | |
| 241 | try ordered_nodes.ensureCapacity(alloc, self.node_count); | 349 | while (fifo.readItem()) |next| { |
| 242 | walkInOrder(&self.root, &ordered_nodes); | 350 | for (next.edges.items) |*edge| { |
| | 351 | try fifo.writeItem(edge.to); |
| | 352 | } |
| | 353 | self.ordered_nodes.appendAssumeCapacity(next); |
| | 354 | } |
| 243 | | 355 | |
| 244 | var offset: usize = 0; | | |
| 245 | var more: bool = true; | 356 | var more: bool = true; |
| 246 | while (more) { | 357 | while (more) { |
| 247 | offset = 0; | 358 | self.size = 0; |
| 248 | more = false; | 359 | more = false; |
| 249 | for (ordered_nodes.items) |node| { | 360 | for (self.ordered_nodes.items) |node| { |
| 250 | const res = node.updateOffset(offset); | 361 | const res = node.finalize(self.size); |
| 251 | offset += res.node_size; | 362 | self.size += res.node_size; |
| 252 | if (res.updated) more = true; | 363 | if (res.updated) more = true; |
| 253 | } | 364 | } |
| 254 | } | 365 | } |
| 255 | | 366 | |
| 256 | try buffer.ensureCapacity(alloc, buffer.items.len + offset); | 367 | self.trie_dirty = false; |
| 257 | for (ordered_nodes.items) |node| { | 368 | } |
| 258 | try node.writeULEB128Mem(buffer); | 369 | |
| | 370 | const ReadError = error{ |
| | 371 | OutOfMemory, |
| | 372 | EndOfStream, |
| | 373 | Overflow, |
| | 374 | }; |
| | 375 | |
| | 376 | /// Parse the trie from a byte stream. |
| | 377 | pub 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. |
| | 386 | pub 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()); |
| 259 | } | 391 | } |
| | 392 | return counting_writer.bytes_written; |
| 260 | } | 393 | } |
| 261 | | 394 | |
| 262 | /// Walks the trie in DFS order gathering all nodes into a linear stream of nodes. | 395 | pub fn deinit(self: *Trie) void { |
| 263 | fn walkInOrder(node: *Node, list: *std.ArrayListUnmanaged(*Node)) void { | 396 | if (self.root) |root| { |
| 264 | list.appendAssumeCapacity(node); | 397 | root.deinit(self.allocator); |
| 265 | for (node.edges.items) |*edge| { | 398 | self.allocator.destroy(root); |
| 266 | walkInOrder(edge.to, list); | | |
| 267 | } | 399 | } |
| | 400 | self.ordered_nodes.deinit(self.allocator); |
| 268 | } | 401 | } |
| 269 | | 402 | |
| 270 | pub fn deinit(self: *Trie, alloc: *Allocator) void { | 403 | fn createRoot(self: *Trie) !void { |
| 271 | self.root.deinit(alloc); | 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 | } |
| 272 | } | 410 | } |
| 273 | | 411 | |
| 274 | test "Trie node count" { | 412 | test "Trie node count" { |
| 275 | var gpa = testing.allocator; | 413 | var gpa = testing.allocator; |
| 276 | var trie: Trie = .{}; | 414 | var trie = Trie.init(gpa); |
| 277 | defer trie.deinit(gpa); | 415 | defer trie.deinit(); |
| 278 | | 416 | |
| 279 | testing.expectEqual(trie.node_count, 1); | 417 | testing.expectEqual(trie.node_count, 0); |
| | 418 | testing.expect(trie.root == null); |
| 280 | | 419 | |
| 281 | try trie.put(gpa, .{ | 420 | try trie.put(.{ |
| 282 | .name = "_main", | 421 | .name = "_main", |
| 283 | .vmaddr_offset = 0, | 422 | .vmaddr_offset = 0, |
| 284 | .export_flags = 0, | 423 | .export_flags = 0, |
| ... | @@ -286,14 +425,14 @@ test "Trie node count" { | ... | @@ -286,14 +425,14 @@ test "Trie node count" { |
| 286 | testing.expectEqual(trie.node_count, 2); | 425 | testing.expectEqual(trie.node_count, 2); |
| 287 | | 426 | |
| 288 | // Inserting the same node shouldn't update the trie. | 427 | // Inserting the same node shouldn't update the trie. |
| 289 | try trie.put(gpa, .{ | 428 | try trie.put(.{ |
| 290 | .name = "_main", | 429 | .name = "_main", |
| 291 | .vmaddr_offset = 0, | 430 | .vmaddr_offset = 0, |
| 292 | .export_flags = 0, | 431 | .export_flags = 0, |
| 293 | }); | 432 | }); |
| 294 | testing.expectEqual(trie.node_count, 2); | 433 | testing.expectEqual(trie.node_count, 2); |
| 295 | | 434 | |
| 296 | try trie.put(gpa, .{ | 435 | try trie.put(.{ |
| 297 | .name = "__mh_execute_header", | 436 | .name = "__mh_execute_header", |
| 298 | .vmaddr_offset = 0x1000, | 437 | .vmaddr_offset = 0x1000, |
| 299 | .export_flags = 0, | 438 | .export_flags = 0, |
| ... | @@ -301,13 +440,13 @@ test "Trie node count" { | ... | @@ -301,13 +440,13 @@ test "Trie node count" { |
| 301 | testing.expectEqual(trie.node_count, 4); | 440 | testing.expectEqual(trie.node_count, 4); |
| 302 | | 441 | |
| 303 | // Inserting the same node shouldn't update the trie. | 442 | // Inserting the same node shouldn't update the trie. |
| 304 | try trie.put(gpa, .{ | 443 | try trie.put(.{ |
| 305 | .name = "__mh_execute_header", | 444 | .name = "__mh_execute_header", |
| 306 | .vmaddr_offset = 0x1000, | 445 | .vmaddr_offset = 0x1000, |
| 307 | .export_flags = 0, | 446 | .export_flags = 0, |
| 308 | }); | 447 | }); |
| 309 | testing.expectEqual(trie.node_count, 4); | 448 | testing.expectEqual(trie.node_count, 4); |
| 310 | try trie.put(gpa, .{ | 449 | try trie.put(.{ |
| 311 | .name = "_main", | 450 | .name = "_main", |
| 312 | .vmaddr_offset = 0, | 451 | .vmaddr_offset = 0, |
| 313 | .export_flags = 0, | 452 | .export_flags = 0, |
| ... | @@ -317,31 +456,28 @@ test "Trie node count" { | ... | @@ -317,31 +456,28 @@ test "Trie node count" { |
| 317 | | 456 | |
| 318 | test "Trie basic" { | 457 | test "Trie basic" { |
| 319 | var gpa = testing.allocator; | 458 | var gpa = testing.allocator; |
| 320 | var trie: Trie = .{}; | 459 | var trie = Trie.init(gpa); |
| 321 | defer trie.deinit(gpa); | 460 | defer trie.deinit(); |
| 322 | | | |
| 323 | // root | | |
| 324 | testing.expect(trie.root.edges.items.len == 0); | | |
| 325 | | 461 | |
| 326 | // root --- _st ---> node | 462 | // root --- _st ---> node |
| 327 | try trie.put(gpa, .{ | 463 | try trie.put(.{ |
| 328 | .name = "_st", | 464 | .name = "_st", |
| 329 | .vmaddr_offset = 0, | 465 | .vmaddr_offset = 0, |
| 330 | .export_flags = 0, | 466 | .export_flags = 0, |
| 331 | }); | 467 | }); |
| 332 | testing.expect(trie.root.edges.items.len == 1); | 468 | testing.expect(trie.root.?.edges.items.len == 1); |
| 333 | testing.expect(mem.eql(u8, trie.root.edges.items[0].label, "_st")); | 469 | testing.expect(mem.eql(u8, trie.root.?.edges.items[0].label, "_st")); |
| 334 | | 470 | |
| 335 | { | 471 | { |
| 336 | // root --- _st ---> node --- art ---> node | 472 | // root --- _st ---> node --- art ---> node |
| 337 | try trie.put(gpa, .{ | 473 | try trie.put(.{ |
| 338 | .name = "_start", | 474 | .name = "_start", |
| 339 | .vmaddr_offset = 0, | 475 | .vmaddr_offset = 0, |
| 340 | .export_flags = 0, | 476 | .export_flags = 0, |
| 341 | }); | 477 | }); |
| 342 | testing.expect(trie.root.edges.items.len == 1); | 478 | testing.expect(trie.root.?.edges.items.len == 1); |
| 343 | | 479 | |
| 344 | const nextEdge = &trie.root.edges.items[0]; | 480 | const nextEdge = &trie.root.?.edges.items[0]; |
| 345 | testing.expect(mem.eql(u8, nextEdge.label, "_st")); | 481 | testing.expect(mem.eql(u8, nextEdge.label, "_st")); |
| 346 | testing.expect(nextEdge.to.edges.items.len == 1); | 482 | testing.expect(nextEdge.to.edges.items.len == 1); |
| 347 | testing.expect(mem.eql(u8, nextEdge.to.edges.items[0].label, "art")); | 483 | testing.expect(mem.eql(u8, nextEdge.to.edges.items[0].label, "art")); |
| ... | @@ -350,14 +486,14 @@ test "Trie basic" { | ... | @@ -350,14 +486,14 @@ test "Trie basic" { |
| 350 | // root --- _ ---> node --- st ---> node --- art ---> node | 486 | // root --- _ ---> node --- st ---> node --- art ---> node |
| 351 | // | | 487 | // | |
| 352 | // | --- main ---> node | 488 | // | --- main ---> node |
| 353 | try trie.put(gpa, .{ | 489 | try trie.put(.{ |
| 354 | .name = "_main", | 490 | .name = "_main", |
| 355 | .vmaddr_offset = 0, | 491 | .vmaddr_offset = 0, |
| 356 | .export_flags = 0, | 492 | .export_flags = 0, |
| 357 | }); | 493 | }); |
| 358 | testing.expect(trie.root.edges.items.len == 1); | 494 | testing.expect(trie.root.?.edges.items.len == 1); |
| 359 | | 495 | |
| 360 | const nextEdge = &trie.root.edges.items[0]; | 496 | const nextEdge = &trie.root.?.edges.items[0]; |
| 361 | testing.expect(mem.eql(u8, nextEdge.label, "_")); | 497 | testing.expect(mem.eql(u8, nextEdge.label, "_")); |
| 362 | testing.expect(nextEdge.to.edges.items.len == 2); | 498 | testing.expect(nextEdge.to.edges.items.len == 2); |
| 363 | testing.expect(mem.eql(u8, nextEdge.to.edges.items[0].label, "st")); | 499 | testing.expect(mem.eql(u8, nextEdge.to.edges.items[0].label, "st")); |
| ... | @@ -368,72 +504,81 @@ test "Trie basic" { | ... | @@ -368,72 +504,81 @@ test "Trie basic" { |
| 368 | } | 504 | } |
| 369 | } | 505 | } |
| 370 | | 506 | |
| 371 | test "Trie.writeULEB128Mem" { | 507 | test "write Trie to a byte stream" { |
| 372 | var gpa = testing.allocator; | 508 | var gpa = testing.allocator; |
| 373 | var trie: Trie = .{}; | 509 | var trie = Trie.init(gpa); |
| 374 | defer trie.deinit(gpa); | 510 | defer trie.deinit(); |
| 375 | | 511 | |
| 376 | try trie.put(gpa, .{ | 512 | try trie.put(.{ |
| 377 | .name = "__mh_execute_header", | 513 | .name = "__mh_execute_header", |
| 378 | .vmaddr_offset = 0, | 514 | .vmaddr_offset = 0, |
| 379 | .export_flags = 0, | 515 | .export_flags = 0, |
| 380 | }); | 516 | }); |
| 381 | try trie.put(gpa, .{ | 517 | try trie.put(.{ |
| 382 | .name = "_main", | 518 | .name = "_main", |
| 383 | .vmaddr_offset = 0x1000, | 519 | .vmaddr_offset = 0x1000, |
| 384 | .export_flags = 0, | 520 | .export_flags = 0, |
| 385 | }); | 521 | }); |
| 386 | | 522 | |
| 387 | var buffer: std.ArrayListUnmanaged(u8) = .{}; | 523 | try trie.finalize(); |
| 388 | defer buffer.deinit(gpa); | 524 | try trie.finalize(); // Finalizing mulitple times is a nop subsequently unless we add new nodes. |
| 389 | | | |
| 390 | try trie.writeULEB128Mem(gpa, &buffer); | | |
| 391 | | 525 | |
| 392 | const exp_buffer = [_]u8{ | 526 | const exp_buffer = [_]u8{ |
| 393 | 0x0, | 527 | 0x0, 0x1, // node root |
| 394 | 0x1, | 528 | 0x5f, 0x0, 0x5, // edge '_' |
| 395 | 0x5f, | 529 | 0x0, 0x2, // non-terminal node |
| 396 | 0x0, | 530 | 0x5f, 0x6d, 0x68, 0x5f, 0x65, 0x78, 0x65, 0x63, 0x75, 0x74, // edge '_mh_execute_header' |
| 397 | 0x5, | 531 | 0x65, 0x5f, 0x68, 0x65, 0x61, 0x64, 0x65, 0x72, 0x0, 0x21, // edge '_mh_execute_header' |
| 398 | 0x0, | 532 | 0x6d, 0x61, 0x69, 0x6e, 0x0, 0x25, // edge 'main' |
| 399 | 0x2, | 533 | 0x2, 0x0, 0x0, 0x0, // terminal node |
| 400 | 0x5f, | 534 | 0x3, 0x0, 0x80, 0x20, 0x0, // terminal node |
| 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, | | |
| 435 | }; | 535 | }; |
| 436 | | 536 | |
| 437 | testing.expect(buffer.items.len == exp_buffer.len); | 537 | var buffer = try gpa.alloc(u8, trie.size); |
| 438 | testing.expect(mem.eql(u8, buffer.items, exp_buffer[0..])); | 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 | |
| | 554 | test "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)); |
| 439 | } | 584 | } |
Jakub Konka <kubkon@jakubkonka.com> 2020-12-10 08:32:59+01:00
GitHub <noreply@github.com> 2020-12-10 08:32:59+01:00
Signed by PGP key 4AEE18F83AFDEB23