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::diagnostic::Diagnostic;
use crate::ir::ir_class::IrClass;
use crate::ir::ir_function::IrFunction;
use crate::symbol_table::SymbolTable;

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 gather_declared_names(
        &mut self,
        symbol_table: &mut SymbolTable,
    ) -> Result<(), Vec<Diagnostic>> {
        symbol_table.push_module_scope("compilation_unit_scope");

        let mut fqn_context = FqnContext::new(); // in the future, we'll push the pkg/ns on here
        let mut diagnostics: Vec<Diagnostic> = vec![];

        self.functions
            .iter_mut()
            .map(|f| f.gather_declared_names(symbol_table, &fqn_context, None))
            .filter_map(Result::err)
            .flatten()
            .for_each(|diagnostic| diagnostics.push(diagnostic));

        self.extern_functions
            .iter_mut()
            .map(|f| f.gather_declared_names(symbol_table, &fqn_context))
            .filter_map(Result::err)
            .flatten()
            .for_each(|diagnostic| diagnostics.push(diagnostic));

        self.classes
            .iter_mut()
            .map(|c| c.gather_declared_names(symbol_table, &mut fqn_context))
            .filter_map(Result::err)
            .flatten()
            .for_each(|diagnostic| diagnostics.push(diagnostic));

        symbol_table.pop_scope();

        if diagnostics.is_empty() {
            Ok(())
        } else {
            Err(diagnostics)
        }
    }

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

        self.functions
            .iter_mut()
            .map(|f| f.check_name_usages(symbol_table, None))
            .filter_map(Result::err)
            .flatten()
            .for_each(|diagnostic| diagnostics.push(diagnostic));

        self.extern_functions
            .iter_mut()
            .map(|f| f.check_name_usages(symbol_table, None))
            .filter_map(Result::err)
            .flatten()
            .for_each(|diagnostic| diagnostics.push(diagnostic));

        self.classes
            .iter_mut()
            .map(|c| c.check_name_usages(symbol_table))
            .filter_map(Result::err)
            .flatten()
            .for_each(|diagnostic| diagnostics.push(diagnostic));

        if diagnostics.is_empty() {
            Ok(())
        } else {
            Err(diagnostics)
        }
    }

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

        self.functions
            .iter_mut()
            .map(|f| f.type_check(symbol_table))
            .filter_map(Result::err)
            .flatten()
            .for_each(|diagnostic| diagnostics.push(diagnostic));

        self.extern_functions
            .iter_mut()
            .map(|f| f.type_check(symbol_table))
            .filter_map(Result::err)
            .flatten()
            .for_each(|diagnostic| diagnostics.push(diagnostic));

        self.classes
            .iter_mut()
            .map(|c| c.type_check(symbol_table))
            .filter_map(Result::err)
            .flatten()
            .for_each(|diagnostic| diagnostics.push(diagnostic));

        if diagnostics.is_empty() {
            Ok(())
        } else {
            Err(diagnostics)
        }
    }

    pub fn to_ir(&self, symbol_table: &SymbolTable) -> (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, None))
            .for_each(|f| functions.push(f));

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

        (classes, functions)
    }
}