tricu

tree.zig (5990B)

---

 const std = @import("std");
 
 pub const NodeTag = enum(u8) {
     leaf = 0,
     stem = 1,
     fork = 2,
     app = 3,
 };
 
 pub const Node = union(NodeTag) {
     leaf,
     stem: struct { child: u32 },
     fork: struct { left: u32, right: u32 },
     app: struct { func: u32, arg: u32 },
 
     pub fn leafNode() Node {
         return .leaf;
     }
 
     pub fn stemNode(child: u32) Node {
         return .{ .stem = .{ .child = child } };
     }
 
     pub fn forkNode(left: u32, right: u32) Node {
         return .{ .fork = .{ .left = left, .right = right } };
     }
 
     pub fn appNode(func: u32, arg: u32) Node {
         return .{ .app = .{ .func = func, .arg = arg } };
     }
 };
 
 pub const NodePool = struct {
     allocator: std.mem.Allocator,
     nodes: std.ArrayList(Node),
 
     pub fn init(allocator: std.mem.Allocator) NodePool {
         return .{
             .allocator = allocator,
             .nodes = .empty,
         };
     }
 
     pub fn deinit(self: *NodePool) void {
         self.nodes.deinit(self.allocator);
     }
 
     pub fn push(self: *NodePool, node: Node) !u32 {
         const idx: u32 = @intCast(self.nodes.items.len);
         try self.nodes.append(self.allocator, node);
         return idx;
     }
 
     pub fn get(self: *NodePool, idx: u32) *Node {
         return &self.nodes.items[idx];
     }
 
     pub fn len(self: *const NodePool) u32 {
         return @intCast(self.nodes.items.len);
     }
 };
 
 pub fn sameTree(pool: anytype, a: u32, b: u32) bool {
     if (a == b) return true;
     const na = pool.nodes.items[a];
     const nb = pool.nodes.items[b];
     if (@intFromEnum(na) != @intFromEnum(nb)) return false;
     return switch (na) {
         .leaf => true,
         .stem => |sa| sameTree(pool, sa.child, nb.stem.child),
         .fork => |fa| sameTree(pool, fa.left, nb.fork.left) and sameTree(pool, fa.right, nb.fork.right),
         .app => |aa| sameTree(pool, aa.func, nb.app.func) and sameTree(pool, aa.arg, nb.app.arg),
     };
 }
 
 const DstArena = @import("arena.zig").Arena;
 
 /// Iterative deep-copy of a DAG from `src` into `dst`.  Uses an explicit
 /// heap-allocated stack so that very deep (e.g. long list) trees do not
 /// blow the native C stack.  Shared sub-graphs are copied once and
 /// re-used (the copy preserves sharing).
 pub fn copyTree(src: []const Node, dst: *DstArena, root: u32) !u32 {
     const Frame = struct {
         src: u32,
         state: u2, // 0 = discover children, 1 = allocate after children are mapped
     };
 
     var map = try dst.allocator.alloc(u32, src.len);
     defer dst.allocator.free(map);
     @memset(std.mem.sliceAsBytes(map), 0xFF);
 
     var stack: std.ArrayList(Frame) = .empty;
     defer stack.deinit(dst.allocator);
     try stack.append(dst.allocator, .{ .src = root, .state = 0 });
 
     while (stack.items.len > 0) {
         const frame = &stack.items[stack.items.len - 1];
         const src_idx = frame.src;
 
         if (map[src_idx] != 0xFFFFFFFF) {
             _ = stack.pop();
             continue;
         }
 
         if (frame.state == 0) {
             frame.state = 1;
             const node = src[src_idx];
             switch (node) {
                 .leaf => {}, // no children, fall through to allocation next iteration
                 .stem => |s| {
                     if (map[s.child] == 0xFFFFFFFF) {
                         try stack.append(dst.allocator, .{ .src = s.child, .state = 0 });
                     }
                 },
                 .fork => |f| {
                     const need_left = map[f.left] == 0xFFFFFFFF;
                     const need_right = map[f.right] == 0xFFFFFFFF;
                     if (need_right) {
                         try stack.append(dst.allocator, .{ .src = f.right, .state = 0 });
                     }
                     if (need_left) {
                         try stack.append(dst.allocator, .{ .src = f.left, .state = 0 });
                     }
                 },
                 .app => |a| {
                     const need_func = map[a.func] == 0xFFFFFFFF;
                     const need_arg = map[a.arg] == 0xFFFFFFFF;
                     if (need_arg) {
                         try stack.append(dst.allocator, .{ .src = a.arg, .state = 0 });
                     }
                     if (need_func) {
                         try stack.append(dst.allocator, .{ .src = a.func, .state = 0 });
                     }
                 },
             }
         } else {
             // All children mapped; allocate this node in dst.
             const node = src[src_idx];
             const dst_idx = switch (node) {
                 .leaf => try dst.alloc(.leaf),
                 .stem => |s| try dst.alloc(.{ .stem = .{ .child = map[s.child] } }),
                 .fork => |f| try dst.alloc(.{ .fork = .{ .left = map[f.left], .right = map[f.right] } }),
                 .app => |a| try dst.alloc(.{ .app = .{ .func = map[a.func], .arg = map[a.arg] } }),
             };
             map[src_idx] = dst_idx;
             _ = stack.pop();
         }
     }
 
     return map[root];
 }
 
 pub fn formatTree(writer: anytype, pool: anytype, idx: u32, depth: usize) !void {
     if (depth > 200) {
         try writer.writeAll("...");
         return;
     }
     const node = pool.nodes.items[idx];
     switch (node) {
         .leaf => try writer.writeAll("Leaf"),
         .stem => |s| {
             try writer.writeAll("Stem(");
             try formatTree(writer, pool, s.child, depth + 1);
             try writer.writeAll(")");
         },
         .fork => |f| {
             try writer.writeAll("Fork(");
             try formatTree(writer, pool, f.left, depth + 1);
             try writer.writeAll(", ");
             try formatTree(writer, pool, f.right, depth + 1);
             try writer.writeAll(")");
         },
         .app => |a| {
             try writer.writeAll("App(");
             try formatTree(writer, pool, a.func, depth + 1);
             try writer.writeAll(", ");
             try formatTree(writer, pool, a.arg, depth + 1);
             try writer.writeAll(")");
         },
     }
 }