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

use crate::ast::assemble_context::AssembleContext;
use crate::ast::statement::Statement;
use crate::constants_table::ConstantsTable;
use crate::diagnostic::Diagnostic;
use crate::ir::Ir;
use crate::ir::ir_constant::IrConstant;
use crate::ir::ir_function::IrFunction;
use crate::ir::ir_statement::IrStatement;
use crate::source_range::SourceRange;
use crate::symbol::FunctionSymbol;
use crate::symbol_table::{SymbolInsertError, SymbolTable};

pub struct Function {
    declared_name: String,
    declared_name_source_range: SourceRange,
    statements: Vec<Statement>,
}

impl Function {
    pub fn new(
        declared_name: &str,
        declared_name_source_range: SourceRange,
        statements: Vec<Statement>,
    ) -> Self {
        Self {
            declared_name: declared_name.to_string(),
            declared_name_source_range,
            statements,
        }
    }

    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(),
            &vec![], // todo
            false,
        ));
        if let Err(symbol_insert_error) = insert_result {
            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(),
                    ));
                }
            }
        }
        symbol_table.push_scope(&format!("function_scope({})", self.declared_name()));
        for statement in &mut self.statements {
            diagnostics.append(&mut statement.gather_declared_names(symbol_table));
        }
        symbol_table.pop_scope();
        diagnostics
    }

    pub fn check_name_usages(&mut self, symbol_table: &SymbolTable) -> Vec<Diagnostic> {
        let mut diagnostics = vec![];
        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![];
        for statement in &mut self.statements {
            diagnostics.append(&mut statement.type_check(symbol_table));
        }
        diagnostics
    }

    pub fn lower_to_ir(&self) -> Vec<Ir> {
        let mut context = FunctionLoweringContext::new();
        for statement in &self.statements {
            statement.lower_to_ir(&mut context);
        }
        let mut irs = vec![];
        for constant in context.take_constants() {
            irs.push(Ir::Constant(constant));
        }
        let ir_function = IrFunction::new(&self.declared_name, context.take_statements());
        irs.push(Ir::Function(ir_function));
        irs
    }

    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));
        for statement in &self.statements {
            statement.assemble(context, symbol_table, constants_table);
        }
        context.complete_function();
    }
}

pub struct FunctionLoweringContext {
    temp_variable_counter: usize,
    constant_counter: usize,
    constants: Vec<IrConstant>,
    statements: Vec<IrStatement>,
}

impl FunctionLoweringContext {
    pub fn new() -> Self {
        Self {
            temp_variable_counter: 0,
            constant_counter: 0,
            constants: vec![],
            statements: vec![],
        }
    }

    pub fn next_temp_variable(&mut self) -> String {
        let temp_variable = format!("t_{}", self.temp_variable_counter);
        self.temp_variable_counter += 1;
        temp_variable
    }

    pub fn next_constant_name(&mut self) -> String {
        let constant_name = format!("%const_{}", self.constant_counter);
        self.constant_counter += 1;
        constant_name
    }

    pub fn add_constant(&mut self, constant: IrConstant) {
        self.constants.push(constant);
    }

    pub fn take_constants(&mut self) -> Vec<IrConstant> {
        std::mem::take(&mut self.constants)
    }

    pub fn add_statement(&mut self, statement: IrStatement) {
        self.statements.push(statement);
    }

    pub fn take_statements(&mut self) -> Vec<IrStatement> {
        std::mem::take(&mut self.statements)
    }
}