cairo_lang_casm/
cell_expression.rs

1use cairo_lang_utils::try_extract_matches;
2use num_bigint::BigInt;
3use num_traits::cast::ToPrimitive;
4
5use crate::ap_change::ApplyApChange;
6use crate::operand::{BinOpOperand, CellRef, DerefOrImmediate, Operation, ResOperand};
7
8#[derive(Clone, Debug, Eq, PartialEq)]
9pub enum CellOperator {
10    Add,
11    Sub,
12    Mul,
13    Div,
14}
15
16impl core::fmt::Display for CellOperator {
17    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
18        match self {
19            CellOperator::Add => write!(f, "+"),
20            CellOperator::Sub => write!(f, "-"),
21            CellOperator::Mul => write!(f, "*"),
22            CellOperator::Div => write!(f, "/"),
23        }
24    }
25}
26
27/// The expression representing a cell in the casm memory.
28#[derive(Clone, Debug, Eq, PartialEq)]
29pub enum CellExpression {
30    Deref(CellRef),
31    /// Represents an expression of the form `[[cell_ref] + offset]`.
32    DoubleDeref(CellRef, i16),
33    Immediate(BigInt),
34    /// Represents an expression of the form `[cell_ref] + [cell_ref]` or `[cell_ref] + imm`.
35    ///
36    /// If `op` is [CellOperator::Div], `b` must not be zero.
37    BinOp {
38        op: CellOperator,
39        a: CellRef,
40        b: DerefOrImmediate,
41    },
42}
43impl CellExpression {
44    pub fn from_res_operand(operand: ResOperand) -> Self {
45        match operand {
46            ResOperand::Deref(cell) => Self::Deref(cell),
47            ResOperand::DoubleDeref(cell, offset) => Self::DoubleDeref(cell, offset),
48            ResOperand::Immediate(imm) => Self::Immediate(imm.value),
49            ResOperand::BinOp(BinOpOperand { op, a, b }) => Self::BinOp {
50                op: match op {
51                    Operation::Add => CellOperator::Add,
52                    Operation::Mul => CellOperator::Mul,
53                },
54                a,
55                b,
56            },
57        }
58    }
59
60    /// Extract the cell reference from the cell expression.
61    pub fn to_deref(&self) -> Option<CellRef> {
62        try_extract_matches!(self, CellExpression::Deref).cloned()
63    }
64
65    /// Extract a deref or immediate from the cell expression.
66    pub fn to_deref_or_immediate(&self) -> Option<DerefOrImmediate> {
67        match self {
68            CellExpression::Deref(cell) => Some(DerefOrImmediate::Deref(*cell)),
69            CellExpression::Immediate(imm) => Some(imm.clone().into()),
70            _ => None,
71        }
72    }
73
74    /// Given `[ref] + offset` returns `([ref], offset)`.
75    pub fn to_deref_with_offset(&self) -> Option<(CellRef, i32)> {
76        match self {
77            CellExpression::Deref(cell) => Some((*cell, 0)),
78            CellExpression::BinOp {
79                op: CellOperator::Add,
80                a: cell,
81                b: DerefOrImmediate::Immediate(offset),
82            } => Some((*cell, offset.value.to_i32()?)),
83            _ => None,
84        }
85    }
86
87    /// Returns the reference as a buffer with at least `required_slack` next cells that can be
88    /// written as an instruction offset.
89    pub fn to_buffer(&self, required_slack: i16) -> Option<CellExpression> {
90        let (base, offset) = self.to_deref_with_offset()?;
91        offset.to_i16()?.checked_add(required_slack)?;
92        if offset == 0 {
93            Some(CellExpression::Deref(base))
94        } else {
95            Some(CellExpression::BinOp {
96                op: CellOperator::Add,
97                a: base,
98                b: DerefOrImmediate::Immediate(offset.into()),
99            })
100        }
101    }
102}
103
104impl ApplyApChange for CellExpression {
105    fn apply_known_ap_change(self, ap_change: usize) -> Option<Self> {
106        Some(match self {
107            CellExpression::Deref(operand) => {
108                CellExpression::Deref(operand.apply_known_ap_change(ap_change)?)
109            }
110            CellExpression::DoubleDeref(operand, offset) => {
111                CellExpression::DoubleDeref(operand.apply_known_ap_change(ap_change)?, offset)
112            }
113            CellExpression::BinOp { op, a, b } => CellExpression::BinOp {
114                op,
115                a: a.apply_known_ap_change(ap_change)?,
116                b: b.apply_known_ap_change(ap_change)?,
117            },
118            expr @ CellExpression::Immediate(_) => expr,
119        })
120    }
121
122    fn can_apply_unknown(&self) -> bool {
123        match self {
124            CellExpression::Deref(operand) | CellExpression::DoubleDeref(operand, _) => {
125                operand.can_apply_unknown()
126            }
127            CellExpression::Immediate(_) => true,
128            CellExpression::BinOp { a, b, .. } => a.can_apply_unknown() && b.can_apply_unknown(),
129        }
130    }
131}
132
133impl core::fmt::Display for CellExpression {
134    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
135        match self {
136            CellExpression::Deref(cell) => write!(f, "{cell}"),
137            CellExpression::DoubleDeref(cell, offset) => write!(f, "[{cell} + {offset}]"),
138            CellExpression::Immediate(imm) => write!(f, "{}", imm),
139            CellExpression::BinOp { op, a, b } => write!(f, "{a} {op} {b}"),
140        }
141    }
142}
cairo_lang_casm/cell_expression.rs

cairo_lang_casm/
cell_expression.rs
