from-scratch/lang/tests/ast.rs

#[path = "shared/shared.rs"]
mod util;

unsafe extern "C" {
    unsafe fn bump_init();

    unsafe fn tokeniser_init_buf(bytes: *const u8, len: usize) -> ();
}

use util::defs::{parse_expr, parse_func, Ast, AstNode};

fn main() {
    unsafe {
        bump_init();
    }
    println!("Bump allocator initialized.");

    let src = b"3 + 4";

    fn print_ast(src: &[u8], parser: impl FnOnce(&mut Ast)) {
        unsafe {
            tokeniser_init_buf(src.as_ptr(), src.len());
            let mut ast = Ast {
                nodes: util::vec::Vec::new(),
            };
            let expr_id = parser(&mut ast);
            println!("{:#}", &ast);
        };
    }

    // print_ast(b"3 + 4", |ast| unsafe {
    //     parse_expr(ast);
    // });
    // print_ast(b"fn main() -> void { return 1 + 2; }", |ast| unsafe {
    //     parse_func(ast);
    // });
    // print_ast(b"fn main() -> void { return (1 + (2)); }", |ast| unsafe {
    //     parse_func(ast);
    // });
    // print_ast(
    //     b"fn main() -> void { return (1 + (2 * 3)) / 4; }",
    //     |ast| unsafe {
    //         parse_func(ast);
    //     },
    // );
    // print_ast(b"fn main() -> void { return 1 + 2 * 3; }", |ast| unsafe {
    //     parse_func(ast);
    // });

    // print_ast(b"fn main() -> void { let x: u32 = 4; }", |ast| unsafe {
    //     parse_func(ast);
    // });
    print_ast(
        b"fn main(a: u32) -> void { let x: u32 = a + 4; }",
        |ast| unsafe {
            parse_func(ast);
        },
    );
}

impl std::fmt::Display for AstNode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        use util::defs::{
            BinaryExpr, AST_ADDRESS_OF, AST_ARG, AST_ASSIGNMENT, AST_BINARY_OP, AST_BLOCK,
            AST_DEREF, AST_FUNCTION, AST_NUMBER, AST_PLACE_TO_VALUE, AST_RETURN_STATEMENT,
            AST_VALUE_TO_PLACE, AST_VAR_DECL, AST_VAR_REF,
        };
        match self.kind {
            AST_NUMBER => {
                write!(f, "Number({})", self.data as usize)
            }
            AST_ARG => {
                let arg = unsafe { self.data.cast::<util::defs::AstArgument>().read() };
                write!(
                    f,
                    "Arg(name: {:?}, arg_type: {:?})",
                    unsafe {
                        std::str::from_utf8(std::slice::from_raw_parts(arg.name, arg.name_len))
                    },
                    arg.arg_type,
                )
            }
            AST_VAR_REF => {
                let var_ref = unsafe { self.data.cast::<util::defs::AstVarRef>().read() };
                if var_ref.resolved != 0 {
                    write!(f, "VarRef({})", var_ref.resolved)
                } else {
                    write!(f, "VarRef(name: {:?})", unsafe {
                        std::str::from_utf8(std::slice::from_raw_parts(
                            var_ref.name,
                            var_ref.name_len,
                        ))
                    },)
                }
            }
            AST_VAR_DECL => {
                let var_decl = unsafe { self.data.cast::<util::defs::AstVarDecl>().read() };
                write!(
                    f,
                    "VarDecl(name: {:?}, var_type: {:?})",
                    unsafe {
                        std::str::from_utf8(std::slice::from_raw_parts(
                            var_decl.name,
                            var_decl.name_len,
                        ))
                    },
                    var_decl.var_type,
                )
            }
            AST_ASSIGNMENT => {
                write!(
                    f,
                    "Assignment(dest: {}, src: {})",
                    self.data as usize, self.extra
                )
            }
            AST_BINARY_OP => {
                let BinaryExpr {
                    left,
                    operator,
                    right,
                } = unsafe { self.data.cast::<util::defs::BinaryExpr>().read() };
                write!(
                    f,
                    "BinaryOp(op: {}, left: {}, right: {})",
                    operator, left, right
                )
            }
            AST_RETURN_STATEMENT => {
                let return_expr_id = self.data as usize;
                write!(f, "ReturnStatement(expr: {})", return_expr_id)
            }
            AST_FUNCTION => {
                let func = unsafe { self.data.cast::<util::defs::AstFunction>().read() };
                write!(
                    f,
                    "Function(name: {:?}, args: {:?}, return_type: {:?}, body: {})",
                    unsafe {
                        std::str::from_utf8(std::slice::from_raw_parts(func.name, func.name_len))
                    },
                    unsafe {
                        std::slice::from_raw_parts(func.args.cast::<u64>(), func.args_len as usize)
                    },
                    func.return_type,
                    func.body
                )
            }
            AST_BLOCK => {
                write!(f, "Block(statements: {:?})", unsafe {
                    std::slice::from_raw_parts(self.data.cast::<u64>(), self.extra as usize)
                })
            }
            AST_PLACE_TO_VALUE => {
                write!(f, "PlaceToValue(place: {})", self.data as usize)
            }
            AST_VALUE_TO_PLACE => {
                write!(f, "ValueToPlace(value: {})", self.data as usize)
            }
            _ => write!(f, "UnknownNode"),
        }
    }
}

impl core::fmt::Display for Ast {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        writeln!(f, "[")?;
        for (i, item) in self.nodes.as_slice().iter().enumerate() {
            if i > 0 {
                writeln!(f, ", ")?;
            }
            write!(f, "\t{i}: {}", item)?;
        }
        writeln!(f, "\n]")
    }
}