deimos-lang/dmc-lib/src/ast/compilation_unit.rs

use crate::ast::class::Class;
use crate::ast::extern_function::ExternFunction;
use crate::ast::fqn_context::FqnContext;
use crate::ast::function::Function;
use crate::ast::helpers::{collect_diagnostics_into_mut, try_insert_symbols_into};
use crate::diagnostic::Diagnostic;
use crate::ir::ir_class::IrClass;
use crate::ir::ir_function::IrFunction;
use crate::symbol_table::SymbolTable;
use crate::types_table::TypesTable;
use crate::{diagnostics_result, handle_diagnostics};

pub struct CompilationUnit {
    functions: Vec<Function>,
    extern_functions: Vec<ExternFunction>,
    classes: Vec<Class>,
}

impl CompilationUnit {
    pub fn new(
        functions: Vec<Function>,
        extern_functions: Vec<ExternFunction>,
        classes: Vec<Class>,
    ) -> Self {
        Self {
            functions,
            extern_functions,
            classes,
        }
    }

    pub fn functions(&self) -> &[Function] {
        &self.functions
    }

    pub fn extern_functions(&self) -> &[ExternFunction] {
        &self.extern_functions
    }

    pub fn classes(&self) -> &[Class] {
        &self.classes
    }

    pub fn init_scopes(&mut self, symbol_table: &mut SymbolTable) {
        let compilation_unit_scope = symbol_table.push_module_scope("compilation_unit_scope");
        for class in &mut self.classes {
            class.init_scopes(symbol_table, compilation_unit_scope);
        }
        for function in &mut self.functions {
            function.init_scopes(symbol_table, compilation_unit_scope);
        }
        for extern_function in &mut self.extern_functions {
            extern_function.init_scopes(symbol_table, compilation_unit_scope);
        }
        symbol_table.pop_scope();
    }

    pub fn gather_symbols_into(
        &self,
        symbol_table: &mut SymbolTable,
    ) -> Result<(), Vec<Diagnostic>> {
        let mut diagnostics = vec![];
        let mut fqn_context = FqnContext::new();
        for class in &self.classes {
            handle_diagnostics!(
                try_insert_symbols_into(class.make_symbols(&mut fqn_context), symbol_table),
                diagnostics
            );
        }
        for function in &self.functions {
            let (_, symbols) = function.make_symbols(&mut fqn_context, false);
            handle_diagnostics!(try_insert_symbols_into(symbols, symbol_table), diagnostics);
        }
        for extern_function in &self.extern_functions {
            let (_, symbols) = extern_function.make_symbols(&mut fqn_context);
            handle_diagnostics!(try_insert_symbols_into(symbols, symbol_table), diagnostics);
        }
        diagnostics_result!(diagnostics)
    }

    pub fn check_names(&self, symbol_table: &mut SymbolTable) -> Result<(), Vec<Diagnostic>> {
        let mut diagnostics = vec![];
        for class in &self.classes {
            diagnostics.append(&mut class.check_names(symbol_table));
            diagnostics.append(&mut class.check_field_initializer_names(symbol_table));
            diagnostics.append(&mut class.analyze_local_names(symbol_table));
        }
        for function in &self.functions {
            diagnostics.append(&mut function.check_names(symbol_table));
            diagnostics.append(&mut function.analyze_static_fn_local_names(symbol_table));
        }
        for extern_function in &self.extern_functions {
            diagnostics.append(&mut extern_function.check_names(symbol_table));
        }
        diagnostics_result!(diagnostics)
    }

    pub fn type_check(
        &mut self,
        symbol_table: &SymbolTable,
        types_table: &mut TypesTable,
    ) -> Result<(), Vec<Diagnostic>> {
        let mut diagnostics: Vec<Diagnostic> = vec![];

        collect_diagnostics_into_mut(
            &mut self.functions,
            |f| f.type_check(symbol_table, types_table),
            &mut diagnostics,
        );

        collect_diagnostics_into_mut(
            &mut self.extern_functions,
            |ef| ef.type_check(symbol_table),
            &mut diagnostics,
        );

        collect_diagnostics_into_mut(
            &mut self.classes,
            |c| c.type_check(symbol_table, types_table),
            &mut diagnostics,
        );

        diagnostics_result!(diagnostics)
    }

    pub fn to_ir(
        &self,
        symbol_table: &SymbolTable,
        types_table: &TypesTable,
    ) -> (Vec<IrClass>, Vec<IrFunction>) {
        let mut functions: Vec<IrFunction> = vec![];
        let mut classes: Vec<IrClass> = vec![];

        self.functions
            .iter()
            .map(|f| f.to_ir(symbol_table, types_table, None))
            .for_each(|f| functions.push(f));

        for class in &self.classes {
            let (class, mut class_functions) = class.to_ir(symbol_table, types_table);
            functions.append(&mut class_functions);
            classes.push(class);
        }

        (classes, functions)
    }
}