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),
    };
}

/// Deep-copy a term from a source node slice into a destination Arena, returning the new index.
/// Uses recursion; assumes the tree is finite and well-formed.
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 = try dst.allocator.alloc(Frame, src.len);
    defer dst.allocator.free(stack);
    var sp: usize = 0;

    stack[sp] = .{ .src = root, .state = 0 };
    sp += 1;

    while (sp > 0) {
        const frame = &stack[sp - 1];
        const src_idx = frame.src;

        if (map[src_idx] != 0xFFFFFFFF) {
            sp -= 1;
            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) {
                        stack[sp] = .{ .src = s.child, .state = 0 };
                        sp += 1;
                    }
                },
                .fork => |f| {
                    const need_left = map[f.left] == 0xFFFFFFFF;
                    const need_right = map[f.right] == 0xFFFFFFFF;
                    if (need_right) {
                        stack[sp] = .{ .src = f.right, .state = 0 };
                        sp += 1;
                    }
                    if (need_left) {
                        stack[sp] = .{ .src = f.left, .state = 0 };
                        sp += 1;
                    }
                },
                .app => |a| {
                    const need_func = map[a.func] == 0xFFFFFFFF;
                    const need_arg = map[a.arg] == 0xFFFFFFFF;
                    if (need_arg) {
                        stack[sp] = .{ .src = a.arg, .state = 0 };
                        sp += 1;
                    }
                    if (need_func) {
                        stack[sp] = .{ .src = a.func, .state = 0 };
                        sp += 1;
                    }
                },
            }
        } 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;
            sp -= 1;
        }
    }

    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(")");
        },
    }
}