rayoko/src/ast_printer.zig

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

usingnamespace @import("ast.zig");
usingnamespace @import("comp_ctx.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: ...) 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: *const 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;

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

            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| {
                print(ident + 1, "({} {})\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(")");
}

fn printTwoExprs(expr_a: *const Expr, expr_b: *const Expr) void {
    std.debug.warn(" ");
    printExpr(expr_a);
    std.debug.warn(" ");
    printExpr(expr_b);
}

const operator_tokens = [_][]const u8{
    "+", "-", "*", "/", "%", ">", ">=", "<", "<=", "==", "&&", "||",
};

const operator_values = [_]BinaryOperator{
    .Add,
    .Sub,
    .Mul,
    .Div,
    .Mod,

    .Greater,
    .GreaterEqual,
    .Less,
    .LessEqual,

    .Equal,
    .And,
    .Or,
};

fn binOpToStr(op: BinaryOperator) ?[]const u8 {
    inline for (operator_values) |val, idx| {
        if (val == op) return operator_tokens[idx];
    }

    return null;
}

fn printBinOp(inner: var) void {
    std.debug.warn("({}", binOpToStr(inner.op));
    printTwoExprs(inner.left, inner.right);
    std.debug.warn(")");
}

const unary_operator_tokens = [_][]const u8{
    "!", "-",
};

const unary_operators = [_]UnaryOperator{
    .Not, .Negate,
};

fn unOpToStr(op: UnaryOperator) ?[]const u8 {
    inline for (unary_operators) |val, idx| {
        if (val == op) return unary_operator_tokens[idx];
    }

    return null;
}

fn printSingleOp(op: UnaryOperator, applied: *const Expr) void {
    printSimpleOp(unOpToStr(op), applied);
}

fn printSimpleOp(op: ?[]const u8, applied: *const Expr) void {
    std.debug.warn("({}", op);
    printExpr(applied);
    std.debug.warn(")");
}

pub fn printExpr(expr: *const Expr) void {
    switch (expr.*) {
        .Binary => |binary| printBinOp(binary),

        .Unary => |unary| printSingleOp(unary.op, unary.right),
        .Grouping => |expr_ptr| printSimpleOp("group", expr_ptr),

        .Literal => |literal| {
            switch (literal) {
                .Bool => |val| std.debug.warn("{}", val),
                .Integer32 => |val| std.debug.warn("{}", val),
                .Integer64 => |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),

        .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.target);
            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: *const Stmt) void {
    switch (stmt.*) {
        .Println => |expr| printSimpleOp("println", expr),
        .Expr => |expr| printExpr(expr),

        .VarDecl => |decl| {
            std.debug.warn("(let {} ", decl.name.lexeme);
            printExpr(decl.value);
            std.debug.warn(")");
        },

        .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.*)),
    }
}

// very bad but be like that
fn retWithName(prefix: []const u8, inner: []const u8) []const u8 {
    var ret_nam_buf = std.heap.direct_allocator.alloc(u8, 256) catch unreachable;
    return std.fmt.bufPrint(ret_nam_buf[0..], "{}({})", prefix, inner) catch unreachable;
}

fn prettyType(typ: SymbolUnderlyingType) []const u8 {
    return switch (typ) {
        .Integer32 => "i32",
        .Integer64 => "i64",
        .Bool => "bool",
        .Double => "double",

        .OpaqueType => |ident| retWithName("opaque", ident),
        .Struct => |ident| retWithName("struct", ident),
        .Enum => |ident| retWithName("enum", ident),
    };
}

pub fn printScope(scope: *Scope, ident: usize) void {
    print(ident, "scope '{}' at addr {}\n", scope.id, @ptrToInt(scope));

    var it = scope.env.iterator();
    while (it.next()) |kv| {
        print(ident + 1, "sym: {}, typ: {}\n", kv.key, prettyType(kv.value));
    }

    for (scope.children.toSlice()) |child| {
        printScope(child, ident + 1);
    }
}

pub fn printContext(ctx: CompilationContext) void {
    var it = ctx.symbol_table.iterator();

    while (it.next()) |kv| {
        switch (kv.value.*) {
            .Function => |fn_sym| {
                std.debug.warn(
                    "function {} returns {}\n",
                    kv.key,
                    prettyType(fn_sym.return_type),
                );

                for (fn_sym.decl.params.toSlice()) |param| {
                    var param_kv = fn_sym.parameters.get(param.name.lexeme).?;
                    std.debug.warn(
                        "\tparameter {} typ {}\n",
                        param_kv.key,
                        prettyType(param_kv.value.typ),
                    );
                }

                // go through scopes
                std.debug.warn("scope info:\n");
                printScope(fn_sym.scope, 1);
            },

            .Struct => |typemap| {
                std.debug.warn("struct '{}'\n", kv.key);
                var map_it = typemap.iterator();
                while (map_it.next()) |map_kv| {
                    std.debug.warn(
                        "\tfield {} type {}\n",
                        map_kv.key,
                        prettyType(map_kv.value),
                    );
                }
            },

            .Variable => std.debug.warn(
                "variable {} type {}\n",
                kv.key,
                kv.value,
            ),

            .Enum => |identmap| {
                std.debug.warn("enum {}:", kv.key);

                var mapit = identmap.iterator();

                while (mapit.next()) |field_kv| {
                    std.debug.warn("\t{} => {}\n", field_kv.key, field_kv.value);
                }
            },

            .Const => |typ| {
                std.debug.warn("const '{}', typ={}\n", kv.key, prettyType(typ));
            },

            else => {
                std.debug.warn("TODO handle print of {}\n", kv.value);
                unreachable;
            },
        }
    }
}