Implement table-based MAST

assembly/src/assembler/span_builder.rs

@@ -1,10 +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::{code_blocks::CodeBlock, mast::MastNode, AdviceInjector, AssemblyOp, Operation};
 
 // SPAN BUILDER
 // ================================================================================================
 
+// TODOP: Rename `BasicBlockBuilder`
 /// A helper struct for constructing SPAN blocks while compiling procedure bodies.
 ///
 /// Operations and decorators can be added to a span builder via various `add_*()` and `push_*()`
@@ -65,7 +66,7 @@
 
 /// Decorators
 impl SpanBuilder {
-    /// Add ths specified decorator to the list of span decorators.
+    /// 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));
     }
@@ -144,4 +145,37 @@
         self.ops.append(&mut self.epilogue);
         self.extract_span_into(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.
+    ///
+    /// This consumes all operations and decorators in the builder, but does not touch the
+    /// operations in the epilogue of the builder.
+    pub fn to_basic_block(&mut self) -> Option<MastNode> {
+        if !self.ops.is_empty() {
+            let ops = self.ops.drain(..).collect();
+            let decorators = self.decorators.drain(..).collect();
+
+            Some(MastNode::new_basic_block_with_decorators(ops, decorators))
+        } 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 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 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) -> Option<MastNode> {
+        self.ops.append(&mut self.epilogue);
+        self.to_basic_block()
+    }
 }

core/src/lib.rs

@@ -82,6 +82,8 @@ pub mod crypto {
     }
 }
 
+pub mod mast;
+
 pub use math::{
     fields::{f64::BaseElement as Felt, QuadExtension},
     polynom, ExtensionOf, FieldElement, StarkField, ToElements,
@@ -91,7 +93,7 @@ pub mod prettier {
     pub use miden_formatting::{prettier::*, pretty_via_display, pretty_via_to_string};
 }
 
-mod program;
+pub mod program;
 pub use program::{blocks as code_blocks, CodeBlockTable, Kernel, Program, ProgramInfo};
 
 mod operations;

core/src/mast/basic_block_node.rs

@@ -0,0 +1,94 @@
+use alloc::vec::Vec;
+use miden_crypto::hash::rpo::RpoDigest;
+
+use crate::{
+    code_blocks::{batch_ops, OpBatch},
+    DecoratorIterator, DecoratorList, Operation,
+};
+
+use super::MerkleTreeNode;
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct BasicBlockNode {
+    /// The primitive operations contained in this basic block.
+    ///
+    /// The operations are broken up into batches of 8 groups,
+    /// with each group containing up to 9 operations, or a
+    /// single immediates. Thus the maximum size of each batch
+    /// is 72 operations. Multiple batches are used for blocks
+    /// consisting of more than 72 operations.
+    op_batches: Vec<OpBatch>,
+    digest: RpoDigest,
+    decorators: DecoratorList,
+}
+
+/// Constructors
+impl BasicBlockNode {
+    /// Returns a new [`BasicBlockNode`] instantiated with the specified operations.
+    ///
+    /// # Errors (TODO)
+    /// Returns an error if:
+    /// - `operations` vector is empty.
+    /// - `operations` vector contains any number of system operations.
+    pub fn new(operations: Vec<Operation>) -> Self {
+        assert!(!operations.is_empty()); // TODO: return error
+        Self::with_decorators(operations, DecoratorList::new())
+    }
+
+    /// Returns a new [`BasicBlockNode`] instantiated with the specified operations and decorators.
+    ///
+    /// # Errors (TODO)
+    /// Returns an error if:
+    /// - `operations` vector is empty.
+    /// - `operations` vector contains any number of system operations.
+    pub fn with_decorators(operations: Vec<Operation>, decorators: DecoratorList) -> Self {
+        assert!(!operations.is_empty()); // TODO: return error
+
+        // validate decorators list (only in debug mode)
+        #[cfg(debug_assertions)]
+        validate_decorators(&operations, &decorators);
+
+        let (op_batches, digest) = batch_ops(operations);
+        Self {
+            op_batches,
+            digest,
+            decorators,
+        }
+    }
+}
+
+impl BasicBlockNode {
+    pub fn op_batches(&self) -> &[OpBatch] {
+        &self.op_batches
+    }
+
+    /// Returns a [`DecoratorIterator`] which allows us to iterate through the decorator list of
+    /// this span block while executing operation batches of this span block
+    pub fn decorator_iter(&self) -> DecoratorIterator {
+        DecoratorIterator::new(&self.decorators)
+    }
+}
+
+impl MerkleTreeNode for BasicBlockNode {
+    fn digest(&self) -> RpoDigest {
+        self.digest
+    }
+}
+
+/// Checks if a given decorators list is valid (only checked in debug mode)
+/// - Assert the decorator list is in ascending order.
+/// - Assert the last op index in decorator list is less than or equal to the number of operations.
+#[cfg(debug_assertions)]
+fn validate_decorators(operations: &[Operation], decorators: &DecoratorList) {
+    if !decorators.is_empty() {
+        // check if decorator list is sorted
+        for i in 0..(decorators.len() - 1) {
+            debug_assert!(decorators[i + 1].0 >= decorators[i].0, "unsorted decorators list");
+        }
+        // assert the last index in decorator list is less than operations vector length
+        debug_assert!(
+            operations.len() >= decorators.last().expect("empty decorators list").0,
+            "last op index in decorator list should be less than or equal to the number of ops"
+        );
+    }
+}

core/src/mast/call_node.rs

@@ -0,0 +1,79 @@
+use miden_crypto::{hash::rpo::RpoDigest, Felt};
+
+use crate::{chiplets::hasher, Operation};
+
+use super::{MastForest, MastNodeId, MerkleTreeNode};
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct CallNode {
+    callee: MastNodeId,
+    is_syscall: bool,
+    digest: RpoDigest,
+}
+
+/// Constants
+impl CallNode {
+    /// The domain of the call block (used for control block hashing).
+    pub const CALL_DOMAIN: Felt = Felt::new(Operation::Call.op_code() as u64);
+    /// The domain of the syscall block (used for control block hashing).
+    pub const SYSCALL_DOMAIN: Felt = Felt::new(Operation::SysCall.op_code() as u64);
+}
+
+/// Constructors
+impl CallNode {
+    /// Returns a new [`CallNode`] instantiated with the specified callee.
+    pub fn new(callee: MastNodeId, mast_forest: &MastForest) -> Self {
+        let digest = {
+            let callee_digest = mast_forest.get_node_by_id(callee).digest();
+
+            hasher::merge_in_domain(&[callee_digest, RpoDigest::default()], Self::CALL_DOMAIN)
+        };
+
+        Self {
+            callee,
+            is_syscall: false,
+            digest,
+        }
+    }
+
+    /// Returns a new [`CallNode`] instantiated with the specified callee and marked as a kernel
+    /// call.
+    pub fn new_syscall(callee: MastNodeId, mast_forest: &MastForest) -> Self {
+        let digest = {
+            let callee_digest = mast_forest.get_node_by_id(callee).digest();
+
+            hasher::merge_in_domain(&[callee_digest, RpoDigest::default()], Self::SYSCALL_DOMAIN)
+        };
+
+        Self {
+            callee,
+            is_syscall: true,
+            digest,
+        }
+    }
+}
+
+impl CallNode {
+    pub fn callee(&self) -> MastNodeId {
+        self.callee
+    }
+
+    pub fn is_syscall(&self) -> bool {
+        self.is_syscall
+    }
+
+    /// Returns the domain of the call node.
+    pub fn hash_domain(&self) -> Felt {
+        if self.is_syscall() {
+            Self::SYSCALL_DOMAIN
+        } else {
+            Self::CALL_DOMAIN
+        }
+    }
+}
+
+impl MerkleTreeNode for CallNode {
+    fn digest(&self) -> RpoDigest {
+        self.digest
+    }
+}

core/src/mast/dyn_node.rs

@@ -0,0 +1,28 @@
+use miden_crypto::{hash::rpo::RpoDigest, Felt};
+
+use crate::Operation;
+
+use super::MerkleTreeNode;
+
+#[derive(Debug, Clone, Default, PartialEq, Eq)]
+pub struct DynNode;
+
+/// Constants
+impl DynNode {
+    /// The domain of the Dyn block (used for control block hashing).
+    pub const DOMAIN: Felt = Felt::new(Operation::Dyn.op_code() as u64);
+}
+
+impl MerkleTreeNode for DynNode {
+    fn digest(&self) -> RpoDigest {
+        // The Dyn node is represented by a constant, which is set to be the hash of two empty
+        // words ([ZERO, ZERO, ZERO, ZERO]) with a domain value of `DYN_DOMAIN`, i.e.
+        // hasher::merge_in_domain(&[Digest::default(), Digest::default()], DynNode::DOMAIN)
+        RpoDigest::new([
+            Felt::new(8115106948140260551),
+            Felt::new(13491227816952616836),
+            Felt::new(15015806788322198710),
+            Felt::new(16575543461540527115),
+        ])
+    }
+}

core/src/mast/join_node.rs

@@ -0,0 +1,49 @@
+use miden_crypto::{hash::rpo::RpoDigest, Felt};
+
+use crate::{chiplets::hasher, Operation};
+
+use super::{MastForest, MastNodeId, MerkleTreeNode};
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct JoinNode {
+    children: [MastNodeId; 2],
+    digest: RpoDigest,
+}
+
+/// Constants
+impl JoinNode {
+    /// The domain of the join block (used for control block hashing).
+    pub const DOMAIN: Felt = Felt::new(Operation::Join.op_code() as u64);
+}
+
+/// Constructors
+impl JoinNode {
+    /// Returns a new [`JoinNode`] instantiated with the specified children nodes.
+    pub fn new(children: [MastNodeId; 2], mast_forest: &MastForest) -> Self {
+        let digest = {
+            let left_child_hash = mast_forest.get_node_by_id(children[0]).digest();
+            let right_child_hash = mast_forest.get_node_by_id(children[1]).digest();
+
+            hasher::merge_in_domain(&[left_child_hash, right_child_hash], Self::DOMAIN)
+        };
+
+        Self { children, digest }
+    }
+}
+
+/// Accessors
+impl JoinNode {
+    pub fn first(&self) -> MastNodeId {
+        self.children[0]
+    }
+
+    pub fn second(&self) -> MastNodeId {
+        self.children[1]
+    }
+}
+
+impl MerkleTreeNode for JoinNode {
+    fn digest(&self) -> RpoDigest {
+        self.digest
+    }
+}

core/src/mast/loop_node.rs

@@ -0,0 +1,42 @@
+use miden_crypto::{hash::rpo::RpoDigest, Felt};
+
+use crate::{chiplets::hasher, Operation};
+
+use super::{MastForest, MastNodeId, MerkleTreeNode};
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct LoopNode {
+    body: MastNodeId,
+    digest: RpoDigest,
+}
+
+/// Constants
+impl LoopNode {
+    /// The domain of the loop block (used for control block hashing).
+    pub const DOMAIN: Felt = Felt::new(Operation::Loop.op_code() as u64);
+}
+
+/// Constructors
+impl LoopNode {
+    pub fn new(body: MastNodeId, mast_forest: &MastForest) -> Self {
+        let digest = {
+            let body_hash = mast_forest.get_node_by_id(body).digest();
+
+            hasher::merge_in_domain(&[body_hash, RpoDigest::default()], Self::DOMAIN)
+        };
+
+        Self { body, digest }
+    }
+}
+
+impl LoopNode {
+    pub fn body(&self) -> MastNodeId {
+        self.body
+    }
+}
+
+impl MerkleTreeNode for LoopNode {
+    fn digest(&self) -> RpoDigest {
+        self.digest
+    }
+}