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

use crate::ast::expression::Expression;
use crate::ast::ir_builder::IrBuilder;
use crate::diagnostic::Diagnostic;
use crate::ir::ir_assign::IrAssign;
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::variable_symbol::VariableSymbol;
use crate::symbol_table::SymbolTable;
use crate::symbol_table::util::try_insert_symbol_into;
use crate::type_info::TypeInfo;
use crate::types_table::TypesTable;
use std::cell::RefCell;
use std::rc::Rc;

pub struct LetStatement {
    declared_name: Rc<str>,
    declared_name_source_range: SourceRange,
    is_mut: bool,
    initializer: Box<Expression>,
    scope_id: Option<usize>,
}

impl LetStatement {
    pub fn new(
        declared_name: &str,
        declared_name_source_range: SourceRange,
        is_mut: bool,
        initializer: Expression,
    ) -> Self {
        Self {
            declared_name: declared_name.into(),
            declared_name_source_range,
            is_mut,
            initializer: initializer.into(),
            scope_id: None,
        }
    }

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

    pub fn initializer(&self) -> &Expression {
        &self.initializer
    }

    pub fn initializer_mut(&mut self) -> &mut Expression {
        &mut self.initializer
    }

    pub fn init_scopes(&mut self, symbol_table: &mut SymbolTable, container_scope: usize) {
        self.scope_id = Some(container_scope);
        self.initializer.init_scopes(symbol_table, container_scope);
    }

    pub fn scope_id(&self) -> usize {
        self.scope_id.unwrap()
    }

    fn make_and_insert_variable_symbol(
        &self,
        symbol_table: &mut SymbolTable,
    ) -> Option<Diagnostic> {
        let variable_symbol = Rc::new(VariableSymbol::new(
            &self.declared_name,
            &self.declared_name_source_range,
            self.is_mut,
            self.scope_id.unwrap(),
        ));
        try_insert_symbol_into(Symbol::Variable(variable_symbol), symbol_table).err()
    }

    pub fn analyze_constructor_local_names(
        &self,
        symbol_table: &mut SymbolTable,
        class_symbol: &ClassSymbol,
    ) -> Vec<Diagnostic> {
        let mut diagnostics = Vec::new();
        diagnostics.append(
            &mut self
                .initializer
                .check_constructor_local_names(symbol_table, class_symbol),
        );
        if let Some(diagnostic) = self.make_and_insert_variable_symbol(symbol_table) {
            diagnostics.push(diagnostic);
        }

        diagnostics
    }

    pub fn analyze_method_local_names(
        &self,
        symbol_table: &mut SymbolTable,
        class_symbol: &ClassSymbol,
    ) -> Vec<Diagnostic> {
        let mut diagnostics = Vec::new();
        diagnostics.append(
            &mut self
                .initializer
                .check_method_local_names(symbol_table, class_symbol),
        );
        if let Some(diagnostic) = self.make_and_insert_variable_symbol(symbol_table) {
            diagnostics.push(diagnostic);
        }
        diagnostics
    }

    pub fn analyze_static_fn_local_names(&self, symbol_table: &mut SymbolTable) -> Vec<Diagnostic> {
        let mut diagnostics = Vec::new();
        diagnostics.append(&mut self.initializer.check_static_fn_local_names(symbol_table));
        if let Some(diagnostic) = self.make_and_insert_variable_symbol(symbol_table) {
            diagnostics.push(diagnostic);
        }
        diagnostics
    }

    pub fn gather_local_type(&self, symbol_table: &SymbolTable, types_table: &mut TypesTable) {
        let initializer_type_info = self
            .initializer
            .type_info(symbol_table, types_table)
            .clone();
        let variable_symbol = symbol_table
            .get_variable_symbol_owned(self.scope_id.unwrap(), &self.declared_name)
            .unwrap();
        types_table
            .variable_types_mut()
            .insert(variable_symbol, initializer_type_info);
    }

    pub fn type_check(
        &mut self,
        symbol_table: &SymbolTable,
        types_table: &mut TypesTable,
    ) -> Result<(), Vec<Diagnostic>> {
        self.initializer.type_check(symbol_table, types_table)?;
        // TODO: this is wrong. We need to check assignability
        let initializer_type_info = self
            .initializer
            .type_info(symbol_table, types_table)
            .clone();
        let variable_symbol = symbol_table
            .get_variable_symbol_owned(self.scope_id.unwrap(), &self.declared_name)
            .unwrap();
        types_table
            .variable_types_mut()
            .insert(variable_symbol, initializer_type_info);
        Ok(())
    }

    fn make_vr_variable(&self, builder: &mut IrBuilder, destination_type: &TypeInfo) -> IrVariable {
        IrVariable::new_vr(
            self.declared_name().into(),
            builder.current_block().id(),
            destination_type,
        )
    }

    fn make_stack_variable(
        &self,
        builder: &mut IrBuilder,
        destination_type: &TypeInfo,
        offset: isize,
    ) -> IrVariable {
        IrVariable::new_stack_with_offset(
            self.declared_name().into(),
            builder.current_block().id(),
            destination_type,
            offset,
        )
    }

    pub fn get_destination_symbol(&self, symbol_table: &SymbolTable) -> Rc<VariableSymbol> {
        symbol_table
            .get_variable_symbol_owned(self.scope_id.unwrap(), &self.declared_name)
            .unwrap()
    }

    pub fn to_ir(
        &self,
        builder: &mut IrBuilder,
        symbol_table: &SymbolTable,
        types_table: &TypesTable,
    ) {
        let init_operation = self
            .initializer
            .to_ir_operation(builder, symbol_table, types_table);

        let destination_symbol = self.get_destination_symbol(symbol_table);

        let destination_type = types_table
            .variable_types()
            .get(&destination_symbol)
            .unwrap();

        let destination_vr_variable = self.make_vr_variable(builder, destination_type);

        let as_rc = Rc::new(RefCell::new(destination_vr_variable));
        let ir_assign = IrAssign::new(as_rc.clone(), init_operation);

        builder
            .local_variables_mut()
            .insert(destination_symbol, as_rc.clone());

        builder
            .current_block_mut()
            .add_statement(IrStatement::Assign(ir_assign));
    }

    pub fn to_repl_ir(
        &self,
        builder: &mut IrBuilder,
        symbol_table: &SymbolTable,
        types_table: &TypesTable,
        destination_stack_offset: isize,
    ) -> Rc<RefCell<IrVariable>> {
        let init_operation = self
            .initializer
            .to_ir_operation(builder, symbol_table, types_table);
        let destination_symbol = self.get_destination_symbol(symbol_table);
        let destination_type = types_table
            .variable_types()
            .get(&destination_symbol)
            .unwrap();
        let destination_stack_variable =
            self.make_stack_variable(builder, destination_type, destination_stack_offset);
        let as_rc = Rc::new(RefCell::new(destination_stack_variable));
        let ir_assign = IrAssign::new(as_rc.clone(), init_operation);
        // do not need to save variable to builder as a new one is created for each repl function
        builder
            .current_block_mut()
            .add_statement(IrStatement::Assign(ir_assign));
        as_rc
    }
}