feat: implement table-based MAST (#1349)

CHANGELOG.md

@@ -6,6 +6,7 @@
 
 - Added error codes support for the `mtree_verify` instruction (#1328).
 - Added support for immediate values for `lt`, `lte`, `gt`, `gte` comparison instructions (#1346).
+- Change MAST to a table-based representation (#1349)
 
 ## 0.9.2 (2024-05-22) - `stdlib` crate only
 - Skip writing MASM documentation to file when building on docs.rs (#1341).

air/src/constraints/chiplets/hasher/mod.rs

@@ -100,8 +100,8 @@ pub fn get_transition_constraint_count() -> usize {
 
 /// Enforces constraints for the hasher chiplet.
 ///
-/// - The `hasher_flag` determines if the hasher chiplet is currently enabled. It should be
-///   computed by the caller and set to `Felt::ONE`
+/// - The `hasher_flag` determines if the hasher chiplet is currently enabled. It should be computed
+///   by the caller and set to `Felt::ONE`
 /// - The `transition_flag` indicates whether this is the last row this chiplet's execution trace,
 ///   and therefore the constraints should not be enforced.
 pub fn enforce_constraints<E: FieldElement<BaseField = Felt>>(

air/src/constraints/chiplets/memory/tests.rs

@@ -106,7 +106,7 @@ enum MemoryTestDeltaType {
 /// - To test a valid write, the MemoryTestDeltaType must be Context or Address and the `old_values`
 ///   and `new_values` must change.
 /// - To test a valid read, the `delta_type` must be Clock and the `old_values` and `new_values`
-/// must be equal.
+///   must be equal.
 fn get_constraint_evaluation(
     selectors: Selectors,
     delta_type: MemoryTestDeltaType,

air/src/constraints/stack/field_ops/mod.rs

@@ -187,8 +187,8 @@ pub fn enforce_incr_constraints<E: FieldElement>(
 /// in the stack with its bitwise not value. Therefore, the following constraints are
 /// enforced:
 /// - The top element should be a binary. It is enforced as a general constraint.
-/// - The first element of the next frame should be a binary not of the first element of
-///   the current frame. s0` + s0 = 1.
+/// - The first element of the next frame should be a binary not of the first element of the current
+///   frame. s0` + s0 = 1.
 pub fn enforce_not_constraints<E: FieldElement>(
     frame: &EvaluationFrame<E>,
     result: &mut [E],
@@ -206,8 +206,8 @@ pub fn enforce_not_constraints<E: FieldElement>(
 
 /// Enforces constraints of the AND operation. The AND operation computes the bitwise and of the
 /// first two elements in the current trace. Therefore, the following constraints are enforced:
-/// - The top two element in the current frame of the stack should be binary. s0^2 - s0 = 0,
-///   s1^2 - s1 = 0. The top element is binary or not is enforced as a general constraint.
+/// - The top two element in the current frame of the stack should be binary. s0^2 - s0 = 0, s1^2 -
+///   s1 = 0. The top element is binary or not is enforced as a general constraint.
 /// - The first element of the next frame should be a binary and of the first two elements in the
 ///   current frame. s0` - s0 * s1 = 0.
 pub fn enforce_and_constraints<E: FieldElement>(
@@ -233,8 +233,8 @@ pub fn enforce_and_constraints<E: FieldElement>(
 
 /// Enforces constraints of the OR operation. The OR operation computes the bitwise or of the
 /// first two elements in the current trace. Therefore, the following constraints are enforced:
-/// - The top two element in the current frame of the stack should be binary. s0^2 - s0 = 0,
-///   s1^2 - s1 = 0. The top element is binary or not is enforced as a general constraint.
+/// - The top two element in the current frame of the stack should be binary. s0^2 - s0 = 0, s1^2 -
+///   s1 = 0. The top element is binary or not is enforced as a general constraint.
 /// - The first element of the next frame should be a binary or of the first two elements in the
 ///   current frame. s0` - ( s0 + s1 - s0 * s1 ) = 0.
 pub fn enforce_or_constraints<E: FieldElement>(

air/src/constraints/stack/overflow/mod.rs

@@ -92,8 +92,7 @@ pub fn enforce_stack_depth_constraints<E: FieldElement>(
 
 /// Enforces constraints on the overflow flag h0. Therefore, the following constraints
 /// are enforced:
-/// - If overflow table has values, then, h0 should be set to ONE, otherwise it should
-///   be ZERO.
+/// - If overflow table has values, then, h0 should be set to ONE, otherwise it should be ZERO.
 pub fn enforce_overflow_flag_constraints<E: FieldElement>(
     frame: &EvaluationFrame<E>,
     result: &mut [E],
@@ -107,8 +106,8 @@ pub fn enforce_overflow_flag_constraints<E: FieldElement>(
 }
 
 /// Enforces constraints on the bookkeeping index `b1`. The following constraints are enforced:
-/// - In the case of a right shift operation, the next b1 index should be updated with current
-///   `clk` value.
+/// - In the case of a right shift operation, the next b1 index should be updated with current `clk`
+///   value.
 /// - In the case of a left shift operation, the last stack item should be set to ZERO when the
 ///   depth of the stack is 16.
 pub fn enforce_overflow_index_constraints<E: FieldElement>(

air/src/constraints/stack/stack_manipulation/mod.rs

@@ -93,8 +93,8 @@ pub fn enforce_pad_constraints<E: FieldElement>(
 /// Enforces constraints of the DUPn and MOVUPn operations. The DUPn operation copies the element
 /// at depth n in the stack and pushes the copy onto the stack, whereas MOVUPn opearation moves the
 /// element at depth n to the top of the stack. Therefore, the following constraints are enforced:
-/// - The top element in the next frame should be equal to the element at depth n in the
-///   current frame. s0` - sn = 0.
+/// - The top element in the next frame should be equal to the element at depth n in the current
+///   frame. s0` - sn = 0.
 pub fn enforce_dup_movup_n_constraints<E: FieldElement>(
     frame: &EvaluationFrame<E>,
     result: &mut [E],
@@ -244,8 +244,8 @@ pub fn enforce_swapwx_constraints<E: FieldElement>(
 
 /// Enforces constraints of the MOVDNn operation. The MOVDNn operation moves the top element
 /// to depth n in the stack. Therefore, the following constraints are enforced:
-/// - The top element in the current frame should be equal to the element at depth n in the
-///   next frame. s0 - sn` = 0.
+/// - The top element in the current frame should be equal to the element at depth n in the next
+///   frame. s0 - sn` = 0.
 pub fn enforce_movdnn_constraints<E: FieldElement>(
     frame: &EvaluationFrame<E>,
     result: &mut [E],

air/src/constraints/stack/u32_ops/mod.rs

@@ -119,8 +119,8 @@ pub fn enforce_u32split_constraints<E: FieldElement<BaseField = Felt>>(
 /// Enforces constraints of the U32ADD operation. The U32ADD operation adds the top two
 /// elements in the current trace of the stack. Therefore, the following constraints are
 /// enforced:
-/// - The aggregation of limbs from the helper registers is equal to the sum of the top two
-///   element in the stack.
+/// - The aggregation of limbs from the helper registers is equal to the sum of the top two element
+///   in the stack.
 pub fn enforce_u32add_constraints<E: FieldElement<BaseField = Felt>>(
     frame: &EvaluationFrame<E>,
     result: &mut [E],

air/src/options.rs

@@ -227,8 +227,8 @@ impl ExecutionOptions {
     ///
     /// In debug mode the VM does the following:
     /// - Executes `debug` instructions (these are ignored in regular mode).
-    /// - Records additional info about program execution (e.g., keeps track of stack state at
-    ///   every cycle of the VM) which enables stepping through the program forward and backward.
+    /// - Records additional info about program execution (e.g., keeps track of stack state at every
+    ///   cycle of the VM) which enables stepping through the program forward and backward.
     pub fn with_debugging(mut self) -> Self {
         self.enable_debugging = true;
         self

assembly/src/assembler/span_builder.rs → assembly/src/assembler/basic_block_builder.rs

@@ -1,8 +1,11 @@
 use super::{AssemblyContext, BodyWrapper, Decorator, DecoratorList, Instruction};
 use alloc::{borrow::Borrow, string::ToString, vec::Vec};
-use vm_core::{code_blocks::CodeBlock, AdviceInjector, AssemblyOp, Operation};
+use vm_core::{
+    mast::{MastForest, MastNode, MastNodeId},
+    AdviceInjector, AssemblyOp, Operation,
+};
 
-// SPAN BUILDER
+// BASIC BLOCK BUILDER
 // ================================================================================================
 
 /// A helper struct for constructing SPAN blocks while compiling procedure bodies.
@@ -13,15 +16,15 @@ use vm_core::{code_blocks::CodeBlock, AdviceInjector, AssemblyOp, Operation};
 /// The same span builder can be used to construct many blocks. It is expected that when the last
 /// SPAN block in a procedure's body is constructed `extract_final_span_into()` will be used.
 #[derive(Default)]
-pub struct SpanBuilder {
+pub struct BasicBlockBuilder {
     ops: Vec<Operation>,
     decorators: DecoratorList,
     epilogue: Vec<Operation>,
     last_asmop_pos: usize,
 }
 
 /// Constructors
-impl SpanBuilder {
+impl BasicBlockBuilder {
     /// Returns a new [SpanBuilder] instantiated with the specified optional wrapper.
     ///
     /// If the wrapper is provided, the prologue of the wrapper is immediately appended to the
@@ -41,7 +44,7 @@ impl SpanBuilder {
 }
 
 /// Operations
-impl SpanBuilder {
+impl BasicBlockBuilder {
     /// Adds the specified operation to the list of span operations.
     pub fn push_op(&mut self, op: Operation) {
         self.ops.push(op);
@@ -64,8 +67,8 @@ impl SpanBuilder {
 }
 
 /// Decorators
-impl SpanBuilder {
-    /// Add ths specified decorator to the list of span decorators.
+impl BasicBlockBuilder {
+    /// Add the specified decorator to the list of span decorators.
     pub fn push_decorator(&mut self, decorator: Decorator) {
         self.decorators.push((self.ops.len(), decorator));
     }
@@ -114,34 +117,40 @@ impl SpanBuilder {
 }
 
 /// Span Constructors
-impl SpanBuilder {
-    /// Creates a new SPAN block from the operations and decorators currently in this builder and
-    /// appends the block to the provided target.
+impl BasicBlockBuilder {
+    /// Creates and returns a new BASIC BLOCK node from the operations and decorators currently in
+    /// this builder. If the builder is empty, then no node is created and `None` is returned.
     ///
     /// This consumes all operations and decorators in the builder, but does not touch the
     /// operations in the epilogue of the builder.
-    pub fn extract_span_into(&mut self, target: &mut Vec<CodeBlock>) {
+    pub fn make_basic_block(&mut self, mast_forest: &mut MastForest) -> Option<MastNodeId> {
         if !self.ops.is_empty() {
             let ops = self.ops.drain(..).collect();
             let decorators = self.decorators.drain(..).collect();
-            target.push(CodeBlock::new_span_with_decorators(ops, decorators));
+
+            let basic_block_node = MastNode::new_basic_block_with_decorators(ops, decorators);
+            let basic_block_node_id = mast_forest.ensure_node(basic_block_node);
+
+            Some(basic_block_node_id)
         } else if !self.decorators.is_empty() {
             // this is a bug in the assembler. we shouldn't have decorators added without their
             // associated operations
             // TODO: change this to an error or allow decorators in empty span blocks
             unreachable!("decorators in an empty SPAN block")
+        } else {
+            None
         }
     }
 
-    /// Creates a new SPAN block from the operations and decorators currently in this builder and
-    /// appends the block to the provided target.
+    /// Creates and returns a new BASIC BLOCK node from the operations and decorators currently in
+    /// this builder. If the builder is empty, then no node is created and `None` is returned.
     ///
-    /// The main differences from the `extract_span_int()` method above are:
-    /// - Operations contained in the epilogue of the span builder are appended to the list of ops
-    ///   which go into the new SPAN block.
-    /// - The span builder is consumed in the process.
-    pub fn extract_final_span_into(mut self, target: &mut Vec<CodeBlock>) {
+    /// The main differences with [`Self::to_basic_block`] are:
+    /// - Operations contained in the epilogue of the builder are appended to the list of ops which
+    ///   go into the new BASIC BLOCK node.
+    /// - The builder is consumed in the process.
+    pub fn into_basic_block(mut self, mast_forest: &mut MastForest) -> Option<MastNodeId> {
         self.ops.append(&mut self.epilogue);
-        self.extract_span_into(target);
+        self.make_basic_block(mast_forest)
     }
 }

assembly/src/assembler/context.rs

@@ -6,7 +6,7 @@ use crate::{
     diagnostics::SourceFile,
     AssemblyError, LibraryPath, RpoDigest, SourceSpan, Span, Spanned,
 };
-use vm_core::code_blocks::CodeBlock;
+use vm_core::mast::{MastForest, MastNodeId};
 
 // ASSEMBLY CONTEXT
 // ================================================================================================
@@ -168,6 +168,7 @@ impl AssemblyContext {
         &mut self,
         callee: &Procedure,
         inlined: bool,
+        mast_forest: &MastForest,
     ) -> Result<(), AssemblyError> {
         let context = self.unwrap_current_procedure_mut();
 
@@ -176,7 +177,7 @@ impl AssemblyContext {
 
         // If the callee is not being inlined, add it to our callset
         if !inlined {
-            context.insert_callee(callee.mast_root());
+            context.insert_callee(callee.mast_root(mast_forest));
         }
 
         Ok(())
@@ -264,11 +265,12 @@ impl ProcedureContext {
         self.visibility.is_syscall()
     }
 
-    pub fn into_procedure(self, code: CodeBlock) -> Box<Procedure> {
-        let procedure = Procedure::new(self.name, self.visibility, self.num_locals as u32, code)
-            .with_span(self.span)
-            .with_source_file(self.source_file)
-            .with_callset(self.callset);
+    pub fn into_procedure(self, body_node_id: MastNodeId) -> Box<Procedure> {
+        let procedure =
+            Procedure::new(self.name, self.visibility, self.num_locals as u32, body_node_id)
+                .with_span(self.span)
+                .with_source_file(self.source_file)
+                .with_callset(self.callset);
         Box::new(procedure)
     }
 }

assembly/src/assembler/instruction/adv_ops.rs

@@ -1,4 +1,4 @@
-use super::{validate_param, SpanBuilder};
+use super::{validate_param, BasicBlockBuilder};
 use crate::{ast::AdviceInjectorNode, AssemblyError, ADVICE_READ_LIMIT};
 use vm_core::Operation;
 
@@ -12,7 +12,7 @@ use vm_core::Operation;
 /// # Errors
 /// Returns an error if the specified number of values to pushed is smaller than 1 or greater
 /// than 16.
-pub fn adv_push(span: &mut SpanBuilder, n: u8) -> Result<(), AssemblyError> {
+pub fn adv_push(span: &mut BasicBlockBuilder, n: u8) -> Result<(), AssemblyError> {
     validate_param(n, 1..=ADVICE_READ_LIMIT)?;
     span.push_op_many(Operation::AdvPop, n as usize);
     Ok(())
@@ -22,6 +22,6 @@ pub fn adv_push(span: &mut SpanBuilder, n: u8) -> Result<(), AssemblyError> {
 // ================================================================================================
 
 /// Appends advice injector decorator to the span.
-pub fn adv_inject(span: &mut SpanBuilder, injector: &AdviceInjectorNode) {
+pub fn adv_inject(span: &mut BasicBlockBuilder, injector: &AdviceInjectorNode) {
     span.push_advice_injector(injector.into());
 }

assembly/src/assembler/instruction/crypto_ops.rs

@@ -1,4 +1,4 @@
-use super::SpanBuilder;
+use super::BasicBlockBuilder;
 use vm_core::{AdviceInjector, Operation::*};
 
 // HASHING
@@ -23,7 +23,7 @@ use vm_core::{AdviceInjector, Operation::*};
 /// 3. Drop D and B to achieve our result [C, ...]
 ///
 /// This operation takes 20 VM cycles.
-pub(super) fn hash(span: &mut SpanBuilder) {
+pub(super) fn hash(span: &mut BasicBlockBuilder) {
     #[rustfmt::skip]
     let ops = [
         // add 4 elements to the stack to be used as the capacity elements for the RPO permutation
@@ -69,7 +69,7 @@ pub(super) fn hash(span: &mut SpanBuilder) {
 /// 4. Drop F and D to return our result [E, ...].
 ///
 /// This operation takes 16 VM cycles.
-pub(super) fn hmerge(span: &mut SpanBuilder) {
+pub(super) fn hmerge(span: &mut BasicBlockBuilder) {
     #[rustfmt::skip]
     let ops = [
         // Add 4 elements to the stack to prepare the capacity portion for the RPO permutation
@@ -110,7 +110,7 @@ pub(super) fn hmerge(span: &mut SpanBuilder) {
 /// - root of the tree, 4 elements.
 ///
 /// This operation takes 9 VM cycles.
-pub(super) fn mtree_get(span: &mut SpanBuilder) {
+pub(super) fn mtree_get(span: &mut BasicBlockBuilder) {
     // stack: [d, i, R, ...]
     // pops the value of the node we are looking for from the advice stack
     read_mtree_node(span);
@@ -140,7 +140,7 @@ pub(super) fn mtree_get(span: &mut SpanBuilder) {
 /// - new root of the tree after the update, 4 elements
 ///
 /// This operation takes 29 VM cycles.
-pub(super) fn mtree_set(span: &mut SpanBuilder) {
+pub(super) fn mtree_set(span: &mut BasicBlockBuilder) {
     // stack: [d, i, R_old, V_new, ...]
 
     // stack: [V_old, R_new, ...] (29 cycles)
@@ -160,7 +160,7 @@ pub(super) fn mtree_set(span: &mut SpanBuilder) {
 /// It is not checked whether the provided roots exist as Merkle trees in the advide providers.
 ///
 /// This operation takes 16 VM cycles.
-pub(super) fn mtree_merge(span: &mut SpanBuilder) {
+pub(super) fn mtree_merge(span: &mut BasicBlockBuilder) {
     // stack input:  [R_rhs, R_lhs, ...]
     // stack output: [R_merged, ...]
 
@@ -192,7 +192,7 @@ pub(super) fn mtree_merge(span: &mut SpanBuilder) {
 /// - new value of the node, 4 elements (only in the case of mtree_set)
 ///
 /// This operation takes 4 VM cycles.
-fn read_mtree_node(span: &mut SpanBuilder) {
+fn read_mtree_node(span: &mut BasicBlockBuilder) {
     // The stack should be arranged in the following way: [d, i, R, ...] so that the decorator
     // can fetch the node value from the root. In the `mtree.get` operation we have the stack in
     // the following format: [d, i, R], whereas in the case of `mtree.set` we would also have the
@@ -210,7 +210,7 @@ fn read_mtree_node(span: &mut SpanBuilder) {
 /// and perform the mutation on the copied tree.
 ///
 /// This operation takes 29 VM cycles.
-fn update_mtree(span: &mut SpanBuilder) {
+fn update_mtree(span: &mut BasicBlockBuilder) {
     // stack: [d, i, R_old, V_new, ...]
     // output: [R_new, R_old, V_new, V_old, ...]