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

use crate::ast::expression::Expression;
use crate::ast::ir_builder::IrBuilder;
use crate::diagnostic::Diagnostic;
use crate::error_codes::BINARY_INCOMPATIBLE_TYPES;
use crate::ir::ir_assign::IrAssign;
use crate::ir::ir_binary_operation::{IrBinaryOperation, IrBinaryOperator};
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_table::SymbolTable;
use crate::type_info::TypeInfo;
use crate::{diagnostics_result, handle_diagnostic, handle_diagnostics, maybe_return_diagnostics};
use std::cell::RefCell;
use std::rc::Rc;

pub enum BinaryOperation {
    Multiply,
    Divide,
    Modulo,
    Add,
    Subtract,
    LeftShift,
    RightShift,
}

pub struct BinaryExpression {
    lhs: Box<Expression>,
    rhs: Box<Expression>,
    op: BinaryOperation,
    source_range: SourceRange,
    type_info: Option<TypeInfo>,
}

impl BinaryExpression {
    pub fn new(
        lhs: Expression,
        rhs: Expression,
        op: BinaryOperation,
        source_range: SourceRange,
    ) -> Self {
        Self {
            lhs: lhs.into(),
            rhs: rhs.into(),
            op,
            source_range,
            type_info: None,
        }
    }

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

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

    pub fn op(&self) -> &BinaryOperation {
        &self.op
    }

    pub fn source_range(&self) -> &SourceRange {
        &self.source_range
    }

    pub fn type_info(&self) -> &TypeInfo {
        self.type_info.as_ref().unwrap()
    }

    pub fn gather_declared_names(
        &mut self,
        symbol_table: &mut SymbolTable,
    ) -> Result<(), Vec<Diagnostic>> {
        let diagnostics = [&mut self.lhs, &mut self.rhs]
            .iter_mut()
            .map(|expression| expression.gather_declared_names(symbol_table))
            .filter_map(|result| result.err())
            .flatten()
            .collect::<Vec<Diagnostic>>();
        if diagnostics.is_empty() {
            Ok(())
        } else {
            Err(diagnostics)
        }
    }

    pub fn check_name_usages(&mut self, symbol_table: &SymbolTable) -> Result<(), Vec<Diagnostic>> {
        let diagnostics: Vec<Diagnostic> = [&mut self.lhs, &mut self.rhs]
            .iter_mut()
            .map(|expression| expression.check_name_usages(symbol_table))
            .filter_map(Result::err)
            .flatten()
            .collect();
        if diagnostics.is_empty() {
            Ok(())
        } else {
            Err(diagnostics)
        }
    }

    fn check_op(
        &mut self,
        check: impl Fn(&TypeInfo, &TypeInfo) -> bool,
        op_result: impl Fn(&TypeInfo, &TypeInfo) -> TypeInfo,
        lazy_diagnostic_message: impl Fn(&TypeInfo, &TypeInfo) -> String,
    ) -> Result<(), Diagnostic> {
        let lhs_type_info = self.lhs.type_info();
        let rhs_type_info = self.rhs.type_info();

        if check(lhs_type_info, rhs_type_info) {
            self.type_info = Some(op_result(lhs_type_info, rhs_type_info));
            Ok(())
        } else {
            let diagnostic = Diagnostic::new(
                &lazy_diagnostic_message(lhs_type_info, rhs_type_info),
                self.source_range.start(),
                self.source_range.end(),
            )
            .with_primary_label_message("Incompatible types for addition.")
            .with_reporter(file!(), line!())
            .with_error_code(BINARY_INCOMPATIBLE_TYPES);
            Err(diagnostic)
        }
    }

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

        handle_diagnostics!(self.lhs.type_check(symbol_table), diagnostics);
        handle_diagnostics!(self.rhs.type_check(symbol_table), diagnostics);

        maybe_return_diagnostics!(diagnostics);

        match &self.op {
            BinaryOperation::Multiply => {
                todo!()
            }
            BinaryOperation::Divide => {
                todo!()
            }
            BinaryOperation::Modulo => {
                todo!()
            }
            BinaryOperation::Add => {
                handle_diagnostic!(
                    self.check_op(
                        |lhs, rhs| lhs.can_add(rhs),
                        |lhs, rhs| lhs.add_result(&rhs),
                        |lhs, rhs| format!("Incompatible types: cannot add {} to {}.", rhs, lhs)
                    ),
                    diagnostics
                );
            }
            BinaryOperation::Subtract => {
                handle_diagnostic!(
                    self.check_op(
                        |lhs, rhs| lhs.can_subtract(rhs),
                        |lhs, rhs| lhs.subtract_result(rhs),
                        |lhs, rhs| format!(
                            "Incompatible types: cannot subtract {} from {}.",
                            rhs, lhs
                        )
                    ),
                    diagnostics
                )
            }
            BinaryOperation::LeftShift => todo!(),
            BinaryOperation::RightShift => todo!(),
        }

        diagnostics_result!(diagnostics)
    }

    pub fn to_ir_operation(
        &self,
        builder: &mut IrBuilder,
        symbol_table: &SymbolTable,
    ) -> IrOperation {
        let lhs = self
            .lhs
            .to_ir_expression(builder, symbol_table)
            .expect("Attempt to use a non-value expression in binary expression.");
        let rhs = self
            .rhs
            .to_ir_expression(builder, symbol_table)
            .expect("Attempt to use a non-value expression in binary expression.");
        let ir_binary_operation = match self.op {
            BinaryOperation::Multiply => {
                IrBinaryOperation::new(lhs, rhs, IrBinaryOperator::Multiply)
            }
            BinaryOperation::Divide => IrBinaryOperation::new(lhs, rhs, IrBinaryOperator::Divide),
            BinaryOperation::Modulo => todo!(),
            BinaryOperation::Add => IrBinaryOperation::new(lhs, rhs, IrBinaryOperator::Add),
            BinaryOperation::Subtract => {
                IrBinaryOperation::new(lhs, rhs, IrBinaryOperator::Subtract)
            }
            BinaryOperation::LeftShift => todo!(),
            BinaryOperation::RightShift => todo!(),
        };
        IrOperation::Binary(ir_binary_operation)
    }

    pub fn to_ir_expression(
        &self,
        builder: &mut IrBuilder,
        symbol_table: &SymbolTable,
    ) -> IrExpression {
        let ir_operation = self.to_ir_operation(builder, symbol_table);
        let t_var = IrVariable::new_vr(
            builder.new_t_var().into(),
            builder.current_block().id(),
            self.type_info(),
        );
        let as_rc = Rc::new(RefCell::new(t_var));
        let ir_assign = IrAssign::new(as_rc.clone(), ir_operation);
        builder
            .current_block_mut()
            .add_statement(IrStatement::Assign(ir_assign));
        IrExpression::Variable(as_rc)
    }
}