Remove uniqueness invariant from MastForestBuilder

assembly/src/assembler/basic_block_builder.rs

@@ -139,7 +139,7 @@ impl BasicBlockBuilder {
             let ops = self.ops.drain(..).collect();
             let decorators = self.decorators.drain(..).collect();
 
-            let basic_block_node_id = mast_forest_builder.ensure_block(ops, Some(decorators))?;
+            let basic_block_node_id = mast_forest_builder.add_block(ops, Some(decorators))?;
 
             Ok(Some(basic_block_node_id))
         } else if !self.decorators.is_empty() {

assembly/src/assembler/instruction/mod.rs

@@ -433,13 +433,19 @@ impl Assembler {
 
             // ----- exec/call instructions -------------------------------------------------------
             Instruction::Exec(ref callee) => {
-                return self.invoke(InvokeKind::Exec, callee, proc_ctx, mast_forest_builder)
+                return self
+                    .invoke(InvokeKind::Exec, callee, proc_ctx, mast_forest_builder)
+                    .map(Into::into);
             },
             Instruction::Call(ref callee) => {
-                return self.invoke(InvokeKind::Call, callee, proc_ctx, mast_forest_builder)
+                return self
+                    .invoke(InvokeKind::Call, callee, proc_ctx, mast_forest_builder)
+                    .map(Into::into);
             },
             Instruction::SysCall(ref callee) => {
-                return self.invoke(InvokeKind::SysCall, callee, proc_ctx, mast_forest_builder)
+                return self
+                    .invoke(InvokeKind::SysCall, callee, proc_ctx, mast_forest_builder)
+                    .map(Into::into);
             },
             Instruction::DynExec => return self.dynexec(mast_forest_builder),
             Instruction::DynCall => return self.dyncall(mast_forest_builder),

assembly/src/assembler/instruction/procedures.rs

@@ -1,5 +1,5 @@
 use smallvec::SmallVec;
-use vm_core::mast::MastNodeId;
+use vm_core::{mast::MastNodeId, utils::Either};
 
 use super::{Assembler, BasicBlockBuilder, Operation};
 use crate::{
@@ -16,10 +16,27 @@ impl Assembler {
         callee: &InvocationTarget,
         proc_ctx: &mut ProcedureContext,
         mast_forest_builder: &mut MastForestBuilder,
-    ) -> Result<Option<MastNodeId>, AssemblyError> {
+    ) -> Result<MastNodeId, AssemblyError> {
         let span = callee.span();
-        let digest = self.resolve_target(kind, callee, proc_ctx, mast_forest_builder)?;
-        self.invoke_mast_root(kind, span, digest, mast_forest_builder)
+        let node_id_or_digest = self.resolve_target(kind, callee, proc_ctx, mast_forest_builder)?;
+        match node_id_or_digest {
+            // TODO(plafer): reconcile with `invoke_mast_root` impl
+            Either::Left(node_id) => match kind {
+                InvokeKind::ProcRef => Ok(node_id),
+                InvokeKind::Exec => {
+                    // We make sure to copy the root node so that the `exec` is associated
+                    // with a different `MastNodeId` than the procedure it is referencing.
+                    // Currently the only purpose of this is so that simple procedures that
+                    // only have an `exec` have a different body node id than the procedure
+                    // they're executing.
+                    let root_node = mast_forest_builder.get_mast_node(node_id).unwrap();
+                    Ok(mast_forest_builder.add_node(root_node.clone())?)
+                },
+                InvokeKind::Call => Ok(mast_forest_builder.add_call(node_id)?),
+                InvokeKind::SysCall => Ok(mast_forest_builder.add_syscall(node_id)?),
+            },
+            Either::Right(digest) => self.invoke_mast_root(kind, span, digest, mast_forest_builder),
+        }
     }
 
     fn invoke_mast_root(
@@ -28,12 +45,12 @@ impl Assembler {
         span: SourceSpan,
         mast_root: RpoDigest,
         mast_forest_builder: &mut MastForestBuilder,
-    ) -> Result<Option<MastNodeId>, AssemblyError> {
+    ) -> Result<MastNodeId, AssemblyError> {
         // Get the procedure from the assembler
         let current_source_file = self.source_manager.get(span.source_id()).ok();
 
         // If the procedure is cached and is a system call, ensure that the call is valid.
-        match mast_forest_builder.find_procedure(&mast_root) {
+        match mast_forest_builder.find_procedure_by_digest(&mast_root) {
             Some(proc) if matches!(kind, InvokeKind::SysCall) => {
                 // Verify if this is a syscall, that the callee is a kernel procedure
                 //
@@ -57,8 +74,8 @@ impl Assembler {
                     })
                     .and_then(|module| {
                         // Note: this module is guaranteed to be of AST variant, since we have the
-                        // AST of a procedure contained in it (i.e. `proc`). Hence, it must be that
-                        // the entire module is in AST representation as well.
+                        // AST of a procedure contained in it (i.e. `proc`). Hence, it must be
+                        // thatthe entire module is in AST representation as well.
                         if module.unwrap_ast().is_kernel() {
                             Ok(())
                         } else {
@@ -74,18 +91,36 @@ impl Assembler {
         }
 
         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`.
             match kind {
-                InvokeKind::Exec | InvokeKind::ProcRef => {
-                    // 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`.
+                InvokeKind::ProcRef => {
                     match mast_forest_builder.find_procedure_node_id(mast_root) {
                         Some(root) => root,
                         None => {
                             // If the MAST root called isn't known to us, make it an external
                             // reference.
-                            mast_forest_builder.ensure_external(mast_root)?
+                            mast_forest_builder.add_external(mast_root)?
+                        },
+                    }
+                },
+                InvokeKind::Exec => {
+                    match mast_forest_builder.find_procedure_node_id(mast_root) {
+                        Some(root) => {
+                            // We make sure to copy the root node so that the `exec` is associated
+                            // with a different `MastNodeId` than the procedure it is referencing.
+                            // Currently the only purpose of this is so that simple procedures that
+                            // only have an `exec` have a different body node id than the procedure
+                            // they're executing.
+                            let root_node = mast_forest_builder.get_mast_node(root).unwrap();
+                            mast_forest_builder.add_node(root_node.clone())?
+                        },
+                        None => {
+                            // If the MAST root called isn't known to us, make it an external
+                            // reference.
+                            mast_forest_builder.add_external(mast_root)?
                         },
                     }
                 },
@@ -95,36 +130,36 @@ impl Assembler {
                         None => {
                             // If the MAST root called isn't known to us, make it an external
                             // reference.
-                            mast_forest_builder.ensure_external(mast_root)?
+                            mast_forest_builder.add_external(mast_root)?
                         },
                     };
 
-                    mast_forest_builder.ensure_call(callee_id)?
+                    mast_forest_builder.add_call(callee_id)?
                 },
                 InvokeKind::SysCall => {
                     let callee_id = match mast_forest_builder.find_procedure_node_id(mast_root) {
                         Some(callee_id) => callee_id,
                         None => {
                             // If the MAST root called isn't known to us, make it an external
                             // reference.
-                            mast_forest_builder.ensure_external(mast_root)?
+                            mast_forest_builder.add_external(mast_root)?
                         },
                     };
 
-                    mast_forest_builder.ensure_syscall(callee_id)?
+                    mast_forest_builder.add_syscall(callee_id)?
                 },
             }
         };
 
-        Ok(Some(mast_root_node_id))
+        Ok(mast_root_node_id)
     }
 
     /// Creates a new DYN block for the dynamic code execution and return.
     pub(super) fn dynexec(
         &self,
         mast_forest_builder: &mut MastForestBuilder,
     ) -> Result<Option<MastNodeId>, AssemblyError> {
-        let dyn_node_id = mast_forest_builder.ensure_dyn()?;
+        let dyn_node_id = mast_forest_builder.add_dyn()?;
 
         Ok(Some(dyn_node_id))
     }
@@ -135,8 +170,8 @@ impl Assembler {
         mast_forest_builder: &mut MastForestBuilder,
     ) -> Result<Option<MastNodeId>, AssemblyError> {
         let dyn_call_node_id = {
-            let dyn_node_id = mast_forest_builder.ensure_dyn()?;
-            mast_forest_builder.ensure_call(dyn_node_id)?
+            let dyn_node_id = mast_forest_builder.add_dyn()?;
+            mast_forest_builder.add_call(dyn_node_id)?
         };
 
         Ok(Some(dyn_call_node_id))
@@ -149,8 +184,15 @@ impl Assembler {
         basic_block_builder: &mut BasicBlockBuilder,
         mast_forest_builder: &MastForestBuilder,
     ) -> Result<(), AssemblyError> {
-        let digest =
-            self.resolve_target(InvokeKind::ProcRef, callee, proc_ctx, mast_forest_builder)?;
+        let digest = match self.resolve_target(
+            InvokeKind::ProcRef,
+            callee,
+            proc_ctx,
+            mast_forest_builder,
+        )? {
+            Either::Left(node_id) => mast_forest_builder.get_mast_node(node_id).unwrap().digest(),
+            Either::Right(digest) => digest,
+        };
         self.procref_mast_root(digest, basic_block_builder)
     }