Skip to main content

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    /// Returns a cell expression which is the sum of `cell` and `value`.
104    /// If the value is zero, returns the cell itself.
105    pub fn add_with_const(cell: CellRef, value: i16) -> CellExpression {
106        if value == 0 {
107            CellExpression::Deref(cell)
108        } else {
109            CellExpression::BinOp {
110                op: CellOperator::Add,
111                a: cell,
112                b: DerefOrImmediate::Immediate(value.into()),
113            }
114        }
115    }
116}
117
118impl ApplyApChange for CellExpression {
119    fn apply_known_ap_change(&mut self, ap_change: usize) -> bool {
120        match self {
121            CellExpression::Deref(operand) => operand.apply_known_ap_change(ap_change),
122            CellExpression::DoubleDeref(operand, _) => operand.apply_known_ap_change(ap_change),
123            CellExpression::BinOp { op: _, a, b } => {
124                a.apply_known_ap_change(ap_change) && b.apply_known_ap_change(ap_change)
125            }
126            CellExpression::Immediate(_) => true,
127        }
128    }
129
130    fn can_apply_unknown(&self) -> bool {
131        match self {
132            CellExpression::Deref(operand) | CellExpression::DoubleDeref(operand, _) => {
133                operand.can_apply_unknown()
134            }
135            CellExpression::Immediate(_) => true,
136            CellExpression::BinOp { a, b, .. } => a.can_apply_unknown() && b.can_apply_unknown(),
137        }
138    }
139}
140
141impl core::fmt::Display for CellExpression {
142    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
143        match self {
144            CellExpression::Deref(cell) => write!(f, "{cell}"),
145            CellExpression::DoubleDeref(cell, offset) => write!(f, "[{cell} + {offset}]"),
146            CellExpression::Immediate(imm) => write!(f, "{imm}"),
147            CellExpression::BinOp { op, a, b } => write!(f, "{a} {op} {b}"),
148        }
149    }
150}