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, insert_declared_types_into, insert_resolved_types_into,
};
use crate::ast::{FunctionReturnTypes, NodesToSymbols, NodesToTypes, SymbolsToTypes};
use crate::compile_pipeline::FileId;
use crate::diagnostic::{Diagnostic, Diagnostics};
use crate::ir::ir_class::IrClass;
use crate::ir::ir_function::IrFunction;
use crate::symbol::Symbol;
use crate::symbol_table::SymbolTable;
use crate::symbol_table::util::try_insert_symbols_into;
use crate::types_table::TypesTable;
use crate::{diagnostics_result, handle_diagnostics};

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

impl CompilationUnit {
    pub fn new(
        file_id: Option<FileId>,
        functions: Vec<Function>,
        extern_functions: Vec<ExternFunction>,
        classes: Vec<Class>,
    ) -> Self {
        Self {
            file_id,
            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 declared_symbols(&self) -> Vec<Symbol> {
        let fqn_context = FqnContext::new();
        [
            self.classes
                .iter()
                .flat_map(|class| class.declared_symbols(&fqn_context))
                .collect::<Vec<_>>(),
            self.extern_functions
                .iter()
                .flat_map(|function| function.declared_symbols(&fqn_context).1)
                .collect(),
            self.functions
                .iter()
                .flat_map(|function| function.declared_symbols(&fqn_context, false).1)
                .collect(),
        ]
        .into_iter()
        .flatten()
        .collect()
    }

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

    pub fn resolve_names(&self, symbol_table: &mut SymbolTable) -> (NodesToSymbols, Diagnostics) {
        let mut diagnostics = Diagnostics::new();
        let mut names_table = NodesToSymbols::new();

        for function in &self.functions {
            let (ns, mut ds) = function.resolve_names_static(symbol_table);
            for (node_id, symbol) in ns.into_iter() {
                names_table.insert(node_id, symbol);
            }
            diagnostics.append(&mut ds);
        }

        for extern_function in &self.extern_functions {
            let (ns, mut ds) = extern_function.resolve_names_static(symbol_table);
            for (node_id, symbol) in ns {
                names_table.insert(node_id, symbol);
            }
            diagnostics.append(&mut ds);
        }

        for class in &self.classes {
            let (ns, mut ds) = class.resolve_names(symbol_table);
            for (node_id, symbol) in ns {
                names_table.insert(node_id, symbol);
            }
            diagnostics.append(&mut ds);
        }

        (names_table, diagnostics)
    }

    #[deprecated]
    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)
    }

    /// Associate each declared symbol with a TypeInfo.
    pub fn declared_types(&self, names_table: &NodesToSymbols) -> (SymbolsToTypes, Diagnostics) {
        let mut diagnostics = Diagnostics::new();
        let mut declared_types = SymbolsToTypes::new();

        for function in &self.functions {
            let (dts, mut ds) = function.declared_types(names_table);
            insert_declared_types_into(dts, &mut declared_types);
            diagnostics.append(&mut ds);
        }

        (declared_types, diagnostics)
    }

    /// Resolve types of all nodes that have an implicit (perhaps not declared) type, checking that
    /// things are assignable, etc., along the way.
    pub fn resolve_types(
        &self,
        names_table: &NodesToSymbols,
        symbols_to_types: &SymbolsToTypes,
    ) -> (SymbolsToTypes, NodesToTypes, Diagnostics) {
        let mut diagnostics = Diagnostics::new();
        let mut resolved_types = NodesToTypes::new();
        let mut symbols_to_types = symbols_to_types.clone();

        for function in &self.functions {
            let (sts, nts, mut ds) = function.resolve_types(names_table, &symbols_to_types);
            insert_declared_types_into(sts, &mut symbols_to_types);
            insert_resolved_types_into(nts, &mut resolved_types);
            diagnostics.append(&mut ds);
        }

        (symbols_to_types, resolved_types, diagnostics)
    }

    #[deprecated]
    pub fn gather_types_into(
        &self,
        symbol_table: &SymbolTable,
        types_table: &mut TypesTable,
    ) -> Result<(), Vec<Diagnostic>> {
        let mut diagnostics = Vec::new();

        for class in &self.classes {
            handle_diagnostics!(class.gather_types(symbol_table, types_table), diagnostics);
        }

        for function in &self.functions {
            function.gather_types(symbol_table, types_table);
        }

        for extern_function in &self.extern_functions {
            extern_function.gather_types(symbol_table, types_table);
        }

        diagnostics_result!(diagnostics)
    }

    #[deprecated]
    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, types_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 lower(
        &self,
        nodes_to_symbols: &NodesToSymbols,
        symbols_to_types: &SymbolsToTypes,
        nodes_to_types: &NodesToTypes,
        function_return_types: &FunctionReturnTypes,
    ) -> (Vec<IrClass>, Vec<IrFunction>) {
        let mut ir_classes = Vec::new();
        let mut ir_functions = Vec::new();

        for function in &self.functions {
            ir_functions.push(function.lower_static(
                nodes_to_symbols,
                symbols_to_types,
                nodes_to_types,
                function_return_types,
            ));
        }

        (ir_classes, ir_functions)
    }

    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)
    }
}