vig/src/ast_printer.zig

const std = @import("std");
const tokens = @import("tokens.zig");
const Token = tokens.Token;

usingnamespace @import("ast.zig");

const warn = std.debug.warn;

fn printIdent(ident: usize) void {
    var i: usize = 0;
    while (i < ident) : (i += 1) {
        std.debug.warn("\t", .{});
    }
}

fn print(ident: usize, comptime fmt: []const u8, args: var) void {
    printIdent(ident);
    std.debug.warn(fmt, args);
}

fn printBlock(ident: usize, block: var, endNewline: bool) void {
    std.debug.warn("(\n", .{});

    for (block.toSlice()) |stmt| {
        printIdent(ident);
        printStmt(ident, stmt);
        std.debug.warn("\n", .{});
    }

    if (endNewline) {
        print(ident - 1, ")\n", .{});
    } else {
        print(ident - 1, ")", .{});
    }
}

pub fn printNode(node: *Node, ident: usize) void {
    switch (node.*) {
        .FnDecl => |decl| {
            const name = decl.func_name.lexeme;

            printIdent(ident);

            const ret_type = decl.return_type.lexeme;

            if (decl.method) |method| {
                const vari = method.variable.lexeme;
                const typ = method.typ.lexeme;

                if (method.mutable) {
                    warn("(method mut {} {} {} {} ", .{ vari, typ, name, ret_type });
                } else {
                    warn("(method {} {} {} {} ", .{ vari, typ, name, ret_type });
                }
            } else {
                warn("(fn {} {} (", .{ name, ret_type });
            }

            for (decl.params.toSlice()) |param| {
                warn("({} {}) ", .{ param.name.lexeme, param.typ.lexeme });
            }

            printBlock(ident + 1, decl.body, false);
            warn(")\n", .{});
        },

        .ConstDecl => |consts| {
            print(ident, "(const (\n", .{});

            for (consts.toSlice()) |const_decl| {
                print(ident + 1, "({} ", .{
                    const_decl.name.lexeme,
                });

                printExpr(const_decl.expr);
                std.debug.warn(")\n", .{});
            }

            print(ident, "))\n", .{});
        },

        .Enum => |decl| {
            print(ident, "(enum {} (\n", .{decl.name.lexeme});

            for (decl.fields.toSlice()) |field| {
                print(ident + 1, "{}\n", .{
                    field.lexeme,
                });
            }

            print(ident, "))\n", .{});
        },

        .Root => {
            for (node.Root.toSlice()) |child| {
                printNode(child, ident + 1);
            }
        },

        .Stmt => |stmt| {
            printIdent(ident);
            printStmt(ident, stmt);
            std.debug.warn("\n", .{});
        },

        .Struct => |struc| {
            print(ident, "(struct {} (\n", .{struc.name.lexeme});
            for (struc.fields.toSlice()) |field| {
                printIdent(ident + 1);
                if (field.mutable) {
                    std.debug.warn("(mut ", .{});
                } else {
                    std.debug.warn("(", .{});
                }

                if (field.public) {
                    std.debug.warn("pub ", .{});
                }

                if (field.mutable_outside) {
                    std.debug.warn("MUT_OUT ", .{});
                }

                std.debug.warn("{} {})\n", .{ field.name.lexeme, field.typ.lexeme });
            }
            print(ident, "))\n", .{});
        },

        else => {
            print(ident, "unknown node: {}\n", .{node});
        },
    }
}

fn parenthetize(name: []const u8, exprs: []*const Expr) void {
    std.debug.warn("({}", .{name});

    for (exprs) |expr| {
        std.debug.warn(" ", .{});
        printExpr(expr);
    }

    std.debug.warn(")", .{});
}

pub fn printExpr(expr: *const Expr) void {
    switch (expr.*) {
        .Binary => |binary| parenthetize(binary.op.lexeme, &[_]*const Expr{ binary.left, binary.right }),
        .Logical => |binary| parenthetize(binary.op.lexeme, &[_]*const Expr{ binary.left, binary.right }),
        .Unary => |unary| parenthetize(unary.op.lexeme, &[_]*const Expr{unary.right}),
        .Grouping => |expr_ptr| parenthetize("group", &[_]*const Expr{expr_ptr}),

        .Literal => |literal| {
            switch (literal) {
                .Bool => |val| std.debug.warn("{}", .{val}),
                .Integer => |val| std.debug.warn("{}", .{val}),
                .Float => |val| std.debug.warn("{}", .{val}),
                .String => |val| std.debug.warn("'{}'", .{val}),
                .Array => |exprs| {
                    parenthetize("array", exprs.toSlice());
                },
                else => |typ| std.debug.warn("UnknownLiteral-{}", .{typ}),
            }
        },

        .Variable => |token| std.debug.warn("{}", .{token.lexeme}),

        .VarDecl => |decl| {
            if (decl.mutable) {
                std.debug.warn("(mut ", .{});
            } else {
                std.debug.warn("(", .{});
            }

            std.debug.warn("let {} ", .{decl.assign.name.lexeme});
            printExpr(decl.assign.value);
            std.debug.warn(")", .{});
        },

        .Assign => |assign| {
            std.debug.warn("(set ", .{});
            std.debug.warn("{} ", .{assign.name.lexeme});
            printExpr(assign.value);
            std.debug.warn(")", .{});
        },

        .Call => |call| {
            std.debug.warn("(", .{});
            printExpr(call.callee);

            for (call.arguments.toSlice()) |arg| {
                std.debug.warn(" ", .{});
                printExpr(arg);
            }

            std.debug.warn(")", .{});
        },

        .Struct => |val| {
            std.debug.warn("({} (", .{val.name.lexeme});

            for (val.inits.toSlice()) |init| {
                std.debug.warn(" ({} ", .{init.field.lexeme});
                printExpr(init.expr);
                std.debug.warn(")", .{});
            }

            std.debug.warn("))", .{});
        },

        .Get => |get| {
            warn("(", .{});
            printExpr(get.struc);
            warn("{}", .{get.name.lexeme});
        },

        .Set => |set| {
            warn("(set ", .{});
            printExpr(set.struc);
            warn(" {} ", .{set.field.lexeme});
            printExpr(set.value);
            warn(")", .{});
        },

        else => std.debug.warn("UnknownExpr-{}", .{@tagName(expr.*)}),
    }
}

pub fn printStmt(ident: usize, stmt: *Stmt) void {
    switch (stmt.*) {
        .Println => |expr| parenthetize("println", &[_]*Expr{expr}),
        .Expr => |expr| printExpr(expr),

        .If => |ifstmt| {
            std.debug.warn("(if ", .{});
            printExpr(ifstmt.condition);
            std.debug.warn(" ", .{});

            printBlock(ident + 1, ifstmt.then_branch, false);
            if (ifstmt.else_branch) |else_branch| {
                std.debug.warn(" else ", .{});
                printBlock(ident + 1, else_branch, false);
            }

            std.debug.warn(")\n", .{});
        },

        .Loop => |loop| {
            std.debug.warn("(loop ", .{});
            if (loop.condition) |cond| {
                printExpr(cond);
            } else {
                std.debug.warn("true", .{});
            }
            std.debug.warn(" ", .{});

            printBlock(ident + 1, loop.then_branch, false);
            std.debug.warn(")\n", .{});
        },

        .For => |forstmt| {
            std.debug.warn("(for ", .{});

            if (forstmt.index) |index| {
                std.debug.warn("({} {}) ", .{ index.lexeme, forstmt.value.lexeme });
            } else {
                std.debug.warn("{} ", .{forstmt.value.lexeme});
            }

            std.debug.warn("{} ", .{forstmt.array.lexeme});

            printBlock(ident + 1, forstmt.block, false);
            std.debug.warn(")\n", .{});
        },

        .Return => |ret| {
            std.debug.warn("(return ", .{});
            printExpr(ret.value);
            std.debug.warn(")\n", .{});
        },

        else => std.debug.warn("UnknownStmt-{}", .{@tagName(stmt.*)}),
    }
}