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

use crate::ast::expression::Expression;
use crate::ast::ir_builder::IrBuilder;
use crate::ast::ir_util::get_or_init_mut_field_pointer_variable;
use crate::diagnostic::{Diagnostic, SecondaryLabel};
use crate::error_codes::{ASSIGN_LHS_IMMUTABLE, ASSIGN_MISMATCHED_TYPES, ASSIGN_NO_L_VALUE};
use crate::ir::ir_assign::IrAssign;
use crate::ir::ir_set_field::IrSetField;
use crate::ir::ir_statement::IrStatement;
use crate::symbol::expressible_symbol::ExpressibleSymbol;
use crate::symbol_table::SymbolTable;

pub struct AssignStatement {
    destination: Box<Expression>,
    value: Box<Expression>,
}

impl AssignStatement {
    pub fn new(destination: Expression, value: Expression) -> Self {
        Self {
            destination: destination.into(),
            value: value.into(),
        }
    }

    pub fn gather_declared_names(
        &mut self,
        symbol_table: &mut SymbolTable,
    ) -> Result<(), Vec<Diagnostic>> {
        let mut diagnostics: Vec<Diagnostic> = vec![];

        match self.value.gather_declared_names(symbol_table) {
            Ok(_) => {}
            Err(mut value_diagnostics) => {
                diagnostics.append(&mut value_diagnostics);
            }
        }

        match self.destination.gather_declared_names(symbol_table) {
            Ok(_) => {}
            Err(mut destination_diagnostics) => {
                diagnostics.append(&mut destination_diagnostics);
            }
        }

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

    pub fn check_name_usages(&mut self, symbol_table: &SymbolTable) -> Result<(), Vec<Diagnostic>> {
        let mut diagnostics: Vec<Diagnostic> = vec![];

        match self.value.check_name_usages(symbol_table) {
            Ok(_) => {}
            Err(mut value_diagnostics) => {
                diagnostics.append(&mut value_diagnostics);
            }
        }

        match self.destination.check_name_usages(symbol_table) {
            Ok(_) => {}
            Err(mut destination_diagnostics) => {
                diagnostics.append(&mut destination_diagnostics);
            }
        }

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

    pub fn type_check(&mut self, symbol_table: &SymbolTable) -> Result<(), Vec<Diagnostic>> {
        let mut diagnostics: Vec<Diagnostic> = vec![];

        match self.value.type_check(symbol_table) {
            Ok(_) => {}
            Err(mut value_diagnostics) => {
                diagnostics.append(&mut value_diagnostics);
            }
        }

        match self.destination.type_check(symbol_table) {
            Ok(_) => {}
            Err(mut destination_diagnostics) => {
                diagnostics.append(&mut destination_diagnostics);
            }
        }

        // check destination is l value
        match &*self.destination {
            Expression::Identifier(identifier) => {
                let expressible_symbol = identifier.expressible_symbol();

                // check mutable
                if !expressible_symbol.is_mut() {
                    let secondary_label = SecondaryLabel::new(
                        expressible_symbol.source_range().start(),
                        expressible_symbol.source_range().end(),
                        Some("Destination (declared here) is immutable.".to_string()),
                    );
                    let diagnostic = Diagnostic::new(
                        "Destination is immutable and not re-assignable.",
                        self.destination.source_range().start(),
                        self.destination.source_range().end(),
                    )
                    .with_primary_label_message("Attempt to mutate immutable destination.")
                    .with_reporter(file!(), line!())
                    .with_error_code(ASSIGN_LHS_IMMUTABLE)
                    .with_secondary_labels(&[secondary_label]);
                    diagnostics.push(diagnostic);
                }

                // check assignable
                let lhs_type = expressible_symbol.type_info();
                let rhs_type = self.value.type_info();
                if !lhs_type.is_assignable_from(rhs_type) {
                    let secondary_label = SecondaryLabel::new(
                        expressible_symbol.source_range().start(),
                        expressible_symbol.source_range().end(),
                        Some(format!(
                            "Destination declared here is of type {}.",
                            lhs_type
                        )),
                    );
                    let diagnostic = Diagnostic::new(
                        &format!(
                            "Mismatched types: right-hand side {} is not assignable to left {}.",
                            rhs_type, lhs_type
                        ),
                        self.destination.source_range().start(),
                        self.value.source_range().end(),
                    )
                    .with_primary_label_message(&format!(
                        "Attempt to assign {} to {}.",
                        rhs_type, lhs_type
                    ))
                    .with_error_code(ASSIGN_MISMATCHED_TYPES)
                    .with_reporter(file!(), line!())
                    .with_secondary_labels(&[secondary_label]);
                    diagnostics.push(diagnostic);
                }
            }
            _ => {
                let diagnostic = Diagnostic::new(
                    "Left-hand side of assign must be an L value.",
                    self.destination.source_range().start(),
                    self.destination.source_range().end(),
                )
                .with_primary_label_message("Must be L value.")
                .with_reporter(file!(), line!())
                .with_error_code(ASSIGN_NO_L_VALUE);
                diagnostics.push(diagnostic);
            }
        }

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

    pub fn to_ir(&self, builder: &mut IrBuilder, symbol_table: &SymbolTable) {
        let destination_symbol = match &*self.destination {
            Expression::Identifier(identifier) => identifier.expressible_symbol(),
            _ => unreachable!("Destination must be a mutable L value"),
        };

        let ir_statement = match destination_symbol {
            ExpressibleSymbol::Field(field_symbol) => {
                let mut_field_pointer_variable =
                    get_or_init_mut_field_pointer_variable(builder, field_symbol).clone();
                let ir_set_field = IrSetField::new(
                    &mut_field_pointer_variable,
                    self.value
                        .to_ir_expression(builder, symbol_table)
                        .expect("Attempt to convert non-value to value"),
                );
                IrStatement::SetField(ir_set_field)
            }
            ExpressibleSymbol::Variable(variable_symbol) => {
                let vr_variable = variable_symbol.borrow().vr_variable().clone();
                let ir_assign = IrAssign::new(
                    vr_variable,
                    self.value.to_ir_operation(builder, symbol_table),
                );
                IrStatement::Assign(ir_assign)
            }
            _ => unreachable!("Destination must be a mutable L value"),
        };

        builder.current_block_mut().add_statement(ir_statement);
    }
}

#[cfg(test)]
mod tests {
    use crate::error_codes::{ASSIGN_LHS_IMMUTABLE, ASSIGN_MISMATCHED_TYPES, ASSIGN_NO_L_VALUE};
    use crate::parser::parse_compilation_unit;
    use crate::symbol_table::SymbolTable;

    #[test]
    fn finds_mismatched_types() {
        let mut compilation_unit = parse_compilation_unit(
            "
            fn main()
                let mut x = 4
                x = \"Hello\"
            end
            ",
        )
        .unwrap();
        let mut symbol_table = SymbolTable::new();
        compilation_unit
            .gather_declared_names(&mut symbol_table)
            .unwrap();
        compilation_unit.check_name_usages(&symbol_table).unwrap();
        match compilation_unit.type_check(&symbol_table) {
            Ok(_) => {
                panic!("Type check missed diagnostic.");
            }
            Err(diagnostics) => {
                assert_eq!(diagnostics.len(), 1);
                assert_eq!(
                    diagnostics[0].error_code().unwrap(),
                    ASSIGN_MISMATCHED_TYPES
                );
            }
        }
    }

    #[test]
    fn finds_no_l_value() {
        let mut compilation_unit = parse_compilation_unit(
            "
            fn main()
                42 = 42
            end
            ",
        )
        .unwrap();
        let mut symbol_table = SymbolTable::new();
        compilation_unit
            .gather_declared_names(&mut symbol_table)
            .unwrap();
        compilation_unit.check_name_usages(&symbol_table).unwrap();
        match compilation_unit.type_check(&symbol_table) {
            Ok(_) => {
                panic!("Type check missed diagnostic.");
            }
            Err(diagnostics) => {
                assert_eq!(diagnostics.len(), 1);
                assert_eq!(diagnostics[0].error_code().unwrap(), ASSIGN_NO_L_VALUE);
            }
        }
    }

    #[test]
    fn finds_immutable_destination() {
        let mut compilation_unit = parse_compilation_unit(
            "
            fn main()
                let x = 42
                x = 43
            end
            ",
        )
        .unwrap();
        let mut symbol_table = SymbolTable::new();
        compilation_unit
            .gather_declared_names(&mut symbol_table)
            .unwrap();
        compilation_unit.check_name_usages(&symbol_table).unwrap();
        match compilation_unit.type_check(&symbol_table) {
            Ok(_) => {
                panic!("Type check missed diagnostic.");
            }
            Err(diagnostics) => {
                assert_eq!(diagnostics.len(), 1);
                assert_eq!(diagnostics[0].error_code().unwrap(), ASSIGN_LHS_IMMUTABLE);
            }
        }
    }
}