fix NameResolver

0xPolygonMiden · Jul 18, 2024 · 9bb9ed5 · 9bb9ed5
1 parent 1b5a1b5
commit 9bb9ed5
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/assembly/src/assembler/module_graph/mod.rs b/assembly/src/assembler/module_graph/mod.rs
@@ -101,9 +101,9 @@ impl WrapperModule {
         }
     }
 
-    pub fn unwrap_ast(&self) -> Arc<Module> {
+    pub fn unwrap_ast(&self) -> &Arc<Module> {
         match self {
-            WrapperModule::Ast(module) => module.clone(),
+            WrapperModule::Ast(module) => module,
             WrapperModule::Exports(_) => {
                 panic!("expected module to be in AST representation, but was compiled")
             }
@@ -558,7 +558,7 @@ impl ModuleGraph {
     /// Fetch a [WrapperProcedure] by [GlobalProcedureIndex], or `None` if index is invalid.
     pub fn get_procedure(&self, id: GlobalProcedureIndex) -> Option<WrapperProcedure> {
         match &self.modules[id.module.as_usize()] {
-            WrapperModule::Ast(m) => m.get(id.index).map(|export| WrapperProcedure::Ast(export)),
+            WrapperModule::Ast(m) => m.get(id.index).map(WrapperProcedure::Ast),
             WrapperModule::Exports(m) => m
                 .procedures
                 .get(id.index.as_usize())

diff --git a/assembly/src/assembler/module_graph/name_resolver.rs b/assembly/src/assembler/module_graph/name_resolver.rs
@@ -177,7 +177,7 @@ impl<'a> NameResolver<'a> {
                                 .get_name(gid.index)
                                 .clone()
                         } else {
-                            self.graph[gid].name().clone()
+                            self.graph.get_procedure_unsafe(gid).name().clone()
                         };
                         Ok(ResolvedTarget::Resolved {
                             gid,
@@ -210,6 +210,7 @@ impl<'a> NameResolver<'a> {
                 .resolve_import(name)
         } else {
             self.graph[caller.module]
+                .unwrap_ast()
                 .resolve_import(name)
                 .map(|import| Span::new(import.span(), import.path()))
         }
@@ -255,7 +256,7 @@ impl<'a> NameResolver<'a> {
                                     .get_name(gid.index)
                                     .clone()
                             } else {
-                                self.graph[gid].name().clone()
+                                self.graph.get_procedure_unsafe(gid).name().clone()
                             };
                             Ok(ResolvedTarget::Resolved {
                                 gid,
@@ -315,7 +316,7 @@ impl<'a> NameResolver<'a> {
         if module_index >= pending_offset {
             self.pending[module_index - pending_offset].resolver.resolve(callee)
         } else {
-            self.graph[module].resolve(callee)
+            self.graph[module].unwrap_ast().resolve(callee)
         }
     }
 
@@ -487,7 +488,7 @@ impl<'a> NameResolver<'a> {
         if module_index >= pending_offset {
             self.pending[module_index - pending_offset].source_file.clone()
         } else {
-            self.graph[module].source_file()
+            self.graph[module].unwrap_ast().source_file()
         }
     }