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

use crate::ast::field::Field;
use crate::ast::ir_builder::IrBuilder;
use crate::ast::parameter::Parameter;
use crate::ast::statement::Statement;
use crate::diagnostic::Diagnostic;
use crate::ir::ir_allocate::IrAllocate;
use crate::ir::ir_assign::IrAssign;
use crate::ir::ir_expression::IrExpression;
use crate::ir::ir_function::IrFunction;
use crate::ir::ir_operation::IrOperation;
use crate::ir::ir_parameter::IrParameter;
use crate::ir::ir_set_field::IrSetField;
use crate::ir::ir_statement::IrStatement;
use crate::ir::ir_variable::IrVariable;
use crate::source_range::SourceRange;
use crate::symbol::Symbol;
use crate::symbol::class_symbol::ClassSymbol;
use crate::symbol::constructor_symbol::ConstructorSymbol;
use crate::symbol::parameter_symbol::ParameterSymbol;
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 Constructor {
    is_public: bool,
    ctor_keyword_source_range: SourceRange,
    parameters: Vec<Parameter>,
    statements: Vec<Statement>,
}

impl Constructor {
    pub fn new(
        is_public: bool,
        ctor_keyword_source_range: SourceRange,
        parameters: Vec<Parameter>,
        statements: Vec<Statement>,
    ) -> Self {
        Self {
            is_public,
            ctor_keyword_source_range,
            parameters,
            statements,
        }
    }

    pub fn statements(&self) -> &[Statement] {
        &self.statements
    }

    pub fn gather_declared_names(
        &mut self,
        symbol_table: &mut SymbolTable,
    ) -> Result<Rc<RefCell<ConstructorSymbol>>, Vec<Diagnostic>> {
        // insert constructor symbol
        let to_insert = ConstructorSymbol::new(self.ctor_keyword_source_range.clone(), false);
        let constructor_symbol =
            symbol_table
                .insert_constructor_symbol(to_insert)
                .map_err(|symbol_insert_error| match symbol_insert_error {
                    SymbolInsertError::AlreadyDeclared(_) => {
                        vec![
                            Diagnostic::new(
                                "Cannot declare more than one constructor.",
                                self.ctor_keyword_source_range.start(),
                                self.ctor_keyword_source_range.end(),
                            )
                            .with_reporter(file!(), line!()),
                        ]
                    }
                })?;

        symbol_table.push_function_scope("ctor_scope");

        let mut parameter_symbols: Vec<Rc<RefCell<ParameterSymbol>>> = vec![];
        let mut parameters_diagnostics = vec![];

        for parameter in &mut self.parameters {
            match parameter.gather_declared_names(symbol_table) {
                Ok(parameter_symbol) => {
                    parameter_symbols.push(parameter_symbol);
                }
                Err(mut ds) => {
                    parameters_diagnostics.append(&mut ds);
                }
            }
        }

        if !parameters_diagnostics.is_empty() {
            symbol_table.pop_scope();
            return Err(parameters_diagnostics);
        } else {
            constructor_symbol
                .borrow_mut()
                .set_parameters(parameter_symbols);
        }

        symbol_table.push_block_scope("ctor_main_block");

        let statements_diagnostics = self
            .statements
            .iter_mut()
            .map(|stmt| stmt.gather_declared_names(symbol_table))
            .filter_map(Result::err)
            .flatten()
            .collect::<Vec<_>>();

        symbol_table.pop_scope(); // block
        symbol_table.pop_scope(); // function

        if statements_diagnostics.is_empty() {
            Ok(constructor_symbol)
        } else {
            Err(statements_diagnostics)
        }
    }

    pub fn check_name_usages(&mut self, symbol_table: &SymbolTable) -> Result<(), Vec<Diagnostic>> {
        let parameters_diagnostics: Vec<Diagnostic> = self
            .parameters
            .iter_mut()
            .map(|param| param.check_name_usages(symbol_table))
            .filter_map(Result::err)
            .flatten()
            .collect();

        if !parameters_diagnostics.is_empty() {
            return Err(parameters_diagnostics);
        }

        let statements_diagnostics: Vec<Diagnostic> = self
            .statements
            .iter_mut()
            .map(|statement| statement.check_name_usages(symbol_table))
            .filter_map(Result::err)
            .flatten()
            .collect();

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

    pub fn type_check(&mut self, symbol_table: &SymbolTable) -> Result<(), Vec<Diagnostic>> {
        let parameters_diagnostics: Vec<Diagnostic> = self
            .parameters
            .iter_mut()
            .map(|param| param.type_check(symbol_table))
            .filter_map(Result::err)
            .flatten()
            .collect();

        if !parameters_diagnostics.is_empty() {
            return Err(parameters_diagnostics);
        }

        let statements_diagnostics: Vec<Diagnostic> = self
            .statements
            .iter_mut()
            .map(|statement| statement.type_check(symbol_table, None))
            .filter_map(Result::err)
            .flatten()
            .collect();

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

    pub fn to_ir(
        &self,
        class_symbol: &Rc<RefCell<ClassSymbol>>,
        fields: &[Field],
        symbol_table: &SymbolTable,
    ) -> IrFunction {
        let mut ir_builder = IrBuilder::new();

        let parameters_count = self.parameters.len();
        let ir_parameters = self
            .parameters
            .iter()
            .enumerate()
            .map(|(i, parameter)| {
                let parameter_symbol = parameter.parameter_symbol().borrow();
                let offset = (parameters_count as isize).neg() + i as isize;
                Rc::new(IrParameter::new(
                    parameter_symbol.declared_name(),
                    parameter_symbol.type_info().clone(),
                    offset,
                ))
            })
            .collect::<Vec<_>>();

        let entry_block_id = ir_builder.new_block();

        // first, allocate the object into a t var
        let alloc_assign_destination = IrVariable::new_vr(
            ir_builder.new_t_var().into(),
            ir_builder.current_block().id(),
            &TypeInfo::ClassInstance(class_symbol.clone()),
        );
        let self_variable = Rc::new(RefCell::new(alloc_assign_destination));
        let alloc_assign = IrAssign::new(
            self_variable.clone(),
            IrOperation::Allocate(IrAllocate::new(class_symbol.borrow().declared_name_owned())),
        );
        ir_builder
            .current_block_mut()
            .add_statement(IrStatement::Assign(alloc_assign));

        // next, initialize fields that have an initializer in their declaration
        for field in fields {
            if let Some(initializer) = field.initializer() {
                let ir_expression = initializer.to_ir(&mut ir_builder, symbol_table).unwrap();
                let ir_set_field = IrSetField::new(
                    &self_variable.clone(),
                    field.field_symbol().borrow().field_index(),
                    ir_expression,
                );
                ir_builder
                    .current_block_mut()
                    .add_statement(IrStatement::SetField(ir_set_field));
            }
        }

        ir_builder.finish_block();
        let entry_block = ir_builder.get_block(entry_block_id);

        IrFunction::new(
            format!("{}::ctor", class_symbol.borrow().declared_name()).into(), // fake function symbol
            &ir_parameters,                                                    // make params
            &TypeInfo::ClassInstance(class_symbol.clone()),
            entry_block.clone(),
        )
    }
}