trying to get familiar with ASTs in Zig again...

Ast.zig

//! Abstract Syntax Tree for yanais source code

const std = @import("std");
const assert = std.debug.assert;
const testing = std.testing;
const Allocator = std.mem.Allocator;

nodes: NodeList.Slice,
extra_data: []const u8,

pub const NodeList = std.MultiArrayList(Node);

pub const Location = struct {
    fileid: u32,
    offset: u32,
};

pub const SubDecLocation = struct {
    line: u32,
    column: u32,
    line_start: u32,
    line_length: u32,
};

pub fn deinit(tree: *Ast, allocator: Allocator) void {
    tree.nodes.deinit(allocator);
    tree.* = undefined;
}

pub const RenderError = error {
    OutOfMemory,
};

pub const Node = struct {
    tag: Tag,
    data: Data,

    // keep location separately, we rarely access it
    // loc: Location,

    pub const Index = u32;

    comptime {
        assert(@sizeOf(Tag) == 1);
    }

    pub const Tag = enum {
        // *
        ty_ty,

        // x: LHS_TY -> RHS_TY(x)
        abs_ty,

        // x -> RHS_VALUE(x)
        // LHS = bounds exceeding capture
        abs,

        // LHS = de-Bruijn index
        bound,

        // invoke(LHS) with(RHS)
        apply,

        // free(pos = LHS, len = RHS)
        free,
    };

    pub const Data = struct {
        lhs: Index = 0,
        rhs: Index = 0,
    };
};

const Ast = @This();

pub const EvalStackEntry = union(enum) {
    ty_ty: void,
    abs_ty: LnR,
    abs: Abs,
    free: []const u8,
    free_app: FreeApp,

    pub const LnR = struct {
        lhs: *EvalStackEntry,
        rhs: *EvalStackEntry,

        pub fn deinit(self: *LnR, allocator: Allocator) void {
            self.lhs.deinit(allocator);
            allocator.destroy(self.lhs);
            self.rhs.deinit(allocator);
            allocator.destroy(self.rhs);
        }
        pub fn clone(self: LnR, allocator: Allocator) error{OutOfMemory}!LnR {
            var lhs = try allocator.create(EvalStackEntry);
            errdefer allocator.destroy(lhs);
            lhs.* = try self.lhs.clone(allocator);
            var rhs = try allocator.create(EvalStackEntry);
            errdefer allocator.destroy(rhs);
            rhs.* = try self.rhs.clone(allocator);
            return .{ .lhs = lhs, .rhs = rhs };
        }
    };

    pub const Abs = struct {
        code: Node.Index,
        capture: []EvalStackEntry,

        pub fn deinit(self: *Abs, allocator: Allocator) void {
            for (self.capture) |*item| item.deinit(allocator);
            allocator.free(self.capture);
            self.* = undefined;
        }

        pub fn clone(self: Abs, allocator: Allocator) error{OutOfMemory}!Abs {
            const capture = try allocator.alloc(EvalStackEntry, self.capture.len);
            var already_init: usize = 0;
            errdefer {
                for (capture[0..already_init]) |*item| item.deinit(allocator);
                allocator.free(capture);
            }
            for (self.capture, 0..) |*item, i| {
                capture[i] = try item.clone(allocator);
                already_init += 1;
            }
            return .{ .code = self.code, .capture = capture };
        }
    };

    pub const FreeApp = struct {
        name: []const u8,
        args: []EvalStackEntry,

        pub fn deinit(self: *FreeApp, allocator: Allocator) void {
            for (self.args) |*item| item.deinit(allocator);
            allocator.free(self.args);
            self.* = undefined;
        }

        pub fn clone(self: FreeApp, allocator: Allocator) error{OutOfMemory}!FreeApp {
            var args = try allocator.alloc(EvalStackEntry, self.args.len);
            var already_init: usize = 0;
            errdefer {
                for (args[0..already_init]) |*item| item.deinit(allocator);
                allocator.free(args);
            }
            for (self.args, 0..) |*item, i| {
                args[i] = try item.clone(allocator);
                already_init += 1;
            }
            return .{ .name = self.name, .args = args };
        }

        // yes I know this is pretty inefficient
        pub fn addArg(self: *FreeApp, allocator: Allocator, item: EvalStackEntry) error{OutOfMemory}!void {
            var args = std.ArrayListUnmanaged(EvalStackEntry).fromOwnedSlice(self.args);
            const res = args.append(allocator, item);
            self.args = if (args.toOwnedSlice(allocator)) |x| x else |_| @panic("unrecoverable memory corruption");
            return res;
        }
    };

    pub fn deinit(self: *EvalStackEntry, allocator: Allocator) void {
        switch (self.*) {
            .ty_ty => {},
            .abs_ty => |*abs| abs.deinit(allocator),
            .abs => |*abs| abs.deinit(allocator),
            .free => |_| {},
            .free_app => |*app| app.deinit(allocator),
        }
        self.* = undefined;
    }

    pub fn clone(self: EvalStackEntry, allocator: Allocator) error{OutOfMemory}!EvalStackEntry {
        return switch (self) {
            .ty_ty => .ty_ty,
            .abs_ty => |abs| .{ .abs_ty = try abs.clone(allocator) },
            .abs => |abs| .{ .abs = try abs.clone(allocator) },
            .free => |f| .{ .free = f },
            .free_app => |app| .{ .free_app = try app.clone(allocator) },
        };
    }
};

pub const EvalStack = std.ArrayList(EvalStackEntry);

pub fn eval(tree: *const Ast, stack: *EvalStack, sel: usize) !void {
    const allocator = stack.allocator;
    assert(sel < tree.nodes.len);

    // std.debug.print("eval @ {d} = {any}\n", .{ sel, tree.nodes.get(sel) });

    switch (tree.nodes.items(.tag)[sel]) {
        .ty_ty => {
            try stack.append(.ty_ty);
        },
        .abs_ty => {
            const dat = tree.nodes.items(.data)[sel];
            try eval(tree, stack, dat.lhs);
            try eval(tree, stack, dat.rhs);

            const lhs = try allocator.create(EvalStackEntry);
            errdefer allocator.destroy(lhs);
            const rhs = try allocator.create(EvalStackEntry);
            rhs.* = stack.pop();
            const lptr = &stack.items[stack.items.len - 1];
            lhs.* = lptr.*;

            lptr.* = .{ .abs_ty = .{ .lhs = lhs, .rhs = rhs } };
        },
        .abs => {
            const dat = tree.nodes.items(.data)[sel];
            const capture = try allocator.alloc(EvalStackEntry, @as(usize, dat.lhs));
            try stack.ensureUnusedCapacity(1);
            var already_init: usize = 0;
            errdefer {
                for (capture[0..already_init]) |*item| { item.deinit(allocator); }
                allocator.free(capture);
            }
            if (capture.len != 0) {
                const lpidx = stack.items.len - 1;
                for (capture, 0..) |*item, i| {
                    item.* = try stack.items[lpidx - i].clone(allocator);
                    already_init += 1;
                }
            }

            stack.appendAssumeCapacity(.{ .abs = .{ .code = dat.rhs, .capture = capture } });
        },
        .bound => {
            const binder = tree.nodes.items(.data)[sel].lhs;
            if (binder < stack.items.len) {
                try stack.append(try stack.items[stack.items.len - binder - 1].clone(allocator));
            } else {
                return error.InvalidBound;
            }
        },
        .apply => {
            const dat = tree.nodes.items(.data)[sel];

            try eval(tree, stack, dat.rhs);
            var rhs = stack.pop();
            errdefer rhs.deinit(allocator);

            try eval(tree, stack, dat.lhs);
            var lhs = stack.pop();
            defer lhs.deinit(allocator);

            switch (lhs) {
                .abs => |*abs| {
                    const truncpos = stack.items.len;
                    const trunclen = abs.capture.len + 1;

                    // move items from capture stack to normal one
                    try stack.ensureUnusedCapacity(trunclen);
                    stack.appendSliceAssumeCapacity(abs.capture);
                    stack.appendAssumeCapacity(rhs);

                    // prevent double-free
                    allocator.free(abs.capture);
                    abs.capture = &[_]EvalStackEntry{};

                    try eval(tree, stack, abs.code);
                    // get rid of now unused stack items
                    for (stack.items[truncpos..truncpos + trunclen]) |*item| { item.deinit(allocator); }
                    try stack.replaceRange(truncpos, trunclen, &[_]EvalStackEntry{});
                },
                .ty_ty, .abs_ty => return error.ApplyInvalidType,
                .free => |name| {
                    const args = try allocator.alloc(EvalStackEntry, 1);
                    errdefer allocator.free(args);
                    args[0] = rhs;
                    try stack.append(.{ .free_app = .{ .name = name, .args = args } });
                },
                .free_app => |*app| {
                    try app.addArg(allocator, rhs);
                    try stack.append(lhs);
                    // prevent double-free
                    lhs = .ty_ty;
                },
            }
        },
        .free => {
            const dat = tree.nodes.items(.data)[sel];
            try stack.append(.{ .free = tree.extra_data[dat.lhs .. dat.lhs + dat.rhs] });
        },
    }
}

pub const std_options = struct {
    const fmt_max_depth = 100;
};

test "simple stack evaluation" {
    const allocator = testing.allocator;

    // create AST
    var nodes_prep = NodeList{};
    try nodes_prep.append(allocator, .{ .tag = .apply, .data = .{ .lhs = 1, .rhs = 8 } });
    try nodes_prep.append(allocator, .{ .tag = .apply, .data = .{ .lhs = 2, .rhs = 7 } });
    try nodes_prep.append(allocator, .{ .tag = .abs, .data = .{ .rhs = 3 } });
    try nodes_prep.append(allocator, .{ .tag = .abs, .data = .{ .lhs = 1, .rhs = 4 } });
    try nodes_prep.append(allocator, .{ .tag = .apply, .data = .{ .lhs = 5, .rhs = 6 } });
    try nodes_prep.append(allocator, .{ .tag = .bound, .data = .{ .lhs = 0 } });
    try nodes_prep.append(allocator, .{ .tag = .bound, .data = .{ .lhs = 1 } });
    try nodes_prep.append(allocator, .{ .tag = .free, .data = .{ .lhs = 0, .rhs = 3 } });
    try nodes_prep.append(allocator, .{ .tag = .free, .data = .{ .lhs = 3, .rhs = 3 } });

    var ast: Ast = .{
        .nodes = nodes_prep.slice(),
        .extra_data = &.{ 1, 2, 3, 4, 5, 6 },
    };
    defer ast.deinit(allocator);

    // setup stack
    var stack = std.ArrayList(EvalStackEntry).init(allocator);
    defer {
        for (stack.items) |*item| { item.deinit(allocator); }
        stack.deinit();
    }

    try eval(&ast, &stack, 0);
    std.debug.print("Resulting stack:\n", .{});
    for (stack.items) |item| {
      std.debug.print("- {any}\n", .{item});
    }
}