refactor: remove root tracking in the ModuleGraph

assembly/src/assembler/mod.rs

@@ -356,8 +356,7 @@ impl Assembler {
         let mut compiled_entrypoint = None;
         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_mast_root(procedure_gid, proc.mast_root())?;
+            if mast_forest_builder.get_procedure(procedure_gid).is_some() {
                 continue;
             }
             let is_entry = entrypoint == Some(procedure_gid);
@@ -384,8 +383,6 @@ impl Assembler {
                 mast_forest_builder.make_root(procedure.body_node_id());
                 compiled_entrypoint = Some(Arc::from(procedure));
             } else {
-                // Make the MAST root available to all dependents
-                self.module_graph.register_mast_root(procedure_gid, procedure.mast_root())?;
                 mast_forest_builder.insert_procedure(procedure_gid, procedure)?;
             }
         }
@@ -536,7 +533,7 @@ impl Assembler {
         };
         let resolved = self.module_graph.resolve_target(&caller, target)?;
         match resolved {
-            ResolvedTarget::Phantom(digest) => Ok(digest),
+            ResolvedTarget::MastRoot(digest) => Ok(digest),
             ResolvedTarget::Exact { gid } | ResolvedTarget::Resolved { gid, .. } => {
                 Ok(mast_forest_builder
                     .get_procedure(gid)

assembly/src/assembler/module_graph/analysis/rewrite_check.rs

@@ -65,7 +65,7 @@ impl<'a, 'b: 'a> RewriteCheckVisitor<'a, 'b> {
         match self.resolver.resolve_target(&caller, target) {
             Err(err) => ControlFlow::Break(Err(err)),
             Ok(ResolvedTarget::Resolved { .. }) => ControlFlow::Break(Ok(true)),
-            Ok(ResolvedTarget::Exact { .. } | ResolvedTarget::Phantom(_)) => {
+            Ok(ResolvedTarget::Exact { .. } | ResolvedTarget::MastRoot(_)) => {
                 ControlFlow::Continue(())
             }
         }

assembly/src/assembler/module_graph/mod.rs

@@ -7,12 +7,10 @@ mod rewrites;
 pub use self::callgraph::{CallGraph, CycleError};
 pub use self::name_resolver::{CallerInfo, ResolvedTarget};
 
-use alloc::{boxed::Box, collections::BTreeMap, sync::Arc, vec::Vec};
+use alloc::{boxed::Box, sync::Arc, vec::Vec};
 use core::ops::Index;
 use vm_core::Kernel;
 
-use smallvec::{smallvec, SmallVec};
-
 use self::{analysis::MaybeRewriteCheck, name_resolver::NameResolver, rewrites::ModuleRewriter};
 use super::{GlobalProcedureIndex, ModuleIndex};
 use crate::{
@@ -44,9 +42,6 @@ pub struct ModuleGraph {
     callgraph: CallGraph,
     /// The computed topological ordering of the call graph
     topo: Vec<GlobalProcedureIndex>,
-    /// The set of MAST roots which have procedure definitions in this graph. There can be
-    /// multiple procedures bound to the same root due to having identical code.
-    roots: BTreeMap<RpoDigest, SmallVec<[GlobalProcedureIndex; 1]>>,
     kernel_index: Option<ModuleIndex>,
     kernel: Kernel,
 }
@@ -342,13 +337,6 @@ impl ModuleGraph {
         self.modules.get(id.module.as_usize()).and_then(|m| m.get(id.index))
     }
 
-    pub fn get_procedure_index_by_digest(
-        &self,
-        digest: &RpoDigest,
-    ) -> Option<GlobalProcedureIndex> {
-        self.roots.get(digest).map(|indices| indices[0])
-    }
-
     #[allow(unused)]
     pub fn callees(&self, gid: GlobalProcedureIndex) -> &[GlobalProcedureIndex] {
         self.callgraph.out_edges(gid)
@@ -364,58 +352,6 @@ impl ModuleGraph {
         resolver.resolve_target(caller, target)
     }
 
-    /// Registers a [RpoDigest] as corresponding to a given [GlobalProcedureIndex].
-    ///
-    /// # SAFETY
-    ///
-    /// It is essential that the caller _guarantee_ that the given digest belongs to the specified
-    /// procedure. It is fine if there are multiple procedures with the same digest, but it _must_
-    /// be the case that if a given digest is specified, it can be used as if it was the definition
-    /// of the referenced procedure, i.e. they are referentially transparent.
-    pub(crate) fn register_mast_root(
-        &mut self,
-        id: GlobalProcedureIndex,
-        digest: RpoDigest,
-    ) -> Result<(), AssemblyError> {
-        use alloc::collections::btree_map::Entry;
-        match self.roots.entry(digest) {
-            Entry::Occupied(ref mut entry) => {
-                let prev_id = entry.get()[0];
-                if prev_id != id {
-                    // Multiple procedures with the same root, but incompatible
-                    let prev = &self.modules[prev_id.module.as_usize()][prev_id.index];
-                    let current = &self.modules[id.module.as_usize()][id.index];
-                    if prev.num_locals() != current.num_locals() {
-                        let prev_module = self.modules[prev_id.module.as_usize()].path();
-                        let prev_name = FullyQualifiedProcedureName {
-                            span: prev.span(),
-                            module: prev_module.clone(),
-                            name: prev.name().clone(),
-                        };
-                        let current_module = self.modules[id.module.as_usize()].path();
-                        let current_name = FullyQualifiedProcedureName {
-                            span: current.span(),
-                            module: current_module.clone(),
-                            name: current.name().clone(),
-                        };
-                        return Err(AssemblyError::ConflictingDefinitions {
-                            first: prev_name,
-                            second: current_name,
-                        });
-                    }
-
-                    // Multiple procedures with the same root, but compatible
-                    entry.get_mut().push(id);
-                }
-            }
-            Entry::Vacant(entry) => {
-                entry.insert(smallvec![id]);
-            }
-        }
-
-        Ok(())
-    }
-
     /// Resolve a [LibraryPath] to a [ModuleIndex] in this graph
     pub fn find_module_index(&self, name: &LibraryPath) -> Option<ModuleIndex> {
         self.modules.iter().position(|m| m.path() == name).map(ModuleIndex::new)

assembly/src/assembler/module_graph/name_resolver.rs

@@ -58,16 +58,16 @@ pub enum ResolvedTarget {
         /// The [InvocationTarget] to use in the caller
         target: InvocationTarget,
     },
-    /// We know the MAST root of the callee, but the procedure is not available, either in the
-    /// procedure cache or in the [ModuleGraph].
-    Phantom(RpoDigest),
+    /// We know the MAST root of the callee, the procedure may or may not be available in the
+    /// [ModuleGraph].
+    MastRoot(RpoDigest),
 }
 
 impl ResolvedTarget {
     pub fn into_global_id(self) -> Option<GlobalProcedureIndex> {
         match self {
             ResolvedTarget::Exact { gid } | ResolvedTarget::Resolved { gid, .. } => Some(gid),
-            ResolvedTarget::Phantom(_) => None,
+            ResolvedTarget::MastRoot(_) => None,
         }
     }
 }
@@ -159,12 +159,7 @@ impl<'a> NameResolver<'a> {
                     target: InvocationTarget::AbsoluteProcedurePath { name, path },
                 })
             }
-            Some(ResolvedProcedure::MastRoot(ref digest)) => {
-                match self.graph.get_procedure_index_by_digest(digest) {
-                    Some(gid) => Ok(ResolvedTarget::Exact { gid }),
-                    None => Ok(ResolvedTarget::Phantom(**digest)),
-                }
-            }
+            Some(ResolvedProcedure::MastRoot(ref digest)) => Ok(ResolvedTarget::MastRoot(**digest)),
             None => Err(AssemblyError::Failed {
                 labels: vec![RelatedLabel::error("undefined procedure")
                     .with_source_file(caller.source_file.clone())
@@ -196,12 +191,7 @@ impl<'a> NameResolver<'a> {
         target: &InvocationTarget,
     ) -> Result<ResolvedTarget, AssemblyError> {
         match target {
-            InvocationTarget::MastRoot(mast_root) => {
-                match self.graph.get_procedure_index_by_digest(mast_root) {
-                    None => Ok(ResolvedTarget::Phantom(mast_root.into_inner())),
-                    Some(gid) => Ok(ResolvedTarget::Exact { gid }),
-                }
-            }
+            InvocationTarget::MastRoot(mast_root) => Ok(ResolvedTarget::MastRoot(**mast_root)),
             InvocationTarget::ProcedureName(ref callee) => self.resolve(caller, callee),
             InvocationTarget::ProcedurePath {
                 ref name,
@@ -350,10 +340,7 @@ impl<'a> NameResolver<'a> {
                     });
                     current_callee = Cow::Owned(fqn);
                 }
-                Some(ResolvedProcedure::MastRoot(ref digest)) => {
-                    if let Some(id) = self.graph.get_procedure_index_by_digest(digest) {
-                        break Ok(id);
-                    }
+                Some(ResolvedProcedure::MastRoot(_)) => {
                     // This is a phantom procedure - we know its root, but do not have its
                     // definition
                     break Err(AssemblyError::Failed {

assembly/src/assembler/module_graph/rewrites/module.rs

@@ -67,7 +67,7 @@ impl<'a, 'b: 'a> ModuleRewriter<'a, 'b> {
         };
         match self.resolver.resolve_target(&caller, target) {
             Err(err) => return ControlFlow::Break(err),
-            Ok(ResolvedTarget::Phantom(_)) => (),
+            Ok(ResolvedTarget::MastRoot(_)) => (),
             Ok(ResolvedTarget::Exact { .. }) => {
                 self.invoked.insert(Invoke {
                     kind,