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

use crate::asm::asm_instruction::{AsmInstruction, LoadConstant, Move, Operand, Pop};
use crate::ast::assemble_context::AssembleContext;
use crate::ast::expression::Expression;
use crate::constants_table::ConstantsTable;
use crate::diagnostic::Diagnostic;
use crate::source_range::SourceRange;
use crate::symbol::{ExpressibleSymbol, VariableSymbol};
use crate::symbol_table::{SymbolInsertError, SymbolTable};

pub struct LetStatement {
    declared_name: String,
    declared_name_source_range: SourceRange,
    initializer: Box<Expression>,
    scope_id: Option<usize>,
}

impl LetStatement {
    pub fn new(
        declared_name: &str,
        declared_name_source_range: SourceRange,
        initializer: Expression,
    ) -> Self {
        Self {
            declared_name: declared_name.to_string(),
            declared_name_source_range,
            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 gather_declared_names(&mut self, symbol_table: &mut SymbolTable) -> Vec<Diagnostic> {
        let mut diagnostics = vec![];
        self.initializer_mut().gather_declared_names(symbol_table);
        let insert_result =
            symbol_table.insert_variable_symbol(VariableSymbol::new(self.declared_name()));
        if let Err(symbol_insert_error) = insert_result {
            match symbol_insert_error {
                SymbolInsertError::AlreadyDeclared(already_declared) => {
                    diagnostics.push(Diagnostic::new(
                        &format!(
                            "Symbol {} already declared in current scope",
                            already_declared.name()
                        ),
                        self.declared_name_source_range.start(),
                        self.declared_name_source_range.end(),
                    ))
                }
            }
        }
        self.scope_id = Some(symbol_table.current_scope_id());
        diagnostics
    }

    pub fn check_name_usages(&mut self, symbol_table: &SymbolTable) -> Vec<Diagnostic> {
        self.initializer.check_name_usages(symbol_table)
    }

    pub fn type_check(&mut self, symbol_table: &SymbolTable) -> Vec<Diagnostic> {
        let initializer_diagnostics = self.initializer.type_check(symbol_table);
        if !initializer_diagnostics.is_empty() {
            return initializer_diagnostics;
        }
        let initializer_type_info = self.initializer.type_info();
        let variable_symbol =
            symbol_table.get_variable_symbol(self.scope_id.unwrap(), &self.declared_name);
        variable_symbol
            .borrow_mut()
            .set_type_info(initializer_type_info);
        vec![]
    }

    pub fn assemble(
        &self,
        context: &mut AssembleContext,
        symbol_table: &SymbolTable,
        constants_table: &mut ConstantsTable,
    ) {
        let destination_register = context.new_local_register();

        // save register to symbol
        let variable_symbol =
            symbol_table.get_variable_symbol(self.scope_id.unwrap(), self.declared_name());
        variable_symbol
            .borrow_mut()
            .set_register(destination_register);

        match self.initializer() {
            Expression::Call(call) => {
                call.assemble(context, symbol_table, constants_table);
                context.instruction(AsmInstruction::Pop(Pop::new(destination_register)));
            }
            Expression::IntegerLiteral(integer_literal) => {
                context.instruction(AsmInstruction::Move(Move::new(
                    Operand::IntegerLiteral(integer_literal.value()),
                    destination_register,
                )));
            }
            Expression::String(string_literal) => {
                let name = constants_table.insert_string(string_literal.content());
                context.instruction(AsmInstruction::LoadConstant(LoadConstant::new(
                    &name,
                    destination_register,
                )));
            }
            Expression::Identifier(identifier) => {
                let expressible_symbol = identifier.expressible_symbol();
                match expressible_symbol {
                    ExpressibleSymbol::Function(_) => {
                        panic!("Moving functions to registers not yet supported");
                    }
                    ExpressibleSymbol::Parameter(parameter_symbol) => {
                        context.instruction(AsmInstruction::Move(Move::new(
                            Operand::StackFrameOffset(
                                parameter_symbol.borrow().stack_frame_offset(),
                            ),
                            destination_register,
                        )));
                    }
                    ExpressibleSymbol::Variable(variable_symbol) => {
                        context.instruction(AsmInstruction::Move(Move::new(
                            Operand::Register(variable_symbol.borrow().register()),
                            destination_register,
                        )));
                    }
                }
            }
            Expression::Additive(additive) => {
                let result_register = additive.assemble(context, symbol_table, constants_table);
                context.instruction(AsmInstruction::Move(Move::new(
                    Operand::Register(result_register),
                    destination_register,
                )));
            }
        }
    }
}