ProcedureInfo: store body_node_id

assembly/src/assembler/mod.rs

@@ -455,7 +455,11 @@ impl Assembler {
         while let Some(procedure_gid) = worklist.pop() {
             // If we have already compiled this procedure, do not recompile
             if let Some(proc) = mast_forest_builder.get_procedure(procedure_gid) {
-                self.module_graph.register_procedure_root(procedure_gid, proc.mast_root())?;
+                self.module_graph.register_procedure_root(
+                    procedure_gid,
+                    proc.body_node_id(),
+                    proc.mast_root(),
+                )?;
                 continue;
             }
             // Fetch procedure metadata from the graph
@@ -494,7 +498,11 @@ impl Assembler {
                     // be added to the forest.
 
                     // Cache the compiled procedure
-                    self.module_graph.register_procedure_root(procedure_gid, procedure.mast_root())?;
+                    self.module_graph.register_procedure_root(
+                        procedure_gid,
+                        procedure.body_node_id(),
+                        procedure.mast_root(),
+                    )?;
                     mast_forest_builder.insert_procedure(procedure_gid, procedure)?;
                 },
                 Export::Alias(proc_alias) => {
@@ -526,7 +534,11 @@ impl Assembler {
                     let procedure = pctx.into_procedure(proc_mast_root, proc_node_id);
 
                     // Make the MAST root available to all dependents
-                    self.module_graph.register_procedure_root(procedure_gid, proc_mast_root)?;
+                    self.module_graph.register_procedure_root(
+                        procedure_gid,
+                        proc_node_id,
+                        proc_mast_root,
+                    )?;
                     mast_forest_builder.insert_procedure(procedure_gid, procedure)?;
                 },
             }

assembly/src/assembler/module_graph/mod.rs

@@ -4,7 +4,12 @@ mod debug;
 mod name_resolver;
 mod rewrites;
 
-use alloc::{boxed::Box, collections::BTreeMap, sync::Arc, vec::Vec};
+use alloc::{
+    boxed::Box,
+    collections::{btree_map::Entry, BTreeMap},
+    sync::Arc,
+    vec::Vec,
+};
 use core::ops::Index;
 
 use smallvec::{smallvec, SmallVec};
@@ -161,7 +166,7 @@ pub struct ModuleGraph {
     /// The set of MAST node ids which have procedure definitions in this graph. There can be
     /// multiple procedures bound to the same root due to having identical code.
     procedure_root_digests: BTreeMap<RpoDigest, SmallVec<[GlobalProcedureIndex; 1]>>,
-    /// The set of MAST node ids which have procedure definitions in this graph. 
+    /// The set of MAST node ids which have procedure definitions in this graph.
     procedure_root_ids: BTreeMap<MastNodeId, GlobalProcedureIndex>,
     kernel_index: Option<ModuleIndex>,
     kernel: Kernel,
@@ -201,7 +206,7 @@ impl ModuleGraph {
         for &module_index in module_indices.iter() {
             for (proc_index, proc) in self[module_index].unwrap_info().clone().procedures() {
                 let gid = module_index + proc_index;
-                self.register_procedure_root(gid, proc.digest)?;
+                self.register_procedure_root(gid, proc.body_node_id, proc.digest)?;
             }
         }
 
@@ -531,7 +536,7 @@ impl ModuleGraph {
     }
 
     /// Returns a procedure index which corresponds to the provided procedure digest.
-    /// 
+    ///
     /// Note that there can be many procedures with the same digest - due to having the same code,
     /// and/or using different decorators which don't affect the MAST root. This method returns an
     /// arbitrary one.
@@ -554,6 +559,7 @@ impl ModuleGraph {
         resolver.resolve_target(caller, target)
     }
 
+    // TODO(plafer): fix docs
     /// Registers a [MastNodeId] as corresponding to a given [GlobalProcedureIndex].
     ///
     /// # SAFETY
@@ -564,21 +570,27 @@ impl ModuleGraph {
     /// of the referenced procedure, i.e. they are referentially transparent.
     pub(crate) fn register_procedure_root(
         &mut self,
-        id: GlobalProcedureIndex,
-        // TODO(plafer): all `procedure_root_id` field, and add to `self.procedure_root_ids`
+        gid: GlobalProcedureIndex,
+        procedure_root_id: MastNodeId,
         procedure_root_digest: RpoDigest,
     ) -> Result<(), AssemblyError> {
-        use alloc::collections::btree_map::Entry;
+        if let Some(old_gid) = self.procedure_root_ids.insert(procedure_root_id, gid) {
+            assert!(
+                old_gid == gid, 
+                "internal error: procedure root id {procedure_root_id:?} is associated with two procedure indices: {old_gid:?} and {gid:?}"
+            );
+        }
+
         match self.procedure_root_digests.entry(procedure_root_digest) {
             Entry::Occupied(ref mut entry) => {
                 let prev_id = entry.get()[0];
-                if prev_id != id {
+                if prev_id != gid {
                     // Multiple procedures with the same root, but compatible
-                    entry.get_mut().push(id);
+                    entry.get_mut().push(gid);
                 }
             },
             Entry::Vacant(entry) => {
-                entry.insert(smallvec![id]);
+                entry.insert(smallvec![gid]);
             },
         }
 

assembly/src/errors.rs

@@ -28,6 +28,7 @@ pub enum AssemblyError {
     #[error("found a cycle in the call graph, involving these procedures: {}", nodes.as_slice().join(", "))]
     #[diagnostic()]
     Cycle { nodes: Vec<String> },
+    // TODO(plafer): remove this error?
     #[error("two procedures found with same mast root, but conflicting definitions ('{first}' and '{second}')")]
     #[diagnostic()]
     ConflictingDefinitions {

assembly/src/library/mod.rs

@@ -73,11 +73,7 @@ impl Library {
     ) -> Result<Self, LibraryError> {
         let digest = compute_content_hash(&exports, &mast_forest);
 
-        Ok(Self {
-            digest,
-            exports,
-            mast_forest,
-        })
+        Ok(Self { digest, exports, mast_forest })
     }
 }
 
@@ -132,13 +128,21 @@ impl Library {
                 .entry(proc_name.module.clone())
                 .and_modify(|compiled_module| {
                     let proc_digest = self.mast_forest[proc_root_node_id].digest();
-                    compiled_module.add_procedure(proc_name.name.clone(), proc_digest);
+                    compiled_module.add_procedure(
+                        proc_name.name.clone(),
+                        proc_root_node_id,
+                        proc_digest,
+                    );
                 })
                 .or_insert_with(|| {
                     let mut module_info = ModuleInfo::new(proc_name.module.clone());
 
                     let proc_digest = self.mast_forest[proc_root_node_id].digest();
-                    module_info.add_procedure(proc_name.name.clone(), proc_digest);
+                    module_info.add_procedure(
+                        proc_name.name.clone(),
+                        proc_root_node_id,
+                        proc_digest,
+                    );
 
                     module_info
                 });
@@ -357,7 +361,7 @@ impl TryFrom<Library> for KernelLibrary {
 
             let proc_digest = library.mast_forest[proc_node_id].digest();
             proc_digests.push(proc_digest);
-            kernel_module.add_procedure(proc_path.name.clone(), proc_digest);
+            kernel_module.add_procedure(proc_path.name.clone(), proc_node_id, proc_digest);
         }
 
         let kernel = Kernel::new(&proc_digests)?;

assembly/src/library/module.rs

@@ -1,4 +1,5 @@
 use alloc::vec::Vec;
+use vm_core::mast::MastNodeId;
 
 use super::LibraryPath;
 use crate::{
@@ -22,8 +23,8 @@ impl ModuleInfo {
     }
 
     /// Adds a procedure to the module.
-    pub fn add_procedure(&mut self, name: ProcedureName, digest: RpoDigest) {
-        self.procedures.push(ProcedureInfo { name, digest });
+    pub fn add_procedure(&mut self, name: ProcedureName, body_node_id: MastNodeId, digest: RpoDigest) {
+        self.procedures.push(ProcedureInfo { name, body_node_id, digest });
     }
 
     /// Returns the module's library path.
@@ -71,5 +72,6 @@ impl ModuleInfo {
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct ProcedureInfo {
     pub name: ProcedureName,
+    pub body_node_id: MastNodeId,
     pub digest: RpoDigest,
 }