const std = @import("std");
const tree = @import("tree.zig");
const Arena = @import("arena.zig").Arena;

pub const ReduceError = error{
    FuelExhausted,
    InvalidApply,
    OutOfMemory,
};

/// Reduce a term to weak head normal form.
pub fn reduce(root: u32, arena: *Arena, fuel: u64) ReduceError!u32 {
    var remaining = fuel;
    return try whnf(root, arena, &remaining);
}

fn whnf(term: u32, arena: *Arena, fuel: *u64) ReduceError!u32 {
    if (fuel.* == 0) return error.FuelExhausted;
    var current = term;

    while (true) {
        switch (arena.get(current).*) {
            .leaf, .stem, .fork => return current,
            .app => |app| {
                const orig = current;
                const func_idx = app.func;
                const arg_idx = app.arg;

                // Reduce function to WHNF
                const f = try whnf(func_idx, arena, fuel);
                if (fuel.* == 0) return error.FuelExhausted;
                fuel.* -= 1;

                switch (arena.get(f).*) {
                    // apply Leaf b = Stem b
                    .leaf => {
                        arena.get(orig).* = .{ .stem = .{ .child = arg_idx } };
                        return orig;
                    },
                    // apply (Stem a) b = Fork a b
                    .stem => |s| {
                        const a = s.child;
                        arena.get(orig).* = .{ .fork = .{ .left = a, .right = arg_idx } };
                        return orig;
                    },
                    .fork => |fork_f| {
                        const left_idx = fork_f.left;
                        const right_idx = fork_f.right;

                        // Reduce left child of Fork
                        const left = try whnf(left_idx, arena, fuel);
                        if (fuel.* == 0) return error.FuelExhausted;
                        fuel.* -= 1;

                        switch (arena.get(left).*) {
                            // apply (Fork Leaf a) _ = a
                            .leaf => {
                                const result = try whnf(right_idx, arena, fuel);
                                if (fuel.* == 0) return error.FuelExhausted;
                                fuel.* -= 1;
                                if (orig != result) {
                                    arena.get(orig).* = arena.get(result).*;
                                }
                                return orig;
                            },
                            // apply (Fork (Stem a) b) c = (a c) (b c)
                            .stem => |s| {
                                const a = s.child;
                                const inner1 = try arena.alloc(.{ .app = .{ .func = a, .arg = arg_idx } });
                                const inner2 = try arena.alloc(.{ .app = .{ .func = right_idx, .arg = arg_idx } });
                                arena.get(orig).* = .{ .app = .{ .func = inner1, .arg = inner2 } };
                                current = orig;
                                if (fuel.* == 0) return error.FuelExhausted;
                                fuel.* -= 1;
                                continue;
                            },
                            .fork => {
                                // Reduce argument
                                const arg = try whnf(arg_idx, arena, fuel);
                                if (fuel.* == 0) return error.FuelExhausted;
                                fuel.* -= 1;

                                switch (arena.get(arg).*) {
                                    // apply (Fork (Fork a b) c) Leaf = a
                                    .leaf => {
                                        const a_idx = arena.get(left).fork.left;
                                        const result = try whnf(a_idx, arena, fuel);
                                        if (fuel.* == 0) return error.FuelExhausted;
                                        fuel.* -= 1;
                                        if (orig != result) {
                                            arena.get(orig).* = arena.get(result).*;
                                        }
                                        return orig;
                                    },
                                    // apply (Fork (Fork a b) c) (Stem u) = b u
                                    .stem => |s| {
                                        const b_idx = arena.get(left).fork.right;
                                        const u = s.child;
                                        arena.get(orig).* = .{ .app = .{ .func = b_idx, .arg = u } };
                                        current = orig;
                                        if (fuel.* == 0) return error.FuelExhausted;
                                        fuel.* -= 1;
                                        continue;
                                    },
                                    // apply (Fork (Fork a b) c) (Fork u v) = (c u) v
                                    .fork => |arg_fork| {
                                        const c_idx = right_idx;
                                        const u = arg_fork.left;
                                        const v = arg_fork.right;
                                        const inner = try arena.alloc(.{ .app = .{ .func = c_idx, .arg = u } });
                                        arena.get(orig).* = .{ .app = .{ .func = inner, .arg = v } };
                                        current = orig;
                                        if (fuel.* == 0) return error.FuelExhausted;
                                        fuel.* -= 1;
                                        continue;
                                    },
                                    .app => return error.InvalidApply,
                                }
                            },
                            .app => return error.InvalidApply,
                        }
                    },
                    .app => return error.InvalidApply,
                }
            },
        }
    }
}