cairo-lang-casm 2.17.0

use cairo_lang_utils::try_extract_matches;
use num_bigint::BigInt;
use num_traits::cast::ToPrimitive;

use crate::ap_change::ApplyApChange;
use crate::operand::{BinOpOperand, CellRef, DerefOrImmediate, Operation, ResOperand};

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum CellOperator {
    Add,
    Sub,
    Mul,
    Div,
}

impl core::fmt::Display for CellOperator {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            CellOperator::Add => write!(f, "+"),
            CellOperator::Sub => write!(f, "-"),
            CellOperator::Mul => write!(f, "*"),
            CellOperator::Div => write!(f, "/"),
        }
    }
}

/// The expression representing a cell in the CASM memory.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum CellExpression {
    Deref(CellRef),
    /// Represents an expression of the form `[[cell_ref] + offset]`.
    DoubleDeref(CellRef, i16),
    Immediate(BigInt),
    /// Represents an expression of the form `[cell_ref] + [cell_ref]` or `[cell_ref] + imm`.
    ///
    /// If `op` is [CellOperator::Div], `b` must not be zero.
    BinOp {
        op: CellOperator,
        a: CellRef,
        b: DerefOrImmediate,
    },
}
impl CellExpression {
    pub fn from_res_operand(operand: ResOperand) -> Self {
        match operand {
            ResOperand::Deref(cell) => Self::Deref(cell),
            ResOperand::DoubleDeref(cell, offset) => Self::DoubleDeref(cell, offset),
            ResOperand::Immediate(imm) => Self::Immediate(imm.value),
            ResOperand::BinOp(BinOpOperand { op, a, b }) => Self::BinOp {
                op: match op {
                    Operation::Add => CellOperator::Add,
                    Operation::Mul => CellOperator::Mul,
                },
                a,
                b,
            },
        }
    }

    /// Extract the cell reference from the cell expression.
    pub fn to_deref(&self) -> Option<CellRef> {
        try_extract_matches!(self, CellExpression::Deref).cloned()
    }

    /// Extract a deref or immediate from the cell expression.
    pub fn to_deref_or_immediate(&self) -> Option<DerefOrImmediate> {
        match self {
            CellExpression::Deref(cell) => Some(DerefOrImmediate::Deref(*cell)),
            CellExpression::Immediate(imm) => Some(imm.clone().into()),
            _ => None,
        }
    }

    /// Given `[ref] + offset` returns `([ref], offset)`.
    pub fn to_deref_with_offset(&self) -> Option<(CellRef, i32)> {
        match self {
            CellExpression::Deref(cell) => Some((*cell, 0)),
            CellExpression::BinOp {
                op: CellOperator::Add,
                a: cell,
                b: DerefOrImmediate::Immediate(offset),
            } => Some((*cell, offset.value.to_i32()?)),
            _ => None,
        }
    }

    /// Returns the reference as a buffer with at least `required_slack` next cells that can be
    /// written as an instruction offset.
    pub fn to_buffer(&self, required_slack: i16) -> Option<CellExpression> {
        let (base, offset) = self.to_deref_with_offset()?;
        offset.to_i16()?.checked_add(required_slack)?;
        if offset == 0 {
            Some(CellExpression::Deref(base))
        } else {
            Some(CellExpression::BinOp {
                op: CellOperator::Add,
                a: base,
                b: DerefOrImmediate::Immediate(offset.into()),
            })
        }
    }

    /// Returns a cell expression which is the sum of `cell` and `value`.
    /// If the value is zero, returns the cell itself.
    pub fn add_with_const(cell: CellRef, value: i16) -> CellExpression {
        if value == 0 {
            CellExpression::Deref(cell)
        } else {
            CellExpression::BinOp {
                op: CellOperator::Add,
                a: cell,
                b: DerefOrImmediate::Immediate(value.into()),
            }
        }
    }
}

impl ApplyApChange for CellExpression {
    fn apply_known_ap_change(&mut self, ap_change: usize) -> bool {
        match self {
            CellExpression::Deref(operand) => operand.apply_known_ap_change(ap_change),
            CellExpression::DoubleDeref(operand, _) => operand.apply_known_ap_change(ap_change),
            CellExpression::BinOp { op: _, a, b } => {
                a.apply_known_ap_change(ap_change) && b.apply_known_ap_change(ap_change)
            }
            CellExpression::Immediate(_) => true,
        }
    }

    fn can_apply_unknown(&self) -> bool {
        match self {
            CellExpression::Deref(operand) | CellExpression::DoubleDeref(operand, _) => {
                operand.can_apply_unknown()
            }
            CellExpression::Immediate(_) => true,
            CellExpression::BinOp { a, b, .. } => a.can_apply_unknown() && b.can_apply_unknown(),
        }
    }
}

impl core::fmt::Display for CellExpression {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            CellExpression::Deref(cell) => write!(f, "{cell}"),
            CellExpression::DoubleDeref(cell, offset) => write!(f, "[{cell} + {offset}]"),
            CellExpression::Immediate(imm) => write!(f, "{imm}"),
            CellExpression::BinOp { op, a, b } => write!(f, "{a} {op} {b}"),
        }
    }
}