From ea44d12d1be8eb17a1555f6ab794621da0212171 Mon Sep 17 00:00:00 2001 From: Jakub Konka Date: Wed, 7 Oct 2020 21:19:45 +0200 Subject: [PATCH] Add writeULEB128Mem test and couple fixes Signed-off-by: Jakub Konka --- src/link/MachO/Trie.zig | 177 ++++++++++++++++++++++++++++------------ 1 file changed, 127 insertions(+), 50 deletions(-) diff --git a/src/link/MachO/Trie.zig b/src/link/MachO/Trie.zig index 24b06c8ba2..f7d37315cf 100644 --- a/src/link/MachO/Trie.zig +++ b/src/link/MachO/Trie.zig @@ -34,6 +34,7 @@ const std = @import("std"); const mem = std.mem; const leb = std.debug.leb; const log = std.log.scoped(.link); +const testing = std.testing; const Allocator = mem.Allocator; pub const Symbol = struct { @@ -67,48 +68,33 @@ const Node = struct { self.edges.deinit(alloc); } - fn put(self: *Node, alloc: *Allocator, fromEdge: ?*Edge, prefix: usize, label: []const u8) !*Node { - // Traverse all edges. + fn put(self: *Node, alloc: *Allocator, label: []const u8) !*Node { + // Check for match with edges from this node. for (self.edges.items) |*edge| { - const match = mem.indexOfDiff(u8, edge.label, label) orelse return self; // Got a full match, don't do anything. - if (match - prefix > 0) { - // If we match, we advance further down the trie. - return edge.to.put(alloc, edge, match, label); - } - } + const match = mem.indexOfDiff(u8, edge.label, label) orelse return edge.to; + if (match == 0) continue; + if (match == edge.label.len) return edge.to.put(alloc, label[match..]); - if (fromEdge) |from| { - if (mem.eql(u8, from.label, label[0..prefix])) { - if (prefix == label.len) return self; + // Found a match, need to splice up nodes. + // From: A -> B + // To: A -> C -> B + const mid = try alloc.create(Node); + mid.* = .{}; + const to_label = edge.label; + const to_node = edge.to; + edge.to = mid; + edge.label = label[0..match]; + + try mid.edges.append(alloc, .{ + .from = mid, + .to = to_node, + .label = to_label[match..], + }); + + if (match == label.len) { + return to_node; } else { - // Fixup nodes. We need to insert an intermediate node between - // from.to and self. - // Is: A -> B - // Should be: A -> C -> B - const mid = try alloc.create(Node); - mid.* = .{}; - const to_label = from.label; - from.to = mid; - from.label = label[0..prefix]; - - try mid.edges.append(alloc, .{ - .from = mid, - .to = self, - .label = to_label, - }); - - if (prefix == label.len) return self; // We're done. - - const new_node = try alloc.create(Node); - new_node.* = .{}; - - try mid.edges.append(alloc, .{ - .from = mid, - .to = new_node, - .label = label, - }); - - return new_node; + return mid.put(alloc, label[match..]); } } @@ -148,7 +134,7 @@ const Node = struct { // Write number of edges (max legal number of edges is 256). try buffer.append(alloc, @intCast(u8, self.edges.items.len)); - var node_offset_info: [@sizeOf(u8)]u64 = undefined; + var node_offset_info: [std.math.maxInt(u8)]u64 = undefined; for (self.edges.items) |edge, i| { // Write edges labels leaving out space in-between to later populate // with offsets to each node. @@ -185,7 +171,7 @@ root: Node, /// This operation may change the layout of the trie by splicing edges in /// certain circumstances. pub fn put(self: *Trie, alloc: *Allocator, symbol: Symbol) !void { - const node = try self.root.put(alloc, null, 0, symbol.name); + const node = try self.root.put(alloc, symbol.name); node.offset = symbol.offset; node.export_flags = symbol.export_flags; } @@ -202,9 +188,7 @@ pub fn deinit(self: *Trie, alloc: *Allocator) void { } test "Trie basic" { - const testing = @import("std").testing; var gpa = testing.allocator; - var trie: Trie = .{ .root = .{}, }; @@ -223,7 +207,7 @@ test "Trie basic" { testing.expect(mem.eql(u8, trie.root.edges.items[0].label, "_st")); { - // root --- _st ---> node --- _start ---> node + // root --- _st ---> node --- art ---> node try trie.put(gpa, .{ .name = "_start", .offset = 0, @@ -234,12 +218,12 @@ test "Trie basic" { const nextEdge = &trie.root.edges.items[0]; testing.expect(mem.eql(u8, nextEdge.label, "_st")); testing.expect(nextEdge.to.edges.items.len == 1); - testing.expect(mem.eql(u8, nextEdge.to.edges.items[0].label, "_start")); + testing.expect(mem.eql(u8, nextEdge.to.edges.items[0].label, "art")); } { - // root --- _ ---> node --- _st ---> node --- _start ---> node + // root --- _ ---> node --- st ---> node --- art ---> node // | - // | --- _main ---> node + // | --- main ---> node try trie.put(gpa, .{ .name = "_main", .offset = 0, @@ -250,10 +234,103 @@ test "Trie basic" { const nextEdge = &trie.root.edges.items[0]; testing.expect(mem.eql(u8, nextEdge.label, "_")); testing.expect(nextEdge.to.edges.items.len == 2); - testing.expect(mem.eql(u8, nextEdge.to.edges.items[0].label, "_st")); - testing.expect(mem.eql(u8, nextEdge.to.edges.items[1].label, "_main")); + testing.expect(mem.eql(u8, nextEdge.to.edges.items[0].label, "st")); + testing.expect(mem.eql(u8, nextEdge.to.edges.items[1].label, "main")); const nextNextEdge = &nextEdge.to.edges.items[0]; - testing.expect(mem.eql(u8, nextNextEdge.to.edges.items[0].label, "_start")); + testing.expect(mem.eql(u8, nextNextEdge.to.edges.items[0].label, "art")); } } + +test "Trie.writeULEB128Mem" { + var gpa = testing.allocator; + var trie: Trie = .{ + .root = .{}, + }; + defer trie.deinit(gpa); + + try trie.put(gpa, .{ + .name = "__mh_execute_header", + .offset = 0, + .export_flags = 0, + }); + try trie.put(gpa, .{ + .name = "_main", + .offset = 0x1000, + .export_flags = 0, + }); + + var buffer: std.ArrayListUnmanaged(u8) = .{}; + defer buffer.deinit(gpa); + + try trie.writeULEB128Mem(gpa, &buffer); + + const exp_buffer = [_]u8{ + 0x0, + 0x1, + 0x5f, + 0x0, + 0xc, + 0x0, + 0x0, + 0x0, + 0x0, + 0x0, + 0x0, + 0x0, + 0x0, + 0x2, + 0x5f, + 0x6d, + 0x68, + 0x5f, + 0x65, + 0x78, + 0x65, + 0x63, + 0x75, + 0x74, + 0x65, + 0x5f, + 0x68, + 0x65, + 0x61, + 0x64, + 0x65, + 0x72, + 0x0, + 0x36, + 0x0, + 0x0, + 0x0, + 0x0, + 0x0, + 0x0, + 0x0, + 0x6d, + 0x61, + 0x69, + 0x6e, + 0x0, + 0x3a, + 0x0, + 0x0, + 0x0, + 0x0, + 0x0, + 0x0, + 0x0, + 0x2, + 0x0, + 0x0, + 0x0, + 0x3, + 0x0, + 0x80, + 0x20, + 0x0, + }; + + testing.expect(buffer.items.len == exp_buffer.len); + testing.expect(mem.eql(u8, buffer.items, exp_buffer[0..])); +}