Introduce MastForestStore

CHANGELOG.md

@@ -7,6 +7,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)
+- Introduce `MastForestStore` (#1359)
 - Adjusted prover's metal acceleration code to work with 0.9 versions of the crates (#1357)
 - Added support for immediate values for `u32lt`, `u32lte`, `u32gt`, `u32gte`, `u32min` and `u32max` comparison instructions (#1358).
 - Added support for the `nop` instruction, which corresponds to the VM opcode of the same name, and has the same semantics. This is implemented for use by compilers primarily.

assembly/src/assembler/basic_block_builder.rs

@@ -1,7 +1,10 @@
-use super::{AssemblyContext, BodyWrapper, Decorator, DecoratorList, Instruction};
+use super::{
+    mast_forest_builder::MastForestBuilder, AssemblyContext, BodyWrapper, Decorator, DecoratorList,
+    Instruction,
+};
 use alloc::{borrow::Borrow, string::ToString, vec::Vec};
 use vm_core::{
-    mast::{MastForest, MastNode, MastNodeId},
+    mast::{MastNode, MastNodeId},
     AdviceInjector, AssemblyOp, Operation,
 };
 
@@ -123,13 +126,16 @@ impl BasicBlockBuilder {
     ///
     /// This consumes all operations and decorators in the builder, but does not touch the
     /// operations in the epilogue of the builder.
-    pub fn make_basic_block(&mut self, mast_forest: &mut MastForest) -> Option<MastNodeId> {
+    pub fn make_basic_block(
+        &mut self,
+        mast_forest_builder: &mut MastForestBuilder,
+    ) -> Option<MastNodeId> {
         if !self.ops.is_empty() {
             let ops = self.ops.drain(..).collect();
             let decorators = self.decorators.drain(..).collect();
 
             let basic_block_node = MastNode::new_basic_block_with_decorators(ops, decorators);
-            let basic_block_node_id = mast_forest.ensure_node(basic_block_node);
+            let basic_block_node_id = mast_forest_builder.ensure_node(basic_block_node);
 
             Some(basic_block_node_id)
         } else if !self.decorators.is_empty() {
@@ -149,8 +155,11 @@ impl BasicBlockBuilder {
     /// - 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> {
+    pub fn into_basic_block(
+        mut self,
+        mast_forest_builder: &mut MastForestBuilder,
+    ) -> Option<MastNodeId> {
         self.ops.append(&mut self.epilogue);
-        self.make_basic_block(mast_forest)
+        self.make_basic_block(mast_forest_builder)
     }
 }

assembly/src/assembler/instruction/mod.rs

@@ -1,13 +1,10 @@
 use super::{
-    ast::InvokeKind, Assembler, AssemblyContext, BasicBlockBuilder, Felt, Instruction, Operation,
-    ONE, ZERO,
+    ast::InvokeKind, mast_forest_builder::MastForestBuilder, Assembler, AssemblyContext,
+    BasicBlockBuilder, Felt, Instruction, Operation, ONE, ZERO,
 };
 use crate::{diagnostics::Report, utils::bound_into_included_u64, AssemblyError};
 use core::ops::RangeBounds;
-use vm_core::{
-    mast::{MastForest, MastNodeId},
-    Decorator,
-};
+use vm_core::{mast::MastNodeId, Decorator};
 
 mod adv_ops;
 mod crypto_ops;
@@ -27,7 +24,7 @@ impl Assembler {
         instruction: &Instruction,
         span_builder: &mut BasicBlockBuilder,
         ctx: &mut AssemblyContext,
-        mast_forest: &mut MastForest,
+        mast_forest_builder: &mut MastForestBuilder,
     ) -> Result<Option<MastNodeId>, AssemblyError> {
         // if the assembler is in debug mode, start tracking the instruction about to be executed;
         // this will allow us to map the instruction to the sequence of operations which were
@@ -36,7 +33,8 @@ impl Assembler {
             span_builder.track_instruction(instruction, ctx);
         }
 
-        let result = self.compile_instruction_impl(instruction, span_builder, ctx, mast_forest)?;
+        let result =
+            self.compile_instruction_impl(instruction, span_builder, ctx, mast_forest_builder)?;
 
         // compute and update the cycle count of the instruction which just finished executing
         if self.in_debug_mode() {
@@ -51,7 +49,7 @@ impl Assembler {
         instruction: &Instruction,
         span_builder: &mut BasicBlockBuilder,
         ctx: &mut AssemblyContext,
-        mast_forest: &mut MastForest,
+        mast_forest_builder: &mut MastForestBuilder,
     ) -> Result<Option<MastNodeId>, AssemblyError> {
         use Operation::*;
 
@@ -376,18 +374,18 @@ impl Assembler {
 
             // ----- exec/call instructions -------------------------------------------------------
             Instruction::Exec(ref callee) => {
-                return self.invoke(InvokeKind::Exec, callee, ctx, mast_forest)
+                return self.invoke(InvokeKind::Exec, callee, ctx, mast_forest_builder)
             }
             Instruction::Call(ref callee) => {
-                return self.invoke(InvokeKind::Call, callee, ctx, mast_forest)
+                return self.invoke(InvokeKind::Call, callee, ctx, mast_forest_builder)
             }
             Instruction::SysCall(ref callee) => {
-                return self.invoke(InvokeKind::SysCall, callee, ctx, mast_forest)
+                return self.invoke(InvokeKind::SysCall, callee, ctx, mast_forest_builder)
             }
-            Instruction::DynExec => return self.dynexec(mast_forest),
-            Instruction::DynCall => return self.dyncall(mast_forest),
+            Instruction::DynExec => return self.dynexec(mast_forest_builder),
+            Instruction::DynCall => return self.dyncall(mast_forest_builder),
             Instruction::ProcRef(ref callee) => {
-                self.procref(callee, ctx, span_builder, mast_forest)?
+                self.procref(callee, ctx, span_builder, mast_forest_builder.forest())?
             }
 
             // ----- debug decorators -------------------------------------------------------------

assembly/src/assembler/instruction/procedures.rs

@@ -1,5 +1,6 @@
 use super::{Assembler, AssemblyContext, BasicBlockBuilder, Operation};
 use crate::{
+    assembler::mast_forest_builder::MastForestBuilder,
     ast::{InvocationTarget, InvokeKind},
     AssemblyError, RpoDigest, SourceSpan, Span, Spanned,
 };
@@ -14,11 +15,11 @@ impl Assembler {
         kind: InvokeKind,
         callee: &InvocationTarget,
         context: &mut AssemblyContext,
-        mast_forest: &mut MastForest,
+        mast_forest_builder: &mut MastForestBuilder,
     ) -> Result<Option<MastNodeId>, AssemblyError> {
         let span = callee.span();
-        let digest = self.resolve_target(kind, callee, context, mast_forest)?;
-        self.invoke_mast_root(kind, span, digest, context, mast_forest)
+        let digest = self.resolve_target(kind, callee, context, mast_forest_builder.forest())?;
+        self.invoke_mast_root(kind, span, digest, context, mast_forest_builder)
     }
 
     fn invoke_mast_root(
@@ -27,7 +28,7 @@ impl Assembler {
         span: SourceSpan,
         mast_root: RpoDigest,
         context: &mut AssemblyContext,
-        mast_forest: &mut MastForest,
+        mast_forest_builder: &mut MastForestBuilder,
     ) -> Result<Option<MastNodeId>, AssemblyError> {
         // Get the procedure from the assembler
         let cache = &self.procedure_cache;
@@ -68,42 +69,59 @@ impl Assembler {
                             })
                         }
                     })?;
-                context.register_external_call(&proc, false, mast_forest)?;
+                context.register_external_call(&proc, false, mast_forest_builder.forest())?;
+            }
+            Some(proc) => {
+                context.register_external_call(&proc, false, mast_forest_builder.forest())?
             }
-            Some(proc) => context.register_external_call(&proc, false, mast_forest)?,
             None if matches!(kind, InvokeKind::SysCall) => {
                 return Err(AssemblyError::UnknownSysCallTarget {
                     span,
-                    source_file: current_source_file,
+                    source_file: current_source_file.clone(),
                     callee: mast_root,
                 });
             }
             None => context.register_phantom_call(Span::new(span, mast_root))?,
         }
 
         let mast_root_node_id = {
-            // Note that here we rely on the fact that we topologically sorted the procedures, such
-            // that when we assemble a procedure, all procedures that it calls will have been
-            // assembled, and hence be present in the `MastForest`. We currently assume that the
-            // `MastForest` contains all the procedures being called; "external procedures" only
-            // known by digest are not currently supported.
-            let callee_id = mast_forest
-                .get_node_id_by_digest(mast_root)
-                .unwrap_or_else(|| panic!("MAST root {} not in MAST forest", mast_root));
-
             match kind {
-                // For `exec`, we return the root of the procedure being exec'd, which has the
-                // effect of inlining it
-                InvokeKind::Exec => callee_id,
-                // For `call`, we just use the corresponding CALL block
+                InvokeKind::Exec => {
+                    // Note that here we rely on the fact that we topologically sorted the
+                    // procedures, such that when we assemble a procedure, all
+                    // procedures that it calls will have been assembled, and
+                    // hence be present in the `MastForest`.
+                    mast_forest_builder.find_procedure_root(mast_root).unwrap_or_else(|| {
+                        // If the MAST root called isn't known to us, make it an external
+                        // reference.
+                        let external_node = MastNode::new_external(mast_root);
+                        mast_forest_builder.ensure_node(external_node)
+                    })
+                }
                 InvokeKind::Call => {
-                    let node = MastNode::new_call(callee_id, mast_forest);
-                    mast_forest.ensure_node(node)
+                    let callee_id =
+                        mast_forest_builder.find_procedure_root(mast_root).unwrap_or_else(|| {
+                            // If the MAST root called isn't known to us, make it an external
+                            // reference.
+                            let external_node = MastNode::new_external(mast_root);
+                            mast_forest_builder.ensure_node(external_node)
+                        });
+
+                    let call_node = MastNode::new_call(callee_id, mast_forest_builder.forest());
+                    mast_forest_builder.ensure_node(call_node)
                 }
-                // For `syscall`, we just use the corresponding SYSCALL block
                 InvokeKind::SysCall => {
-                    let node = MastNode::new_syscall(callee_id, mast_forest);
-                    mast_forest.ensure_node(node)
+                    let callee_id =
+                        mast_forest_builder.find_procedure_root(mast_root).unwrap_or_else(|| {
+                            // If the MAST root called isn't known to us, make it an external
+                            // reference.
+                            let external_node = MastNode::new_external(mast_root);
+                            mast_forest_builder.ensure_node(external_node)
+                        });
+
+                    let syscall_node =
+                        MastNode::new_syscall(callee_id, mast_forest_builder.forest());
+                    mast_forest_builder.ensure_node(syscall_node)
                 }
             }
         };
@@ -114,23 +132,23 @@ impl Assembler {
     /// Creates a new DYN block for the dynamic code execution and return.
     pub(super) fn dynexec(
         &self,
-        mast_forest: &mut MastForest,
+        mast_forest_builder: &mut MastForestBuilder,
     ) -> Result<Option<MastNodeId>, AssemblyError> {
-        let dyn_node_id = mast_forest.ensure_node(MastNode::Dyn);
+        let dyn_node_id = mast_forest_builder.ensure_node(MastNode::Dyn);
 
         Ok(Some(dyn_node_id))
     }
 
     /// Creates a new CALL block whose target is DYN.
     pub(super) fn dyncall(
         &self,
-        mast_forest: &mut MastForest,
+        mast_forest_builder: &mut MastForestBuilder,
     ) -> Result<Option<MastNodeId>, AssemblyError> {
         let dyn_call_node_id = {
-            let dyn_node_id = mast_forest.ensure_node(MastNode::Dyn);
-            let dyn_call_node = MastNode::new_call(dyn_node_id, mast_forest);
+            let dyn_node_id = mast_forest_builder.ensure_node(MastNode::Dyn);
+            let dyn_call_node = MastNode::new_call(dyn_node_id, mast_forest_builder.forest());
 
-            mast_forest.ensure_node(dyn_call_node)
+            mast_forest_builder.ensure_node(dyn_call_node)
         };
 
         Ok(Some(dyn_call_node_id))

assembly/src/assembler/mast_forest_builder.rs

@@ -0,0 +1,74 @@
+use core::ops::Index;
+
+use alloc::collections::BTreeMap;
+use vm_core::{
+    crypto::hash::RpoDigest,
+    mast::{MastForest, MastNode, MastNodeId, MerkleTreeNode},
+};
+
+/// Builder for a [`MastForest`].
+#[derive(Clone, Debug, Default)]
+pub struct MastForestBuilder {
+    mast_forest: MastForest,
+    node_id_by_hash: BTreeMap<RpoDigest, MastNodeId>,
+}
+
+impl MastForestBuilder {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn build(self) -> MastForest {
+        self.mast_forest
+    }
+}
+
+/// Accessors
+impl MastForestBuilder {
+    /// Returns the underlying [`MastForest`] being built
+    pub fn forest(&self) -> &MastForest {
+        &self.mast_forest
+    }
+
+    /// Returns the [`MastNodeId`] of the procedure associated with a given digest, if any.
+    #[inline(always)]
+    pub fn find_procedure_root(&self, digest: RpoDigest) -> Option<MastNodeId> {
+        self.mast_forest.find_procedure_root(digest)
+    }
+}
+
+/// Mutators
+impl MastForestBuilder {
+    /// Adds a node to the forest, and returns the [`MastNodeId`] associated with it.
+    ///
+    /// If a [`MastNode`] which is equal to the current node was previously added, the previously
+    /// returned [`MastNodeId`] will be returned. This enforces this invariant that equal
+    /// [`MastNode`]s have equal [`MastNodeId`]s.
+    pub fn ensure_node(&mut self, node: MastNode) -> MastNodeId {
+        let node_digest = node.digest();
+
+        if let Some(node_id) = self.node_id_by_hash.get(&node_digest) {
+            // node already exists in the forest; return previously assigned id
+            *node_id
+        } else {
+            let new_node_id = self.mast_forest.add_node(node);
+            self.node_id_by_hash.insert(node_digest, new_node_id);
+
+            new_node_id
+        }
+    }
+
+    /// Marks the given [`MastNodeId`] as being the root of a procedure.
+    pub fn make_root(&mut self, new_root_id: MastNodeId) {
+        self.mast_forest.make_root(new_root_id)
+    }
+}
+
+impl Index<MastNodeId> for MastForestBuilder {
+    type Output = MastNode;
+
+    #[inline(always)]
+    fn index(&self, node_id: MastNodeId) -> &Self::Output {
+        &self.mast_forest[node_id]
+    }
+}