chore: Address review comments

assembly/src/assembler/mast_forest_merger_tests.rs

@@ -31,11 +31,19 @@ fn mast_forest_merge_assembler() {
   export.foo
       push.19
   end
+
+  export.qux
+      swap drop
+  end
 "#;
 
     let lib_b = r#"
   use.lib::mod
 
+  export.qux_duplicate
+      swap drop
+  end
+
   export.bar
       push.2
       if.true
@@ -55,7 +63,7 @@ fn mast_forest_merge_assembler() {
         for root in forest.procedure_roots() {
             let original_digest = forest.nodes()[root.as_usize()].digest();
             let new_root = root_maps.map_root(forest_idx, root).unwrap();
-            let new_digest = forest.nodes()[new_root.as_usize()].digest();
+            let new_digest = merged.nodes()[new_root.as_usize()].digest();
             assert_eq!(original_digest, new_digest);
         }
     }

core/src/mast/merger/mod.rs

@@ -357,9 +357,8 @@ impl MastForestRootMap {
 ///
 /// This type is meant to encapsulates some guarantees:
 ///
-/// - Indexing into the vector for any ID is safe if that ID is valid for the corresponding forest,
-///   which is enforced in the `from_u32_safe` functions (as long as they are used with the correct
-///   forest). Despite that, we still cannot index unconditionally in case a node with invalid
+/// - Indexing into the vector for any ID is safe if that ID is valid for the corresponding forest.
+///   Despite that, we still cannot index unconditionally in case a node with invalid
 ///   [`DecoratorId`]s is passed to `merge`.
 /// - The entry itself can be either None or Some. However:
 ///   - For `DecoratorId`s we iterate and insert all decorators into this map before retrieving any

core/src/mast/mod.rs

@@ -532,8 +532,8 @@ impl MastNodeId {
         }
     }
 
-    #[cfg(test)]
-    pub fn new_unsafe(value: u32) -> Self {
+    /// Returns a new [`MastNodeId`] from the given `value` without checking its validity.
+    pub(crate) fn new_unsafe(value: u32) -> Self {
         Self(value)
     }
 

core/src/mast/multi_forest_node_iterator.rs

@@ -102,8 +102,9 @@ impl<'forest> MultiMastForestNodeIter<'forest> {
 
         for (forest_idx, forest) in mast_forests.iter().enumerate() {
             for (node_idx, node) in forest.nodes().iter().enumerate() {
-                let node_id = MastNodeId::from_u32_safe(node_idx as u32, mast_forests[forest_idx])
-                    .expect("the passed id should be a valid node in the forest");
+                // SAFETY: The passed id comes from the iterator over the nodes, so we never exceed
+                // the forest's number of nodes.
+                let node_id = MastNodeId::new_unsafe(node_idx as u32);
                 if !node.is_external() {
                     non_external_nodes.insert(node.digest(), (forest_idx, node_id));
                 }
@@ -335,6 +336,9 @@ pub(crate) enum MultiMastForestIteratorItem {
     },
 }
 
+// TESTS
+// ================================================================================================
+
 #[cfg(test)]
 mod tests {
     use miden_crypto::hash::rpo::RpoDigest;