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::ast::ir_builder::IrBuilder;
use crate::constants_table::ConstantsTable;
use crate::diagnostic::Diagnostic;
use crate::ir::ir_assign::IrAssign;
use crate::ir::ir_expression::IrExpression;
use crate::ir::ir_operation::IrOperation;
use crate::ir::ir_statement::IrStatement;
use crate::ir::ir_variable::IrVariable;
use crate::source_range::SourceRange;
use crate::symbol::{ExpressibleSymbol, VariableSymbol};
use crate::symbol_table::{SymbolInsertError, SymbolTable};
use std::cell::RefCell;
use std::rc::Rc;

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 to_ir(&self, builder: &mut IrBuilder, symbol_table: &SymbolTable) {
        let init_operation = match self.initializer() {
            Expression::Call(call) => IrOperation::Call(call.to_ir(builder, symbol_table)),
            Expression::IntegerLiteral(integer_literal) => {
                IrOperation::Load(IrExpression::Int(integer_literal.value()))
            }
            Expression::String(string_literal) => {
                IrOperation::Load(IrExpression::String(string_literal.content().into()))
            }
            Expression::Identifier(identifier) => {
                IrOperation::Load(identifier.expressible_symbol().ir_expression())
            }
            Expression::Additive(additive_expression) => {
                IrOperation::Add(additive_expression.to_ir(builder, symbol_table))
            }
        };

        let destination_symbol =
            symbol_table.get_variable_symbol(self.scope_id.unwrap(), &self.declared_name);
        let destination_vr_variable = IrVariable::new_vr(
            self.declared_name().into(),
            builder.current_block().id(),
            self.initializer.type_info(),
        );
        let as_rc = Rc::new(RefCell::new(destination_vr_variable));
        let ir_assign = IrAssign::new(as_rc.clone(), init_operation);
        destination_symbol.borrow_mut().set_vr_variable(as_rc);

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

    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,
                )));
            }
        }
    }
}