cairo_lang_casm/
builder.rs

1use cairo_lang_utils::casts::IntoOrPanic;
2use cairo_lang_utils::extract_matches;
3use cairo_lang_utils::small_ordered_map::{Entry, SmallOrderedMap};
4use num_bigint::BigInt;
5use num_traits::{One, Zero};
6
7use crate::ap_change::ApplyApChange;
8use crate::cell_expression::{CellExpression, CellOperator};
9use crate::deref_or_immediate;
10use crate::hints::Hint;
11use crate::instructions::{
12    AddApInstruction, AssertEqInstruction, Blake2sCompressInstruction, CallInstruction,
13    Instruction, InstructionBody, JnzInstruction, JumpInstruction, RetInstruction,
14};
15use crate::operand::{BinOpOperand, CellRef, DerefOrImmediate, Operation, Register, ResOperand};
16
17#[cfg(test)]
18#[path = "builder_test.rs"]
19mod test;
20
21/// Variables for CASM builder, representing a `CellExpression`.
22#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
23pub struct Var(usize);
24
25/// The kind of an assert_eq action.
26///
27/// To be used in `assert_vars_eq`.
28pub enum AssertEqKind {
29    Felt252,
30    QM31,
31}
32
33/// The state of the variables at some line.
34#[derive(Clone, Debug, Default)]
35pub struct State {
36    /// The value per variable.
37    vars: SmallOrderedMap<Var, CellExpression>,
38    /// The number of allocated variables from the beginning of the run.
39    allocated: i16,
40    /// The AP change since the beginning of the run.
41    pub ap_change: usize,
42    /// The number of CASM steps since the beginning of the run.
43    pub steps: usize,
44}
45impl State {
46    /// Returns the value, in relation to the initial ap value.
47    fn get_value(&self, var: Var) -> CellExpression {
48        self.vars[&var].clone()
49    }
50
51    /// Returns the value, in relation to the current ap value.
52    pub fn get_adjusted(&self, var: Var) -> CellExpression {
53        self.get_value(var).unchecked_apply_known_ap_change(self.ap_change)
54    }
55
56    /// Returns the value, assuming it is a direct cell reference.
57    pub fn get_adjusted_as_cell_ref(&self, var: Var) -> CellRef {
58        extract_matches!(self.get_adjusted(var), CellExpression::Deref)
59    }
60
61    /// Validates that the state is valid, as it had enough ap change.
62    fn validate_finality(&self) {
63        assert!(
64            self.ap_change >= self.allocated.into_or_panic(),
65            "Not enough instructions to update ap. Add an `ap += *` instruction. ap_change: {}, \
66             allocated: {}",
67            self.ap_change,
68            self.allocated,
69        );
70    }
71
72    /// Intersect the states of branches leading to the same label, validating that the states can
73    /// intersect.
74    fn intersect(&mut self, other: &Self, read_only: bool) {
75        assert_eq!(self.ap_change, other.ap_change, "Merged branches not aligned on AP change.");
76        assert_eq!(
77            self.allocated, other.allocated,
78            "Merged branches not aligned on number of allocations."
79        );
80        if read_only {
81            assert!(self.steps >= other.steps, "Finalized branch cannot be updated.");
82        } else {
83            self.steps = self.steps.max(other.steps);
84        }
85        // Allowing removal of variables as valid code won't be producible in that case anyway, and
86        // otherwise merging branches becomes very difficult.
87        self.vars.retain(|var, value| {
88            other
89                .vars
90                .get(var)
91                .map(|x| assert_eq!(x, value, "Var mismatch between branches."))
92                .is_some()
93        });
94    }
95}
96
97/// A statement added to the builder.
98enum Statement {
99    /// A final instruction, no need for further editing.
100    Final(Instruction),
101    /// A jump or call command, requires fixing the actual target label.
102    Jump(&'static str, Instruction),
103    /// A target label for jumps, additionally with an offset for defining the label in a position
104    /// relative to the current statement.
105    Label(&'static str, usize),
106}
107
108/// The builder result.
109pub struct CasmBuildResult<const BRANCH_COUNT: usize> {
110    /// The actual CASM code.
111    pub instructions: Vec<Instruction>,
112    /// The state and relocations per branch.
113    pub branches: [(State, Vec<usize>); BRANCH_COUNT],
114}
115
116/// Information about the state of a label.
117struct LabelInfo {
118    /// The state at a point of jumping into the label.
119    state: State,
120    /// Whether the label state is read-only.
121    read_only: bool,
122}
123
124/// Builder to more easily write CASM code.
125///
126/// Allows CASM building without specifically thinking about ap changes and the sizes of opcodes.
127/// Wrong usages of it would panic instead of returning a result, as this builder assumes we are in
128/// a post validation of parameters stage.
129pub struct CasmBuilder {
130    /// The information about the labels, per label.
131    label_info: SmallOrderedMap<&'static str, LabelInfo>,
132    /// The state at the last added statement.
133    main_state: State,
134    /// The added statements.
135    statements: Vec<Statement>,
136    /// The current set of added hints.
137    current_hints: Vec<Hint>,
138    /// The number of vars created. Used to not reuse var names.
139    var_count: usize,
140    /// Is the current state reachable.
141    /// Example for unreachable state is after a unconditional jump, before any label is stated.
142    reachable: bool,
143}
144impl CasmBuilder {
145    /// Finalizes the builder, with the requested labels as the returning branches.
146    /// "Fallthrough" is a special case for the fallthrough case.
147    pub fn build<const BRANCH_COUNT: usize>(
148        mut self,
149        branch_names: [&str; BRANCH_COUNT],
150    ) -> CasmBuildResult<BRANCH_COUNT> {
151        assert!(
152            self.current_hints.is_empty(),
153            "Build cannot be called with hints as the last addition."
154        );
155        let label_offsets = self.compute_label_offsets();
156        if self.reachable {
157            self.label_info
158                .insert("Fallthrough", LabelInfo { state: self.main_state, read_only: false });
159        }
160        let mut instructions = vec![];
161        let mut branch_relocations: [Vec<usize>; BRANCH_COUNT] = core::array::from_fn(|_| vec![]);
162        let mut offset = 0;
163        for statement in self.statements {
164            match statement {
165                Statement::Final(inst) => {
166                    offset += inst.body.op_size();
167                    instructions.push(inst);
168                }
169                Statement::Jump(label, mut inst) => {
170                    match label_offsets.get(&label) {
171                        Some(label_offset) => match &mut inst.body {
172                            InstructionBody::Jnz(JnzInstruction {
173                                jump_offset: DerefOrImmediate::Immediate(value),
174                                ..
175                            })
176                            | InstructionBody::Jump(JumpInstruction {
177                                target: DerefOrImmediate::Immediate(value),
178                                ..
179                            })
180                            | InstructionBody::Call(CallInstruction {
181                                target: DerefOrImmediate::Immediate(value),
182                                ..
183                            }) => {
184                                // Updating the value, instead of assigning into it, to avoid
185                                // allocating a BigInt since it is already 0.
186                                value.value += *label_offset as i128 - offset as i128;
187                            }
188
189                            _ => unreachable!("Only jump or call statements should be here."),
190                        },
191                        None => {
192                            let idx = branch_names.iter().position(|name| name == &label).unwrap();
193                            branch_relocations[idx].push(instructions.len())
194                        }
195                    }
196                    offset += inst.body.op_size();
197                    instructions.push(inst);
198                }
199                Statement::Label(name, _) => {
200                    self.label_info.remove(&name);
201                }
202            }
203        }
204        let branches = core::array::from_fn(|i| {
205            let label = branch_names[i];
206            let info = self
207                .label_info
208                .remove(&label)
209                .unwrap_or_else(|| panic!("Requested a non existing final label: {label:?}."));
210            info.state.validate_finality();
211            (info.state, core::mem::take(&mut branch_relocations[i]))
212        });
213        assert!(self.label_info.is_empty(), "Did not use all branches.");
214        CasmBuildResult { instructions, branches }
215    }
216
217    /// Returns the current ap change of the builder.
218    /// Useful for manual ap change handling.
219    pub fn curr_ap_change(&self) -> usize {
220        self.main_state.ap_change
221    }
222
223    /// Computes the code offsets of all the labels.
224    fn compute_label_offsets(&self) -> SmallOrderedMap<&'static str, usize> {
225        let mut label_offsets = SmallOrderedMap::default();
226        let mut offset = 0;
227        for statement in &self.statements {
228            match statement {
229                Statement::Final(inst) | Statement::Jump(_, inst) => {
230                    offset += inst.body.op_size();
231                }
232                Statement::Label(name, extra_offset) => {
233                    label_offsets.insert(*name, offset + extra_offset);
234                }
235            }
236        }
237        label_offsets
238    }
239
240    /// Adds a variable pointing to `value`.
241    pub fn add_var(&mut self, value: CellExpression) -> Var {
242        let var = Var(self.var_count);
243        self.var_count += 1;
244        self.main_state.vars.insert(var, value);
245        var
246    }
247
248    /// Allocates a new variable in memory, either local (FP-based) or temp (AP-based).
249    pub fn alloc_var(&mut self, local_var: bool) -> Var {
250        let var = self.add_var(CellExpression::Deref(CellRef {
251            offset: self.main_state.allocated,
252            register: if local_var { Register::FP } else { Register::AP },
253        }));
254        self.main_state.allocated += 1;
255        var
256    }
257
258    /// Returns an additional variable pointing to the same value.
259    pub fn duplicate_var(&mut self, var: Var) -> Var {
260        self.add_var(self.get_value(var, false))
261    }
262
263    /// Adds a hint, generated from `inputs` which are cell refs or immediates and `outputs` which
264    /// must be cell refs.
265    pub fn add_hint<
266        const INPUTS_COUNT: usize,
267        const OUTPUTS_COUNT: usize,
268        THint: Into<Hint>,
269        F: FnOnce([ResOperand; INPUTS_COUNT], [CellRef; OUTPUTS_COUNT]) -> THint,
270    >(
271        &mut self,
272        f: F,
273        inputs: [Var; INPUTS_COUNT],
274        outputs: [Var; OUTPUTS_COUNT],
275    ) {
276        self.current_hints.push(
277            f(
278                inputs.map(|v| match self.get_value(v, true) {
279                    CellExpression::Deref(cell) => ResOperand::Deref(cell),
280                    CellExpression::DoubleDeref(cell, offset) => {
281                        ResOperand::DoubleDeref(cell, offset)
282                    }
283                    CellExpression::Immediate(imm) => imm.into(),
284                    CellExpression::BinOp { op, a: other, b } => match op {
285                        CellOperator::Add => {
286                            ResOperand::BinOp(BinOpOperand { op: Operation::Add, a: other, b })
287                        }
288                        CellOperator::Mul => {
289                            ResOperand::BinOp(BinOpOperand { op: Operation::Mul, a: other, b })
290                        }
291                        CellOperator::Sub | CellOperator::Div => {
292                            panic!("hints to non ResOperand references are not supported.")
293                        }
294                    },
295                }),
296                outputs.map(|v| self.as_cell_ref(v, true)),
297            )
298            .into(),
299        );
300    }
301
302    /// Adds an assertion that `dst = res`.
303    /// `dst` must be a cell reference.
304    pub fn assert_vars_eq(&mut self, dst: Var, res: Var, kind: AssertEqKind) {
305        let a = self.as_cell_ref(dst, true);
306        let b = self.get_value(res, true);
307        let (a, b) = match b {
308            CellExpression::Deref(cell) => (a, ResOperand::Deref(cell)),
309            CellExpression::DoubleDeref(cell, offset) => (a, ResOperand::DoubleDeref(cell, offset)),
310            CellExpression::Immediate(imm) => (a, imm.into()),
311            CellExpression::BinOp { op, a: other, b } => match op {
312                CellOperator::Add => {
313                    (a, ResOperand::BinOp(BinOpOperand { op: Operation::Add, a: other, b }))
314                }
315                CellOperator::Mul => {
316                    (a, ResOperand::BinOp(BinOpOperand { op: Operation::Mul, a: other, b }))
317                }
318                CellOperator::Sub => {
319                    (other, ResOperand::BinOp(BinOpOperand { op: Operation::Add, a, b }))
320                }
321                CellOperator::Div => {
322                    (other, ResOperand::BinOp(BinOpOperand { op: Operation::Mul, a, b }))
323                }
324            },
325        };
326        let inner = AssertEqInstruction { a, b };
327        let instruction = self.next_instruction(
328            match kind {
329                AssertEqKind::Felt252 => InstructionBody::AssertEq(inner),
330                AssertEqKind::QM31 => InstructionBody::QM31AssertEq(inner),
331            },
332            true,
333        );
334        self.statements.push(Statement::Final(instruction));
335    }
336
337    /// Writes and increments a buffer.
338    /// Useful for RangeCheck and similar buffers.
339    /// `buffer` must be a cell reference, or a cell reference with a small added constant.
340    /// `value` must be a cell reference.
341    pub fn buffer_write_and_inc(&mut self, buffer: Var, value: Var) {
342        let (cell, offset) = self.buffer_get_and_inc(buffer);
343        let location = self.add_var(CellExpression::DoubleDeref(cell, offset));
344        self.assert_vars_eq(value, location, AssertEqKind::Felt252);
345    }
346
347    /// Writes `var` as a new tempvar and returns it as a variable, unless its value is already
348    /// `deref` (which includes trivial calculations as well) so it instead returns a variable
349    /// pointing to that location.
350    pub fn maybe_add_tempvar(&mut self, var: Var) -> Var {
351        self.add_var(CellExpression::Deref(match self.get_value(var, false) {
352            CellExpression::Deref(cell) => cell,
353            CellExpression::BinOp {
354                op: CellOperator::Add | CellOperator::Sub,
355                a,
356                b: DerefOrImmediate::Immediate(imm),
357            } if imm.value.is_zero() => a,
358            CellExpression::BinOp {
359                op: CellOperator::Mul | CellOperator::Div,
360                a,
361                b: DerefOrImmediate::Immediate(imm),
362            } if imm.value.is_one() => a,
363            _ => {
364                let temp = self.alloc_var(false);
365                self.assert_vars_eq(temp, var, AssertEqKind::Felt252);
366                return temp;
367            }
368        }))
369    }
370
371    /// Increments a buffer and allocates and returns variable pointing to its previous value.
372    pub fn get_ref_and_inc(&mut self, buffer: Var) -> Var {
373        let (cell, offset) = self.as_cell_ref_plus_const(buffer, 0, false);
374        self.main_state.vars.insert(
375            buffer,
376            CellExpression::BinOp {
377                op: CellOperator::Add,
378                a: cell,
379                b: deref_or_immediate!(BigInt::from(offset) + 1),
380            },
381        );
382        self.add_var(CellExpression::DoubleDeref(cell, offset))
383    }
384
385    /// Increments a buffer and returning the previous value it pointed to.
386    /// Useful for writing, reading and referencing values.
387    /// `buffer` must be a cell reference, or a cell reference with a small added constant.
388    fn buffer_get_and_inc(&mut self, buffer: Var) -> (CellRef, i16) {
389        let (base, offset) = match self.get_value(buffer, false) {
390            CellExpression::Deref(cell) => (cell, 0),
391            CellExpression::BinOp {
392                op: CellOperator::Add,
393                a,
394                b: DerefOrImmediate::Immediate(imm),
395            } => (a, imm.value.try_into().expect("Too many buffer writes.")),
396            _ => panic!("Not a valid buffer."),
397        };
398        self.main_state.vars.insert(
399            buffer,
400            CellExpression::BinOp {
401                op: CellOperator::Add,
402                a: base,
403                b: deref_or_immediate!(offset + 1),
404            },
405        );
406        (base, offset)
407    }
408
409    /// Increases AP by `size`.
410    pub fn add_ap(&mut self, size: usize) {
411        let instruction = self.next_instruction(
412            InstructionBody::AddAp(AddApInstruction { operand: BigInt::from(size).into() }),
413            false,
414        );
415        self.statements.push(Statement::Final(instruction));
416        self.main_state.ap_change += size;
417    }
418
419    /// Increases the AP change by `size`, without adding an instruction.
420    pub fn increase_ap_change(&mut self, amount: usize) {
421        self.main_state.ap_change += amount;
422        self.main_state.allocated += amount.into_or_panic::<i16>();
423    }
424
425    /// Returns a variable that is the `op` of `lhs` and `rhs`.
426    /// `lhs` must be a cell reference and `rhs` must be deref or immediate.
427    pub fn bin_op(&mut self, op: CellOperator, lhs: Var, rhs: Var) -> Var {
428        let (a, b) = match self.get_value(lhs, false) {
429            // Regular `bin_op` support.
430            CellExpression::Deref(cell) => (cell, self.as_deref_or_imm(rhs, false)),
431            // `add_with_const` + `imm` support.
432            CellExpression::BinOp {
433                op: CellOperator::Add,
434                a,
435                b: DerefOrImmediate::Immediate(imm),
436            } if op == CellOperator::Add => (
437                a,
438                DerefOrImmediate::Immediate(
439                    (imm.value
440                        + extract_matches!(self.get_value(rhs, false), CellExpression::Immediate))
441                    .into(),
442                ),
443            ),
444            _ => panic!(
445                "`bin_op` is supported only between a `deref` and a `deref_or_imm`, or a \
446                 `add_with_const` and `imm`."
447            ),
448        };
449        self.add_var(CellExpression::BinOp { op, a, b })
450    }
451
452    /// Returns a variable that is `[[var] + offset]`.
453    /// `var` must be a cell reference, or a cell ref plus a small constant.
454    pub fn double_deref(&mut self, var: Var, offset: i16) -> Var {
455        let (cell, full_offset) = self.as_cell_ref_plus_const(var, offset, false);
456        self.add_var(CellExpression::DoubleDeref(cell, full_offset))
457    }
458
459    /// Sets the label to have the set states, otherwise tests if the state matches the existing one
460    /// by merging.
461    fn set_or_test_label_state(&mut self, label: &'static str, state: State) {
462        match self.label_info.entry(label) {
463            Entry::Occupied(e) => {
464                let info = e.into_mut();
465                info.state.intersect(&state, info.read_only);
466            }
467            Entry::Vacant(e) => {
468                e.insert(LabelInfo { state, read_only: false });
469            }
470        }
471    }
472
473    /// Add a statement to jump to `label`.
474    pub fn jump(&mut self, label: &'static str) {
475        let instruction = self.next_instruction(
476            InstructionBody::Jump(JumpInstruction {
477                target: deref_or_immediate!(0),
478                relative: true,
479            }),
480            true,
481        );
482        self.statements.push(Statement::Jump(label, instruction));
483        let mut state = State::default();
484        core::mem::swap(&mut state, &mut self.main_state);
485        self.set_or_test_label_state(label, state);
486        self.reachable = false;
487    }
488
489    /// Add a statement to jump to `label` if `condition != 0`.
490    /// `condition` must be a cell reference.
491    pub fn jump_nz(&mut self, condition: Var, label: &'static str) {
492        let cell = self.as_cell_ref(condition, true);
493        let instruction = self.next_instruction(
494            InstructionBody::Jnz(JnzInstruction {
495                condition: cell,
496                jump_offset: deref_or_immediate!(0),
497            }),
498            true,
499        );
500        self.statements.push(Statement::Jump(label, instruction));
501        self.set_or_test_label_state(label, self.main_state.clone());
502    }
503
504    /// Add a statement performing Blake2s compression.
505    ///
506    /// `state` must be a cell reference to a pointer to `[u32; 8]` as the hash state.
507    /// `byte_count` must be a cell reference to the number of bytes in the message.
508    /// `message` must be a cell reference to a pointer to `[u32; 16]` as the message.
509    /// `finalize` should be `true` if this is the final compression.
510    ///
511    /// Additionally the pointer to the output hash state should be on the top of the stack.
512    pub fn blake2s_compress(&mut self, state: Var, byte_count: Var, message: Var, finalize: bool) {
513        let instruction = self.next_instruction(
514            InstructionBody::Blake2sCompress(Blake2sCompressInstruction {
515                state: self.as_cell_ref(state, true),
516                byte_count: self.as_cell_ref(byte_count, true),
517                message: self.as_cell_ref(message, true),
518                finalize,
519            }),
520            true,
521        );
522        assert!(instruction.inc_ap);
523        assert_eq!(
524            self.main_state.allocated as usize, self.main_state.ap_change,
525            "Output var must be top of stack"
526        );
527        self.statements.push(Statement::Final(instruction));
528    }
529
530    /// Adds a label here named `name`.
531    pub fn label(&mut self, name: &'static str) {
532        if self.reachable {
533            self.set_or_test_label_state(name, self.main_state.clone());
534        }
535        let info = self
536            .label_info
537            .get_mut(&name)
538            .unwrap_or_else(|| panic!("No known value for state on reaching {name}."));
539        info.read_only = true;
540        self.main_state = info.state.clone();
541        self.statements.push(Statement::Label(name, 0));
542        self.reachable = true;
543    }
544
545    /// Adds a label `name` in distance `offset` from the current point.
546    /// Useful for calling code outside of the builder's context.
547    pub fn future_label(&mut self, name: &'static str, offset: usize) {
548        self.statements.push(Statement::Label(name, offset));
549    }
550
551    /// Rescoping the values, while ignoring all vars not stated in `vars` and giving the vars on
552    /// the left side the values of the vars on the right side.
553    pub fn rescope<const VAR_COUNT: usize>(&mut self, vars: [(Var, Var); VAR_COUNT]) {
554        self.main_state.validate_finality();
555        let values =
556            vars.map(|(new_var, value_var)| (new_var, self.main_state.get_adjusted(value_var)));
557        self.main_state.ap_change = 0;
558        self.main_state.allocated = 0;
559        self.main_state.vars.clear();
560        self.main_state.vars.extend(values);
561    }
562
563    /// Adds a call command to 'label'. All AP based variables are passed to the called function
564    /// state and dropped from the calling function state.
565    pub fn call(&mut self, label: &'static str) {
566        self.main_state.validate_finality();
567        // Vars to be passed to the called function state.
568        let mut function_vars = SmallOrderedMap::<Var, CellExpression>::default();
569        // FP based vars which will remain in the current state.
570        let mut main_vars = SmallOrderedMap::<Var, CellExpression>::default();
571        let ap_change = self.main_state.ap_change;
572        let cell_to_var_flags = |cell: &CellRef| {
573            if cell.register == Register::AP { (true, false) } else { (false, true) }
574        };
575        for (var, value) in self.main_state.vars.iter() {
576            let (function_var, main_var) = match value {
577                CellExpression::DoubleDeref(cell, _) | CellExpression::Deref(cell) => {
578                    cell_to_var_flags(cell)
579                }
580                CellExpression::Immediate(_) => (true, true),
581                CellExpression::BinOp { op: _, a, b } => match b {
582                    DerefOrImmediate::Deref(cell) => {
583                        if a.register == cell.register {
584                            cell_to_var_flags(cell)
585                        } else {
586                            // Mixed FP and AP based, dropped from both states.
587                            (false, false)
588                        }
589                    }
590                    DerefOrImmediate::Immediate(_) => (true, true),
591                },
592            };
593            if function_var {
594                // Apply ap change (+2 because of the call statement) and change to FP based before
595                // the function call.
596                let mut value = value.clone().unchecked_apply_known_ap_change(ap_change + 2);
597                match &mut value {
598                    CellExpression::DoubleDeref(cell, _) | CellExpression::Deref(cell) => {
599                        cell.register = Register::FP
600                    }
601                    CellExpression::Immediate(_) => {}
602                    CellExpression::BinOp { a, b, .. } => {
603                        a.register = Register::FP;
604                        match b {
605                            DerefOrImmediate::Deref(cell) => cell.register = Register::FP,
606                            DerefOrImmediate::Immediate(_) => {}
607                        }
608                    }
609                }
610                function_vars.insert(*var, value);
611            }
612            if main_var {
613                main_vars.insert(*var, value.clone());
614            }
615        }
616
617        let instruction = self.next_instruction(
618            InstructionBody::Call(CallInstruction {
619                relative: true,
620                target: deref_or_immediate!(0),
621            }),
622            false,
623        );
624        self.statements.push(Statement::Jump(label, instruction));
625
626        self.main_state.vars = main_vars;
627        self.main_state.allocated = 0;
628        self.main_state.ap_change = 0;
629        let function_state = State { vars: function_vars, ..Default::default() };
630        self.set_or_test_label_state(label, function_state);
631    }
632
633    /// A return statement in the code.
634    pub fn ret(&mut self) {
635        self.main_state.validate_finality();
636        let instruction = self.next_instruction(InstructionBody::Ret(RetInstruction {}), false);
637        self.statements.push(Statement::Final(instruction));
638        self.reachable = false;
639    }
640
641    /// The number of steps at the last added statement.
642    pub fn steps(&self) -> usize {
643        self.main_state.steps
644    }
645
646    /// Resets the steps counter.
647    pub fn reset_steps(&mut self) {
648        self.main_state.steps = 0;
649    }
650
651    /// Create an assert that would always fail.
652    pub fn fail(&mut self) {
653        let cell = CellRef { offset: -1, register: Register::FP };
654        let instruction = self.next_instruction(
655            InstructionBody::AssertEq(AssertEqInstruction {
656                a: cell,
657                b: ResOperand::BinOp(BinOpOperand {
658                    op: Operation::Add,
659                    a: cell,
660                    b: DerefOrImmediate::Immediate(BigInt::one().into()),
661                }),
662            }),
663            false,
664        );
665        self.statements.push(Statement::Final(instruction));
666        self.mark_unreachable();
667    }
668
669    /// Marks the current state as unreachable.
670    /// Useful for removing bookkeeping after unsatisfiable conditions.
671    pub fn mark_unreachable(&mut self) {
672        self.reachable = false;
673    }
674
675    /// Returns `var`s value, with fixed ap if `adjust_ap` is true.
676    pub fn get_value(&self, var: Var, adjust_ap: bool) -> CellExpression {
677        if adjust_ap { self.main_state.get_adjusted(var) } else { self.main_state.get_value(var) }
678    }
679
680    /// Returns `var`s value as a cell reference, with fixed ap if `adjust_ap` is true.
681    fn as_cell_ref(&self, var: Var, adjust_ap: bool) -> CellRef {
682        extract_matches!(self.get_value(var, adjust_ap), CellExpression::Deref)
683    }
684
685    /// Returns `var`s value as a cell reference or immediate, with fixed ap if `adjust_ap` is true.
686    fn as_deref_or_imm(&self, var: Var, adjust_ap: bool) -> DerefOrImmediate {
687        match self.get_value(var, adjust_ap) {
688            CellExpression::Deref(cell) => DerefOrImmediate::Deref(cell),
689            CellExpression::Immediate(imm) => DerefOrImmediate::Immediate(imm.into()),
690            CellExpression::DoubleDeref(_, _) | CellExpression::BinOp { .. } => {
691                panic!("wrong usage.")
692            }
693        }
694    }
695
696    /// Returns `var`s value as a cell reference plus a small const offset, with fixed ap if
697    /// `adjust_ap` is true.
698    fn as_cell_ref_plus_const(
699        &self,
700        var: Var,
701        additional_offset: i16,
702        adjust_ap: bool,
703    ) -> (CellRef, i16) {
704        match self.get_value(var, adjust_ap) {
705            CellExpression::Deref(cell) => (cell, additional_offset),
706            CellExpression::BinOp {
707                op: CellOperator::Add,
708                a,
709                b: DerefOrImmediate::Immediate(imm),
710            } => (
711                a,
712                (imm.value + additional_offset).try_into().expect("Offset too large for deref."),
713            ),
714            _ => panic!("Not a valid ptr."),
715        }
716    }
717
718    /// Returns an instruction wrapping the instruction body, and updates the state.
719    /// If `inc_ap_supported` may add an `ap++` to the instruction.
720    fn next_instruction(&mut self, body: InstructionBody, inc_ap_supported: bool) -> Instruction {
721        assert!(self.reachable, "Cannot add instructions at unreachable code.");
722        let inc_ap =
723            inc_ap_supported && self.main_state.allocated as usize > self.main_state.ap_change;
724        if inc_ap {
725            self.main_state.ap_change += 1;
726        }
727        self.main_state.steps += 1;
728        let mut hints = vec![];
729        core::mem::swap(&mut hints, &mut self.current_hints);
730        Instruction { body, inc_ap, hints }
731    }
732}
733
734impl Default for CasmBuilder {
735    fn default() -> Self {
736        Self {
737            label_info: Default::default(),
738            main_state: Default::default(),
739            statements: Default::default(),
740            current_hints: Default::default(),
741            var_count: Default::default(),
742            reachable: true,
743        }
744    }
745}
746
747#[macro_export]
748macro_rules! casm_build_extend {
749    ($builder:expr,) => {};
750    ($builder:expr, tempvar $var:ident; $($tok:tt)*) => {
751        let $var = $builder.alloc_var(false);
752        $crate::casm_build_extend!($builder, $($tok)*)
753    };
754    ($builder:expr, localvar $var:ident; $($tok:tt)*) => {
755        let $var = $builder.alloc_var(true);
756        $crate::casm_build_extend!($builder, $($tok)*)
757    };
758    ($builder:expr, ap += $value:expr; $($tok:tt)*) => {
759        $builder.add_ap($value);
760        $crate::casm_build_extend!($builder, $($tok)*)
761    };
762    ($builder:expr, const $imm:ident = $value:expr; $($tok:tt)*) => {
763        let $imm = $builder.add_var($crate::cell_expression::CellExpression::Immediate(($value).into()));
764        $crate::casm_build_extend!($builder, $($tok)*)
765    };
766    ($builder:expr, assert $dst:ident = $res:ident; $($tok:tt)*) => {
767        $builder.assert_vars_eq($dst, $res, $crate::builder::AssertEqKind::Felt252);
768        $crate::casm_build_extend!($builder, $($tok)*)
769    };
770    ($builder:expr, assert $dst:ident = $a:ident + $b:ident; $($tok:tt)*) => {
771        {
772            let __sum = $builder.bin_op($crate::cell_expression::CellOperator::Add, $a, $b);
773            $builder.assert_vars_eq($dst, __sum, $crate::builder::AssertEqKind::Felt252);
774        }
775        $crate::casm_build_extend!($builder, $($tok)*)
776    };
777    ($builder:expr, assert $dst:ident = $a:ident * $b:ident; $($tok:tt)*) => {
778        {
779            let __product = $builder.bin_op($crate::cell_expression::CellOperator::Mul, $a, $b);
780            $builder.assert_vars_eq($dst, __product, $crate::builder::AssertEqKind::Felt252);
781        }
782        $crate::casm_build_extend!($builder, $($tok)*)
783    };
784    ($builder:expr, assert $dst:ident = $a:ident - $b:ident; $($tok:tt)*) => {
785        {
786            let __diff = $builder.bin_op($crate::cell_expression::CellOperator::Sub, $a, $b);
787            $builder.assert_vars_eq($dst, __diff, $crate::builder::AssertEqKind::Felt252);
788        }
789        $crate::casm_build_extend!($builder, $($tok)*)
790    };
791    ($builder:expr, assert $dst:ident = $a:ident / $b:ident; $($tok:tt)*) => {
792        {
793            let __division = $builder.bin_op($crate::cell_expression::CellOperator::Div, $a, $b);
794            $builder.assert_vars_eq($dst, __division, $crate::builder::AssertEqKind::Felt252);
795        }
796        $crate::casm_build_extend!($builder, $($tok)*)
797    };
798    ($builder:expr, assert $dst:ident = $buffer:ident [ $offset:expr ] ; $($tok:tt)*) => {
799        {
800            let __deref = $builder.double_deref($buffer, $offset);
801            $builder.assert_vars_eq($dst, __deref, $crate::builder::AssertEqKind::Felt252);
802        }
803        $crate::casm_build_extend!($builder, $($tok)*)
804    };
805    ($builder:expr, assert $dst:ident = * $buffer:ident; $($tok:tt)*) => {
806        {
807            let __deref = $builder.double_deref($buffer, 0);
808            $builder.assert_vars_eq($dst, __deref, $crate::builder::AssertEqKind::Felt252);
809        }
810        $crate::casm_build_extend!($builder, $($tok)*)
811    };
812    ($builder:expr, assert $value:ident = * ( $buffer:ident ++ ); $($tok:tt)*) => {
813        $builder.buffer_write_and_inc($buffer, $value);
814        $crate::casm_build_extend!($builder, $($tok)*)
815    };
816    ($builder:expr, tempvar $var:ident = $value:ident; $($tok:tt)*) => {
817        $crate::casm_build_extend!($builder, tempvar $var; assert $var = $value; $($tok)*);
818    };
819    ($builder:expr, tempvar $var:ident = $lhs:ident + $rhs:ident; $($tok:tt)*) => {
820        $crate::casm_build_extend!($builder, tempvar $var; assert $var = $lhs + $rhs; $($tok)*);
821    };
822    ($builder:expr, tempvar $var:ident = $lhs:ident * $rhs:ident; $($tok:tt)*) => {
823        $crate::casm_build_extend!($builder, tempvar $var; assert $var = $lhs * $rhs; $($tok)*);
824    };
825    ($builder:expr, tempvar $var:ident = $lhs:ident - $rhs:ident; $($tok:tt)*) => {
826        $crate::casm_build_extend!($builder, tempvar $var; assert $var = $lhs - $rhs; $($tok)*);
827    };
828    ($builder:expr, tempvar $var:ident = $lhs:ident / $rhs:ident; $($tok:tt)*) => {
829        $crate::casm_build_extend!($builder, tempvar $var; assert $var = $lhs / $rhs; $($tok)*);
830    };
831    ($builder:expr, tempvar $var:ident = * ( $buffer:ident ++ ); $($tok:tt)*) => {
832        $crate::casm_build_extend!($builder, tempvar $var; assert $var = *($buffer++); $($tok)*);
833    };
834    ($builder:expr, tempvar $var:ident = $buffer:ident [ $offset:expr ]; $($tok:tt)*) => {
835        $crate::casm_build_extend!($builder, tempvar $var; assert $var = $buffer[$offset]; $($tok)*);
836    };
837    ($builder:expr, tempvar $var:ident = * $buffer:ident ; $($tok:tt)*) => {
838        $crate::casm_build_extend!($builder, tempvar $var; assert $var = *$buffer; $($tok)*);
839    };
840    ($builder:expr, maybe_tempvar $var:ident = $value:ident; $($tok:tt)*) => {
841        let $var = $builder.maybe_add_tempvar($value);
842        $crate::casm_build_extend!($builder, $($tok)*);
843    };
844    ($builder:expr, maybe_tempvar $var:ident = $lhs:ident + $rhs:ident; $($tok:tt)*) => {
845        $crate::casm_build_extend! {$builder,
846            let $var = $lhs + $rhs;
847            maybe_tempvar $var = $var;
848            $($tok)*
849        };
850    };
851    ($builder:expr, maybe_tempvar $var:ident = $lhs:ident * $rhs:ident; $($tok:tt)*) => {
852        $crate::casm_build_extend! {$builder,
853            let $var = $lhs * $rhs;
854            maybe_tempvar $var = $var;
855            $($tok)*
856        };
857    };
858    ($builder:expr, maybe_tempvar $var:ident = $lhs:ident - $rhs:ident; $($tok:tt)*) => {
859        $crate::casm_build_extend! {$builder,
860            let $var = $lhs - $rhs;
861            maybe_tempvar $var = $var;
862            $($tok)*
863        };
864    };
865    ($builder:expr, maybe_tempvar $var:ident = $lhs:ident / $rhs:ident; $($tok:tt)*) => {
866        $crate::casm_build_extend! {$builder,
867            let $var = $lhs / $rhs;
868            maybe_tempvar $var = $var;
869            $($tok)*
870        };
871    };
872    ($builder:expr, localvar $var:ident = $value:ident; $($tok:tt)*) => {
873        $crate::casm_build_extend!($builder, localvar $var; assert $var = $value; $($tok)*);
874    };
875    ($builder:expr, localvar $var:ident = $lhs:ident + $rhs:ident; $($tok:tt)*) => {
876        $crate::casm_build_extend!($builder, localvar $var; assert $var = $lhs + $rhs; $($tok)*);
877    };
878    ($builder:expr, localvar $var:ident = $lhs:ident * $rhs:ident; $($tok:tt)*) => {
879        $crate::casm_build_extend!($builder, localvar $var; assert $var = $lhs * $rhs; $($tok)*);
880    };
881    ($builder:expr, localvar $var:ident = $lhs:ident - $rhs:ident; $($tok:tt)*) => {
882        $crate::casm_build_extend!($builder, localvar $var; assert $var = $lhs - $rhs; $($tok)*);
883    };
884    ($builder:expr, localvar $var:ident = $lhs:ident / $rhs:ident; $($tok:tt)*) => {
885        $crate::casm_build_extend!($builder, localvar $var; assert $var = $lhs / $rhs; $($tok)*);
886    };
887    ($builder:expr, localvar $var:ident = * ( $buffer:ident ++ ); $($tok:tt)*) => {
888        $crate::casm_build_extend!($builder, localvar $var; assert $var = *($buffer++); $($tok)*);
889    };
890    ($builder:expr, localvar $var:ident = $buffer:ident [ $offset:expr ]; $($tok:tt)*) => {
891        $crate::casm_build_extend!($builder, localvar $var; assert $var = $buffer[$offset]; $($tok)*);
892    };
893    ($builder:expr, localvar $var:ident = * $buffer:ident ; $($tok:tt)*) => {
894        $crate::casm_build_extend!($builder, localvar $var; assert $var = *$buffer; $($tok)*);
895    };
896    ($builder:expr, let $dst:ident = $a:ident + $b:ident; $($tok:tt)*) => {
897        let $dst = $builder.bin_op($crate::cell_expression::CellOperator::Add, $a, $b);
898        $crate::casm_build_extend!($builder, $($tok)*)
899    };
900    ($builder:expr, let $dst:ident = $a:ident * $b:ident; $($tok:tt)*) => {
901        let $dst = $builder.bin_op($crate::cell_expression::CellOperator::Mul, $a, $b);
902        $crate::casm_build_extend!($builder, $($tok)*)
903    };
904    ($builder:expr, let $dst:ident = $a:ident - $b:ident; $($tok:tt)*) => {
905        let $dst = $builder.bin_op($crate::cell_expression::CellOperator::Sub, $a, $b);
906        $crate::casm_build_extend!($builder, $($tok)*)
907    };
908    ($builder:expr, let $dst:ident = $a:ident / $b:ident; $($tok:tt)*) => {
909        let $dst = $builder.bin_op($crate::cell_expression::CellOperator::Div, $a, $b);
910        $crate::casm_build_extend!($builder, $($tok)*)
911    };
912    ($builder:expr, let $dst:ident = * ( $buffer:ident ++ ); $($tok:tt)*) => {
913        let $dst = $builder.get_ref_and_inc($buffer);
914        $crate::casm_build_extend!($builder, $($tok)*)
915    };
916    ($builder:expr, let $dst:ident = $buffer:ident [ $offset:expr ] ; $($tok:tt)*) => {
917        let $dst = $builder.double_deref($buffer, $offset);
918        $crate::casm_build_extend!($builder, $($tok)*)
919    };
920    ($builder:expr, let $dst:ident = *$buffer:ident; $($tok:tt)*) => {
921        let $dst = $builder.double_deref($buffer, 0);
922        $crate::casm_build_extend!($builder, $($tok)*)
923    };
924    ($builder:expr, let $dst:ident = $src:ident; $($tok:tt)*) => {
925        let $dst = $builder.duplicate_var($src);
926        $crate::casm_build_extend!($builder, $($tok)*)
927    };
928    ($builder:expr, jump $target:ident; $($tok:tt)*) => {
929        $builder.jump(core::stringify!($target));
930        $crate::casm_build_extend!($builder, $($tok)*)
931    };
932    ($builder:expr, jump $target:ident if $condition:ident != 0; $($tok:tt)*) => {
933        $builder.jump_nz($condition, core::stringify!($target));
934        $crate::casm_build_extend!($builder, $($tok)*)
935    };
936    ($builder:expr, let ($($var_name:ident),*) = call $target:ident; $($tok:tt)*) => {
937        $builder.call(core::stringify!($target));
938
939        let __var_count = {0i16 $(+ (stringify!($var_name), 1i16).1)*};
940        let mut __var_index = 0;
941        $(
942            let $var_name = $builder.add_var($crate::cell_expression::CellExpression::Deref($crate::operand::CellRef {
943                offset: __var_index - __var_count,
944                register: $crate::operand::Register::AP,
945            }));
946            __var_index += 1;
947        )*
948        $crate::casm_build_extend!($builder, $($tok)*)
949    };
950    ($builder:expr, ret; $($tok:tt)*) => {
951        $builder.ret();
952        $crate::casm_build_extend!($builder, $($tok)*)
953    };
954    ($builder:expr, $label:ident: $($tok:tt)*) => {
955        $builder.label(core::stringify!($label));
956        $crate::casm_build_extend!($builder, $($tok)*)
957    };
958    ($builder:expr, fail; $($tok:tt)*) => {
959        $builder.fail();
960        $crate::casm_build_extend!($builder, $($tok)*)
961    };
962    ($builder:expr, unsatisfiable_assert $dst:ident = $res:ident; $($tok:tt)*) => {
963        $crate::casm_build_extend!($builder, assert $dst = $res;);
964        $builder.mark_unreachable();
965        $crate::casm_build_extend!($builder, $($tok)*)
966    };
967    ($builder:expr, hint $hint_head:ident$(::$hint_tail:ident)+ {
968            $($input_name:ident $(: $input_value:ident)?),*
969        } into {
970            $($output_name:ident $(: $output_value:ident)?),*
971        }; $($tok:tt)*) => {
972        $builder.add_hint(
973            |[$($input_name),*], [$($output_name),*]| $hint_head$(::$hint_tail)+ {
974                $($input_name,)* $($output_name,)*
975            },
976            [$($crate::casm_build_hint_param_value!($input_name  $(: $input_value)?),)*],
977            [$($crate::casm_build_hint_param_value!($output_name  $(: $output_value)?),)*],
978        );
979        $crate::casm_build_extend!($builder, $($tok)*)
980    };
981    ($builder:expr, hint $hint_name:ident { $($inputs:tt)* } into { $($outputs:tt)* }; $($tok:tt)*) => {
982        $crate::casm_build_extend! {$builder,
983            hint $crate::hints::CoreHint::$hint_name { $($inputs)* } into { $($outputs)* };
984            $($tok)*
985        }
986    };
987    ($builder:expr, hint $hint_head:ident$(::$hint_tail:ident)* { $($inputs:tt)* }; $($tok:tt)*) => {
988        $crate::casm_build_extend! {$builder,
989            hint $hint_head$(::$hint_tail)* { $($inputs)* } into {};
990            $($tok)*
991        }
992    };
993    ($builder:expr, hint $hint_head:ident$(::$hint_tail:ident)* into { $($outputs:tt)* }; $($tok:tt)*) => {
994        $crate::casm_build_extend! {$builder,
995            hint $hint_head$(::$hint_tail)* {} into { $($outputs)* };
996            $($tok)*
997        }
998    };
999    ($builder:expr, rescope { $($new_var:ident = $value_var:ident),* }; $($tok:tt)*) => {
1000        $builder.rescope([$(($new_var, $value_var)),*]);
1001        $crate::casm_build_extend!($builder, $($tok)*)
1002    };
1003    ($builder:expr, #{ validate steps == $count:expr; } $($tok:tt)*) => {
1004        assert_eq!($builder.steps(), $count);
1005        $crate::casm_build_extend!($builder, $($tok)*)
1006    };
1007    ($builder:expr, #{ steps = 0; } $($tok:tt)*) => {
1008        $builder.reset_steps();
1009        $crate::casm_build_extend!($builder, $($tok)*)
1010    };
1011    ($builder:expr, #{ $counter:ident += steps; steps = 0; } $($tok:tt)*) => {
1012        $counter += $builder.steps() as i32;
1013        $builder.reset_steps();
1014        $crate::casm_build_extend!($builder, $($tok)*)
1015    };
1016}
1017
1018#[macro_export]
1019macro_rules! casm_build_hint_param_value {
1020    ($_name:ident : $value:ident) => {
1021        $value
1022    };
1023    ($name:ident) => {
1024        $name
1025    };
1026}