diff --git a/assembly/src/assembler/mast_forest_merger_tests.rs b/assembly/src/assembler/mast_forest_merger_tests.rs index 63197cfbd..96e533992 100644 --- a/assembly/src/assembler/mast_forest_merger_tests.rs +++ b/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); } } diff --git a/core/src/mast/merger/mod.rs b/core/src/mast/merger/mod.rs index 73daf2d20..7802f7a01 100644 --- a/core/src/mast/merger/mod.rs +++ b/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 diff --git a/core/src/mast/mod.rs b/core/src/mast/mod.rs index 172817af2..e9cf9f36d 100644 --- a/core/src/mast/mod.rs +++ b/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) } diff --git a/core/src/mast/multi_forest_node_iterator.rs b/core/src/mast/multi_forest_node_iterator.rs index 2a1dd9be0..c6a734216 100644 --- a/core/src/mast/multi_forest_node_iterator.rs +++ b/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;