Merge pull request #1023 from 0xPolygonMiden/frisitano-dup-mast-fix

Fix: Duplicate MAST root different procedure Id
0xPolygonMiden · Jul 24, 2023 · 579fcf4 · 579fcf4
2 parents 7588767 + f15e84e
commit 579fcf4
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 1 deletion.
diff --git a/assembly/src/assembler/procedure_cache.rs b/assembly/src/assembler/procedure_cache.rs
@@ -72,11 +72,12 @@ impl ProcedureCache {
                 if proc.num_locals() != cached_proc.num_locals() {
                     Err(AssemblyError::conflicting_num_locals(proc.name()))
                 } else {
+                    self.proc_id_map.insert(*proc.id(), proc.mast_root());
                     Ok(())
                 }
             }
             Entry::Vacant(entry) => {
-                self.proc_id_map.entry(*proc.id()).or_insert(proc.mast_root());
+                self.proc_id_map.insert(*proc.id(), proc.mast_root());
                 entry.insert(proc.into_inner());
                 Ok(())
             }

diff --git a/assembly/src/tests.rs b/assembly/src/tests.rs
@@ -663,6 +663,48 @@ fn program_with_one_import_and_hex_call() {
     assert_eq!(expected, format!("{program}"));
 }
 
+#[test]
+fn program_with_two_imported_procs_with_same_mast_root() {
+    const NAMESPACE: &str = "dummy";
+    const MODULE: &str = "math::u256";
+    const PROCEDURE: &str = r#"
+        export.iszero_unsafe_dup
+            eq.0
+            repeat.7
+                swap
+                eq.0
+                and
+            end
+        end
+
+        export.iszero_unsafe
+            eq.0
+            repeat.7
+                swap
+                eq.0
+                and
+            end
+        end"#;
+
+    let namespace = LibraryNamespace::try_from(NAMESPACE.to_string()).unwrap();
+    let path = LibraryPath::try_from(MODULE.to_string()).unwrap().prepend(&namespace).unwrap();
+    let ast = ModuleAst::parse(PROCEDURE).unwrap();
+    let modules = vec![Module { path, ast }];
+    let library = DummyLibrary::new(namespace, modules);
+
+    let assembler = super::Assembler::default().with_library(&library).unwrap();
+    let source = format!(
+        r#"
+        use.{NAMESPACE}::{MODULE}
+        begin
+            push.4 push.3
+            exec.u256::iszero_unsafe
+            exec.u256::iszero_unsafe_dup
+        end"#
+    );
+    assert!(assembler.compile(source).is_ok());
+}
+
 #[test]
 fn program_with_reexported_proc_in_same_library() {
     // exprted proc is in same library