from-scratch/lang/tests/ast.rs

#![feature(debug_closure_helpers)]

#[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) -> u64) {
        unsafe {
            tokeniser_init_buf(src.as_ptr(), src.len());
            let mut ast = Ast {
                nodes: util::vec::Vec::new(),
            };
            let expr_id = parser(&mut ast);
            eprintln!("Parsed expression ID: {}", expr_id);
            println!("{:#}", &ast);

            // unsafe extern "C" fn visit_node(_this: *mut (), ast: *mut Ast, node_id: u64) {
            //     let ast = unsafe { &*ast };
            //     let node = ast.nodes.get(node_id as usize).unwrap();
            //     eprintln!("Visiting node {node_id}: {node}");
            // }

            // util::defs::ast_walk_for_each(&mut ast, expr_id, core::ptr::null_mut(), visit_node);

            let mut symtable = core::mem::MaybeUninit::<util::defs::SymbolTable>::uninit();
            util::defs::ast_build_symtable(&mut ast, expr_id, &mut symtable);
            let symtable = symtable.assume_init();
            use util::DisplayedSliceExt;
            println!(
                "Symbol Table: {:#?}",
                symtable.symtable.as_slice().displayed()
            );
        };
    }

    // 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) },
    );
    print_ast(
        b"fn main(a: u32) -> void {
let x: u32 = a + 4;
let y: *u32 = &x;
return *y;
}",
        |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_DEREF => {
                write!(f, "Deref(expr: {})", self.data as usize)
            }
            AST_ADDRESS_OF => {
                write!(f, "AddressOf(expr: {})", 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]")
    }
}

impl core::fmt::Display for util::defs::SymEntry {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("SymEntry")
            .field_with("key", |f| {
                f.debug_struct("Key")
                    .field_with("kind", |f| {
                        f.write_str(match self.key.kind {
                            util::defs::SYM_KEY_SCOPE => "Scope",
                            util::defs::SYM_KEY_SCOPE_NAME => "ScopeName",
                            util::defs::SYM_KEY_PARENT_SCOPE => "ParentScope",
                            util::defs::SYM_KEY_ARG => "Argument",
                            util::defs::SYM_KEY_VAR => "Variable",
                            _ => "Unknown",
                        })
                    })
                    .field("scope", &self.key.scope_index)
                    .field("span", &self.key.span)
                    .field_with("ident", |f| {
                        f.write_str(unsafe {
                            &core::str::from_utf8_unchecked(core::slice::from_raw_parts(
                                self.key.ident,
                                self.key.ident_len,
                            ))
                        })
                    })
                    .finish()
            })
            .field_with("value", |f| {
                let stct = &mut f.debug_struct("Value");
                if self.extra == 0 {
                    stct.field("ast_index", &self.index).finish()
                } else if self.index != 0 {
                    stct.field_with("ident", |f| {
                        f.write_str(unsafe {
                            core::str::from_utf8_unchecked(core::slice::from_raw_parts(
                                self.index as *const u8,
                                self.extra as usize,
                            ))
                        })
                    })
                    .finish()
                } else {
                    stct.field("index", &self.index)
                        .field("extra", &self.extra)
                        .finish()
                }
            })
            .finish()
    }
}