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

use crate::ast::ir_builder::IrBuilder;
use crate::ast::parameter::Parameter;
use crate::ast::statement::Statement;
use crate::ast::type_use::TypeUse;
use crate::diagnostic::Diagnostic;
use crate::ir::ir_function::IrFunction;
use crate::ir::ir_parameter::IrParameter;
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![];
        for parameter in &mut self.parameters {
            match parameter.gather_declared_names(symbol_table) {
                Ok(parameter_symbol) => {
                    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();

        // parameters
        for (i, parameter) in self.parameters.iter().enumerate() {
            let parameter_symbol = parameter.parameter_symbol();
            let stack_offset = (self.parameters.len() as isize).neg() + (i as isize);
            let ir_parameter = IrParameter::new(
                parameter_symbol.borrow().name(),
                parameter_symbol.borrow().type_info().clone(),
                stack_offset,
            );
            let as_rc = Rc::new(ir_parameter);
            builder.parameters_mut().push(as_rc.clone());
            parameter_symbol.borrow_mut().set_ir_parameter(as_rc);
        }

        let entry_block_id = builder.new_block();

        let return_type_info = self
            .function_symbol
            .as_ref()
            .unwrap()
            .borrow()
            .return_type();

        let should_return_value = !matches!(return_type_info, TypeInfo::Void);

        for (i, statement) in self.statements.iter().enumerate() {
            let is_last = i == self.statements.len() - 1;
            statement.to_ir(&mut builder, symbol_table, should_return_value && is_last);
        }
        builder.finish_block();

        let entry_block = builder.get_block(entry_block_id).clone();
        IrFunction::new(
            self.function_symbol.as_ref().unwrap().clone(),
            builder.parameters(),
            return_type_info,
            entry_block,
        )
    }
}