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

use crate::asm::asm_instruction::{AsmInstruction, Return};
use crate::ast::assemble_context::AssembleContext;
use crate::ast::ir_builder::IrBuilder;
use crate::ast::parameter::Parameter;
use crate::ast::statement::Statement;
use crate::ast::type_use::TypeUse;
use crate::constants_table::ConstantsTable;
use crate::diagnostic::Diagnostic;
use crate::ir::ir_function::IrFunction;
use crate::source_range::SourceRange;
use crate::symbol::FunctionSymbol;
use crate::symbol_table::{SymbolInsertError, SymbolTable};
use crate::type_info::TypeInfo;
use std::cell::RefCell;
use std::ops::Neg;
use std::rc::Rc;

pub struct Function {
    declared_name: String,
    declared_name_source_range: SourceRange,
    parameters: Vec<Parameter>,
    return_type: Option<TypeUse>,
    statements: Vec<Statement>,
    function_symbol: Option<Rc<RefCell<FunctionSymbol>>>,
}

impl Function {
    pub fn new(
        declared_name: &str,
        declared_name_source_range: SourceRange,
        parameters: Vec<Parameter>,
        return_type: Option<TypeUse>,
        statements: Vec<Statement>,
    ) -> Self {
        Self {
            declared_name: declared_name.to_string(),
            declared_name_source_range,
            parameters,
            return_type,
            statements,
            function_symbol: None,
        }
    }

    pub fn declared_name(&self) -> &str {
        &self.declared_name
    }

    pub fn statements(&self) -> Vec<&Statement> {
        self.statements.iter().collect()
    }

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

        // insert function symbol
        let insert_result = symbol_table.insert_function_symbol(FunctionSymbol::new(
            self.declared_name(),
            false,
            self.return_type
                .as_ref()
                .map(|return_type| return_type.to_type_info())
                .unwrap_or(TypeInfo::Void),
        ));

        // get function symbol if successful
        let function_symbol = match insert_result {
            Ok(function_symbol) => function_symbol,
            Err(symbol_insert_error) => {
                return match symbol_insert_error {
                    SymbolInsertError::AlreadyDeclared(already_declared) => {
                        diagnostics.push(Diagnostic::new(
                            &format!(
                                "Function {} already declared in current scope",
                                already_declared.name()
                            ),
                            self.declared_name_source_range.start(),
                            self.declared_name_source_range.end(),
                        ));
                        diagnostics
                    }
                };
            }
        };

        // save function symbol for later
        self.function_symbol = Some(function_symbol.clone());

        // handle parameters
        symbol_table.push_scope(&format!("parameter_scope({})", self.declared_name));
        let mut parameter_symbols = vec![];
        let parameter_count = self.parameters.len();
        let stack_frame_offset_base = (parameter_count as isize).neg();
        for (i, parameter) in self.parameters.iter_mut().enumerate() {
            match parameter.gather_declared_names(symbol_table) {
                Ok(parameter_symbol) => {
                    parameter_symbol
                        .borrow_mut()
                        .set_stack_frame_offset(stack_frame_offset_base + (i as isize));
                    parameter_symbols.push(parameter_symbol);
                }
                Err(mut parameter_diagnostics) => {
                    diagnostics.append(&mut parameter_diagnostics);
                }
            }
        }
        function_symbol
            .borrow_mut()
            .set_parameters(parameter_symbols);

        symbol_table.push_scope(&format!("body_scope({})", self.declared_name));
        for statement in &mut self.statements {
            diagnostics.append(&mut statement.gather_declared_names(symbol_table));
        }
        symbol_table.pop_scope(); // body
        symbol_table.pop_scope(); // params
        diagnostics
    }

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

        // parameters
        diagnostics.append(
            &mut self
                .parameters
                .iter_mut()
                .map(|parameter| parameter.check_name_usages(symbol_table))
                .filter_map(|result| result.err())
                .flatten()
                .collect(),
        );

        // statements
        for statement in &mut self.statements {
            diagnostics.append(&mut statement.check_name_usages(symbol_table));
        }
        diagnostics
    }

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

        // parameters
        diagnostics.append(
            &mut self
                .parameters
                .iter_mut()
                .map(|parameter| parameter.type_check(symbol_table))
                .filter_map(|result| result.err())
                .flatten()
                .collect(),
        );

        // statements
        for statement in &mut self.statements {
            diagnostics.append(&mut statement.type_check(symbol_table));
        }

        diagnostics
    }

    pub fn to_ir(&self, symbol_table: &SymbolTable) -> IrFunction {
        let mut builder = IrBuilder::new();
        let entry_block_id = builder.new_block();
        for (i, statement) in self.statements.iter().enumerate() {
            let is_last = i == self.statements.len() - 1;
            statement.to_ir(&mut builder, symbol_table, is_last);
        }
        builder.finish_block();

        let entry_block = builder.get_block(entry_block_id);
        IrFunction::new(Rc::from(self.declared_name()), entry_block.clone())
    }

    pub fn assemble(
        &self,
        context: &mut AssembleContext,
        symbol_table: &SymbolTable,
        constants_table: &mut ConstantsTable,
    ) {
        context.new_function(&self.declared_name, &self.declared_name_source_range);
        context.new_block(&format!("{}_enter", self.declared_name));
        let function_symbol = self
            .function_symbol
            .as_ref()
            .expect("function_symbol not initialized; did you type check yet?");
        for (i, statement) in self.statements.iter().enumerate() {
            let is_last = i == self.statements.len() - 1;
            statement.assemble(
                context,
                symbol_table,
                constants_table,
                function_symbol,
                is_last,
            );
        }
        // return
        context.instruction(AsmInstruction::Return(Return::new(self.parameters.len())));

        context.complete_function();
    }
}