refactor: Build dense MAST forests through finalization

CHANGELOG.md

@@ -13,6 +13,7 @@
 - Added trusted sparse MAST forest serialization for trace replay payloads ([#3313](https://github.com/0xMiden/miden-vm/pull/3313)).
 - [BREAKING] MAST forest payloads now include sorted root and dependency digests, so commitment inputs round trip ([#3294](https://github.com/0xMiden/miden-vm/pull/3294)).
 - [BREAKING] MAST forest and package code digests now commit to both public roots and external dependencies ([#3311](https://github.com/0xMiden/miden-vm/pull/3311)).
+- [BREAKING] Dense `MastForest`s are now finalized into canonical node order, and `MastForestContributor` no longer appends nodes directly ([#3329](https://github.com/0xMiden/miden-vm/pull/3329)).
 
 ## v0.24.0 (2026-06-24)
 

core/src/mast/dense_builder.rs

@@ -0,0 +1,105 @@
+use alloc::vec::Vec;
+
+use crate::{
+    advice::AdviceMap,
+    mast::{
+        MastForest, MastForestError, MastForestParts, MastNode, MastNodeBuilder, MastNodeContext,
+        MastNodeId,
+    },
+    utils::{DenseIdMap, Idx, IndexVec},
+};
+
+/// Construction surface for dense MAST forests.
+///
+/// The builder may append nodes while a forest is under construction. The value returned by
+/// [`Self::finish`] is a finalized [`MastForest`] with canonicalized dense node order.
+#[derive(Debug, Default)]
+pub struct DenseMastForestBuilder {
+    nodes: IndexVec<MastNodeId, MastNode>,
+    roots: Vec<MastNodeId>,
+    advice_map: AdviceMap,
+}
+
+impl DenseMastForestBuilder {
+    pub fn new() -> Self {
+        Self {
+            nodes: IndexVec::new(),
+            roots: Vec::new(),
+            advice_map: AdviceMap::default(),
+        }
+    }
+
+    pub fn push_node_builder(
+        &mut self,
+        builder: MastNodeBuilder,
+    ) -> Result<MastNodeId, MastForestError> {
+        let node = builder.build(self)?;
+        self.push_linked_node(node)
+    }
+
+    pub fn push_node(
+        &mut self,
+        builder: impl Into<MastNodeBuilder>,
+    ) -> Result<MastNodeId, MastForestError> {
+        self.push_node_builder(builder.into())
+    }
+
+    fn push_linked_node(&mut self, node: MastNode) -> Result<MastNodeId, MastForestError> {
+        self.nodes.push(node).map_err(|_| MastForestError::TooManyNodes)
+    }
+
+    pub fn get_node_by_id(&self, node_id: MastNodeId) -> Option<&MastNode> {
+        self.nodes.get(node_id)
+    }
+
+    pub fn mark_root(&mut self, root: MastNodeId) {
+        assert!(root.to_usize() < self.nodes.len());
+
+        if !self.roots.contains(&root) {
+            self.roots.push(root);
+        }
+    }
+
+    pub(crate) fn merge_advice_map(
+        &mut self,
+        advice_map: &AdviceMap,
+    ) -> Result<(), MastForestError> {
+        self.advice_map
+            .merge(advice_map)
+            .map_err(|((key, _prev), _new)| MastForestError::AdviceMapKeyCollisionOnMerge(key))
+    }
+
+    pub fn finish(self) -> Result<MastForest, MastForestError> {
+        self.finish_with_id_map().map(|(forest, _remapping)| forest)
+    }
+
+    pub fn finish_with_id_map(
+        self,
+    ) -> Result<(MastForest, DenseIdMap<MastNodeId, MastNodeId>), MastForestError> {
+        MastForest::from_parts_with_id_map(MastForestParts {
+            nodes: self.nodes,
+            roots: self.roots,
+            advice_map: self.advice_map,
+        })
+    }
+}
+
+impl MastNodeContext for DenseMastForestBuilder {
+    fn node_count(&self) -> usize {
+        self.nodes.len()
+    }
+
+    fn get_node_by_id(&self, node_id: MastNodeId) -> Option<&MastNode> {
+        self.nodes.get(node_id)
+    }
+}
+
+impl From<IndexVec<MastNodeId, MastNode>> for DenseMastForestBuilder {
+    fn from(nodes: IndexVec<MastNodeId, MastNode>) -> Self {
+        Self {
+            nodes,
+            roots: Vec::new(),
+            advice_map: AdviceMap::default(),
+        }
+    }
+}

core/src/mast/merger/mod.rs

@@ -3,8 +3,8 @@ use alloc::{collections::BTreeMap, vec::Vec};
 use crate::{
     Word,
     mast::{
-        MastForest, MastForestContributor, MastForestError, MastNode, MastNodeBuilder, MastNodeId,
-        MultiMastForestIteratorItem, MultiMastForestNodeIter,
+        DenseMastForestBuilder, MastForest, MastForestContributor, MastForestError, MastNode,
+        MastNodeBuilder, MastNodeId, MultiMastForestIteratorItem, MultiMastForestNodeIter,
     },
     utils::{DenseIdMap, IndexVec},
 };
@@ -16,7 +16,7 @@ mod tests;
 ///
 /// This functionality is exposed via [`MastForest::merge`]. See its documentation for more details.
 pub(crate) struct MastForestMerger {
-    mast_forest: MastForest,
+    mast_forest: DenseMastForestBuilder,
     // Internal indices needed for efficient duplicate checking.
     //
     // These are always in-sync with the nodes in `mast_forest`, i.e. all nodes added to the
@@ -55,13 +55,16 @@ impl MastForestMerger {
         let mut merger = Self {
             node_id_by_hash: BTreeMap::new(),
             hash_by_node_id: IndexVec::new(),
-            mast_forest: MastForest::new(),
+            mast_forest: DenseMastForestBuilder::new(),
             node_id_mappings,
         };
 
         merger.merge_inner(forests.clone())?;
 
         let Self { mast_forest, node_id_mappings, .. } = merger;
+        let (mast_forest, final_id_remapping) = mast_forest.finish_with_id_map()?;
+        let node_id_mappings =
+            Self::remap_finalized_node_ids(node_id_mappings, &final_id_remapping);
 
         let root_maps = MastForestRootMap::from_node_id_map(node_id_mappings, forests);
 
@@ -143,10 +146,7 @@ impl MastForestMerger {
     }
 
     fn merge_advice_map(&mut self, other_forest: &MastForest) -> Result<(), MastForestError> {
-        self.mast_forest
-            .advice_map
-            .merge(&other_forest.advice_map)
-            .map_err(|((key, _prev), _new)| MastForestError::AdviceMapKeyCollisionOnMerge(key))
+        self.mast_forest.merge_advice_map(other_forest.advice_map())
     }
 
     fn merge_node(
@@ -182,7 +182,7 @@ impl MastForestMerger {
             None => {
                 // If no node with a matching fingerprint exists, then the merging node is
                 // unique and we can add it to the merged forest using builders.
-                let new_node_id = remapped_builder.add_to_forest(&mut self.mast_forest)?;
+                let new_node_id = self.mast_forest.push_node_builder(remapped_builder)?;
                 self.node_id_mappings[forest_idx].insert(merging_id, new_node_id);
 
                 self.node_id_by_hash.insert(node_fingerprint, new_node_id);
@@ -232,6 +232,27 @@ impl MastForestMerger {
     ) -> Result<MastNodeBuilder, MastForestError> {
         super::build_node_with_remapped_ids(merging_id, src, original_forest, nmap)
     }
+
+    fn remap_finalized_node_ids(
+        mut node_id_mappings: Vec<DenseIdMap<MastNodeId, MastNodeId>>,
+        final_id_remapping: &DenseIdMap<MastNodeId, MastNodeId>,
+    ) -> Vec<DenseIdMap<MastNodeId, MastNodeId>> {
+        for node_id_mapping in &mut node_id_mappings {
+            for source_index in 0..node_id_mapping.len() {
+                let source_id = MastNodeId::new_unchecked(
+                    source_index.try_into().expect("source node index exceeds u32"),
+                );
+                if let Some(builder_id) = node_id_mapping.get(source_id) {
+                    let finalized_id = final_id_remapping
+                        .get(builder_id)
+                        .expect("every builder node id should map to a finalized node id");
+                    node_id_mapping.insert(source_id, finalized_id);
+                }
+            }
+        }
+
+        node_id_mappings
+    }
 }
 
 // MAST FOREST ROOT MAP

core/src/mast/merger/tests.rs

@@ -6,8 +6,7 @@ use crate::{
     advice::AdviceMap,
     mast::{
         BasicBlockNode, BasicBlockNodeBuilder, CallNodeBuilder, DynNodeBuilder,
-        ExternalNodeBuilder, LoopNodeBuilder, OpBatch,
-        node::{MastForestContributor, MastNodeExt},
+        ExternalNodeBuilder, LoopNodeBuilder, OpBatch, node::MastNodeExt,
     },
     operations::Operation,
     utils::Idx,
@@ -174,7 +173,7 @@ fn mast_node_partial_eq_is_structural_for_same_digest_blocks() {
 /// +
 /// [Block(bar), Call(bar)]
 /// =
-/// [Block(foo), Call(foo), Block(bar), Call(bar)]
+/// [Block(foo), Block(bar), Call(foo), Call(bar)]
 #[test]
 fn mast_forest_merge_remap() {
     let mut forest_a = MastForest::new();
@@ -195,17 +194,20 @@ fn mast_forest_merge_remap() {
     assert_matches!(&merged.nodes()[0], MastNode::Block(merged_block)
         if merged_block == &expected_foo_block);
 
-    assert_matches!(&merged.nodes()[1], MastNode::Call(call_node) if 0u32 == u32::from(call_node.callee()));
-
     let expected_bar_block = block_bar().build().unwrap();
-    assert_matches!(&merged.nodes()[2], MastNode::Block(merged_block)
+    assert_matches!(&merged.nodes()[1], MastNode::Block(merged_block)
         if merged_block == &expected_bar_block);
-    assert_matches!(&merged.nodes()[3], MastNode::Call(call_node) if 2u32 == u32::from(call_node.callee()));
 
-    assert_eq!(u32::from(root_maps.map_root(0, &id_call_a).unwrap()), 1u32);
-    assert_eq!(u32::from(root_maps.map_root(1, &id_call_b).unwrap()), 3u32);
+    let mapped_call_a = root_maps.map_root(0, &id_call_a).unwrap();
+    let mapped_call_b = root_maps.map_root(1, &id_call_b).unwrap();
+
+    assert_matches!(&merged[mapped_call_a], MastNode::Call(call_node)
+        if 0u32 == u32::from(call_node.callee()));
+    assert_matches!(&merged[mapped_call_b], MastNode::Call(call_node)
+        if 1u32 == u32::from(call_node.callee()));
 
     assert_child_id_lt_parent_id(&merged);
+    merged.validate_dense_node_order().unwrap();
 }
 
 /// Tests that Forest_A + Forest_A = Forest_A (i.e. duplicates are removed).