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#[derive(Clone, Debug, Eq, PartialEq)]
29pub enum CellExpression {
30 Deref(CellRef),
31 DoubleDeref(CellRef, i16),
33 Immediate(BigInt),
34 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 pub fn to_deref(&self) -> Option<CellRef> {
62 try_extract_matches!(self, CellExpression::Deref).cloned()
63 }
64
65 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 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 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 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}