diff --git a/CHANGELOG.md b/CHANGELOG.md index ffbefa7633..e041444a3d 100644 --- a/CHANGELOG.md +++ b/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) diff --git a/core/src/mast/dense_builder.rs b/core/src/mast/dense_builder.rs new file mode 100644 index 0000000000..c639c841f6 --- /dev/null +++ b/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, + roots: Vec, + 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 { + let node = builder.build(self)?; + self.push_linked_node(node) + } + + pub fn push_node( + &mut self, + builder: impl Into, + ) -> Result { + self.push_node_builder(builder.into()) + } + + fn push_linked_node(&mut self, node: MastNode) -> Result { + 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 { + self.finish_with_id_map().map(|(forest, _remapping)| forest) + } + + pub fn finish_with_id_map( + self, + ) -> Result<(MastForest, DenseIdMap), 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> for DenseMastForestBuilder { + fn from(nodes: IndexVec) -> Self { + Self { + nodes, + roots: Vec::new(), + advice_map: AdviceMap::default(), + } + } +} diff --git a/core/src/mast/merger/mod.rs b/core/src/mast/merger/mod.rs index fde6db521c..833f6861ae 100644 --- a/core/src/mast/merger/mod.rs +++ b/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 { super::build_node_with_remapped_ids(merging_id, src, original_forest, nmap) } + + fn remap_finalized_node_ids( + mut node_id_mappings: Vec>, + final_id_remapping: &DenseIdMap, + ) -> Vec> { + 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 diff --git a/core/src/mast/merger/tests.rs b/core/src/mast/merger/tests.rs index f2a23ac57c..62d227a28e 100644 --- a/core/src/mast/merger/tests.rs +++ b/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). diff --git a/core/src/mast/mod.rs b/core/src/mast/mod.rs index 45811253d2..08efe1353f 100644 --- a/core/src/mast/mod.rs +++ b/core/src/mast/mod.rs @@ -60,8 +60,8 @@ pub(crate) use node::collect_immediate_placements; pub use node::{ BasicBlockNode, BasicBlockNodeBuilder, CallNode, CallNodeBuilder, DynNode, DynNodeBuilder, ExternalNode, ExternalNodeBuilder, JoinNode, JoinNodeBuilder, LoopNode, LoopNodeBuilder, - MastForestContributor, MastNode, MastNodeBuilder, MastNodeExt, OP_BATCH_SIZE, OP_GROUP_SIZE, - OpBatch, SplitNode, SplitNodeBuilder, + MastForestContributor, MastNode, MastNodeBuilder, MastNodeContext, MastNodeExt, OP_BATCH_SIZE, + OP_GROUP_SIZE, OpBatch, SplitNode, SplitNodeBuilder, }; use crate::{ @@ -78,6 +78,9 @@ pub use serialization::{ MastForestWireView, MastNodeEntry, MastNodeInfo, SparseMastForestReadOptions, }; +mod dense_builder; +pub use dense_builder::DenseMastForestBuilder; + mod untrusted; pub use untrusted::{UntrustedMastForest, UntrustedMastForestReadOptions}; @@ -169,10 +172,27 @@ impl MastForest { Self::from_parts(MastForestParts { nodes, roots, advice_map }) } + /// Builds a [`MastForest`] from raw parts and returns the node ID remapping applied during + /// canonicalization. + #[doc(hidden)] + pub fn from_raw_parts_with_id_map( + nodes: IndexVec, + roots: Vec, + advice_map: AdviceMap, + ) -> Result<(Self, DenseIdMap), MastForestError> { + Self::from_parts_with_id_map(MastForestParts { nodes, roots, advice_map }) + } + /// Builds a [`MastForest`] from completed parts. pub(crate) fn from_parts(parts: MastForestParts) -> Result { + Self::from_parts_with_id_map(parts).map(|(forest, _remapping)| forest) + } + + pub(crate) fn from_parts_with_id_map( + parts: MastForestParts, + ) -> Result<(Self, DenseIdMap), MastForestError> { validate_mast_forest_parts_bounds(&parts)?; - let parts = canonicalize_dense_parts(parts)?; + let (parts, id_remapping) = canonicalize_dense_parts(parts)?; let forest = Self { commitment: compute_mast_forest_commitment( @@ -188,7 +208,7 @@ impl MastForest { forest.validate_dense_node_order()?; forest.validate()?; forest.validate_node_hashes()?; - Ok(forest) + Ok((forest, id_remapping)) } pub(in crate::mast) fn from_trusted_deserialization_parts( @@ -226,7 +246,9 @@ fn validate_mast_forest_parts_bounds(parts: &MastForestParts) -> Result<(), Mast Ok(()) } -fn canonicalize_dense_parts(parts: MastForestParts) -> Result { +fn canonicalize_dense_parts( + parts: MastForestParts, +) -> Result<(MastForestParts, DenseIdMap), MastForestError> { let node_count = parts.nodes.len(); let mut ordered_ids = (0..node_count) .map(|index| MastNodeId::new_unchecked(index as u32)) @@ -257,7 +279,7 @@ fn canonicalize_dense_parts(parts: MastForestParts) -> Result Result) -> Word { - let digests: Vec = if validate_dense_node_order(nodes).is_ok() { - nodes - .iter() - .take_while(|node| node.is_external()) - .map(MastNodeExt::digest) - .collect() - } else { - // Construction-phase forests may still be append-ordered. Finalized dense forests are - // canonicalized by `from_parts()` or rejected by dense deserialization. - let mut digests: Vec = nodes - .iter() - .filter(|node| node.is_external()) - .map(MastNodeExt::digest) - .collect(); - digests.sort_unstable(); - digests - }; + let digests: Vec = nodes + .iter() + .take_while(|node| node.is_external()) + .map(MastNodeExt::digest) + .collect(); Poseidon2::merge_many(&digests) } diff --git a/core/src/mast/multi_forest_node_iterator.rs b/core/src/mast/multi_forest_node_iterator.rs index f45e9c6faa..3e299d2fc8 100644 --- a/core/src/mast/multi_forest_node_iterator.rs +++ b/core/src/mast/multi_forest_node_iterator.rs @@ -312,7 +312,7 @@ mod tests { Word, mast::{ BasicBlockNodeBuilder, CallNodeBuilder, ExternalNodeBuilder, JoinNodeBuilder, - MastForestContributor, SplitNodeBuilder, + SplitNodeBuilder, }, operations::Operation, }; diff --git a/core/src/mast/node/basic_block_node/arbitrary.rs b/core/src/mast/node/basic_block_node/arbitrary.rs index 58d981c01e..54e4beb512 100644 --- a/core/src/mast/node/basic_block_node/arbitrary.rs +++ b/core/src/mast/node/basic_block_node/arbitrary.rs @@ -11,8 +11,8 @@ use crate::{ Felt, Word, advice::AdviceMap, mast::{ - CallNodeBuilder, DynNodeBuilder, ExternalNodeBuilder, JoinNodeBuilder, LoopNodeBuilder, - SplitNodeBuilder, + CallNodeBuilder, DenseMastForestBuilder, DynNodeBuilder, ExternalNodeBuilder, + JoinNodeBuilder, LoopNodeBuilder, SplitNodeBuilder, }, operations::{AssemblyOp, Operation}, program::{Kernel, Program}, @@ -400,14 +400,14 @@ impl Arbitrary for MastForest { syscall_indices, external_digests, )| { - let mut forest = MastForest::new(); + let mut forest = DenseMastForestBuilder::new(); + let empty_forest = MastForest::new(); // 2) Add basic blocks and collect their IDs let mut basic_block_ids = Vec::new(); for block in basic_blocks { - let builder = block.to_builder(&forest); - let node_id = - builder.add_to_forest(&mut forest).expect("Failed to add block"); + let builder = block.to_builder(&empty_forest); + let node_id = forest.push_node(builder).expect("Failed to add block"); basic_block_ids.push(node_id); } @@ -420,7 +420,7 @@ impl Arbitrary for MastForest { let left_id = all_node_ids[left_idx]; let right_id = all_node_ids[right_idx]; if let Ok(join_id) = - JoinNodeBuilder::new([left_id, right_id]).add_to_forest(&mut forest) + forest.push_node(JoinNodeBuilder::new([left_id, right_id])) { all_node_ids.push(join_id); } @@ -432,8 +432,8 @@ impl Arbitrary for MastForest { if true_idx < all_node_ids.len() && false_idx < all_node_ids.len() { let true_id = all_node_ids[true_idx]; let false_id = all_node_ids[false_idx]; - if let Ok(split_id) = SplitNodeBuilder::new([true_id, false_id]) - .add_to_forest(&mut forest) + if let Ok(split_id) = + forest.push_node(SplitNodeBuilder::new([true_id, false_id])) { all_node_ids.push(split_id); } @@ -444,9 +444,7 @@ impl Arbitrary for MastForest { for &body_idx in &loop_indices { if body_idx < all_node_ids.len() { let body_id = all_node_ids[body_idx]; - if let Ok(loop_id) = - LoopNodeBuilder::new(body_id).add_to_forest(&mut forest) - { + if let Ok(loop_id) = forest.push_node(LoopNodeBuilder::new(body_id)) { all_node_ids.push(loop_id); } } @@ -457,7 +455,7 @@ impl Arbitrary for MastForest { if callee_idx < all_node_ids.len() { let callee_id = all_node_ids[callee_idx]; let call_id = - CallNodeBuilder::new(callee_id).add_to_forest(&mut forest).unwrap(); + forest.push_node(CallNodeBuilder::new(callee_id)).unwrap(); all_node_ids.push(call_id); } } @@ -468,9 +466,8 @@ impl Arbitrary for MastForest { for &callee_idx in &syscall_indices { if callee_idx < all_node_ids.len() { let callee_id = all_node_ids[callee_idx]; - let syscall_id = CallNodeBuilder::new_syscall(callee_id) - .add_to_forest(&mut forest) - .unwrap(); + let syscall_id = + forest.push_node(CallNodeBuilder::new_syscall(callee_id)).unwrap(); all_node_ids.push(syscall_id); } } @@ -483,8 +480,7 @@ impl Arbitrary for MastForest { continue; } - if let Ok(external_id) = - ExternalNodeBuilder::new(digest).add_to_forest(&mut forest) + if let Ok(external_id) = forest.push_node(ExternalNodeBuilder::new(digest)) { all_node_ids.push(external_id); } @@ -494,9 +490,9 @@ impl Arbitrary for MastForest { // WARNING: These leave junk on the stack and cannot execute properly for i in 0..num_dyns { let dyn_id = if i % 2 == 0 { - DynNodeBuilder::new_dyn().add_to_forest(&mut forest).unwrap() + forest.push_node(DynNodeBuilder::new_dyn()).unwrap() } else { - DynNodeBuilder::new_dyncall().add_to_forest(&mut forest).unwrap() + forest.push_node(DynNodeBuilder::new_dyncall()).unwrap() }; all_node_ids.push(dyn_id); } @@ -506,13 +502,14 @@ impl Arbitrary for MastForest { let mut root_digest_set = BTreeSet::new(); for (i, &node_id) in all_node_ids.iter().enumerate() { if i % (all_node_ids.len() / num_roots.max(1)) == 0 - && root_digest_set.insert(forest[node_id].digest()) + && root_digest_set + .insert(forest.get_node_by_id(node_id).unwrap().digest()) { - forest.make_root(node_id); + forest.mark_root(node_id); } } - forest + forest.finish().expect("generated MAST forest should be valid") }, ) .boxed() @@ -598,25 +595,16 @@ impl Arbitrary for Program { }) .prop_map(|node| { // Create a new MastForest - let mut forest = MastForest::new(); + let mut builder = DenseMastForestBuilder::new(); + let empty_forest = MastForest::new(); // Add the node to the forest using builder - let builder = node.to_builder(&forest); - let node_id = builder.add_to_forest(&mut forest).expect("Failed to add node"); - - // Since we added a node, it should be available as a procedure root - // If not, we need to make it a root manually - let entrypoint = if forest.num_procedures() > 0 { - forest.procedure_roots()[0] - } else { - // Make the node a root manually - forest.make_root(node_id); - // After making it a root, it should be a procedure - if forest.num_procedures() == 0 { - panic!("Failed to create a valid procedure from node"); - } - forest.procedure_roots()[0] - }; + let node_builder = node.to_builder(&empty_forest); + let node_id = builder.push_node(node_builder).expect("Failed to add node"); + builder.mark_root(node_id); + let (forest, remapping) = + builder.finish_with_id_map().expect("generated program forest should be valid"); + let entrypoint = remapping.get(node_id).expect("entrypoint should be retained"); Program::new(Arc::new(forest), entrypoint) }) diff --git a/core/src/mast/node/basic_block_node/mod.rs b/core/src/mast/node/basic_block_node/mod.rs index b7d550741d..c2c444b4b6 100644 --- a/core/src/mast/node/basic_block_node/mod.rs +++ b/core/src/mast/node/basic_block_node/mod.rs @@ -1,10 +1,12 @@ use alloc::{boxed::Box, string::String, vec::Vec}; use core::{fmt, iter::repeat_n}; +#[cfg(any(test, feature = "arbitrary"))] +use crate::mast::MastNode; use crate::{ Felt, Word, ZERO, chiplets::hasher, - mast::{MastForest, MastForestError, MastNode, MastNodeId}, + mast::{MastForest, MastForestError, MastNodeId}, operations::Operation, prettier::PrettyPrint, serde::Serializable, @@ -16,7 +18,7 @@ pub use op_batch::OpBatch; use op_batch::OpBatchAccumulator; pub(crate) use op_batch::collect_immediate_placements; -use super::{MastForestContributor, MastNodeExt}; +use super::{MastForestContributor, MastNodeContext, MastNodeExt}; #[cfg(any(test, feature = "arbitrary"))] pub mod arbitrary; @@ -821,100 +823,23 @@ impl BasicBlockNodeBuilder { Ok(BasicBlockNode { op_batches, digest }) } +} - /// Add this node to a forest using relaxed validation. - /// - /// This method is used during deserialization where nodes may reference child nodes - /// that haven't been added to the forest yet. The child node IDs have already been - /// validated against the expected final node count during the `try_into_mast_node_builder` - /// step, so we can safely skip validation here. - /// - /// Note: This is not part of the `MastForestContributor` trait because it's only - /// intended for internal use during deserialization. - /// - /// For BasicBlockNode, this is equivalent to the normal `add_to_forest` since basic blocks - /// don't have child nodes to validate. - pub(in crate::mast) fn add_to_forest_relaxed( - self, - forest: &mut MastForest, - ) -> Result { - // Process based on operation data type - let (op_batches, digest) = match self.operation_data { - OperationData::Raw { operations } => { - if operations.is_empty() { - return Err(MastForestError::EmptyBasicBlock); - } - - // Batch operations (adds padding NOOPs) - let (op_batches, computed_digest) = batch_and_hash_ops(&operations); - - // Use the forced digest if provided, otherwise use the computed digest - let digest = self.digest.unwrap_or(computed_digest); - - (op_batches, digest) - }, - OperationData::Batched { op_batches } => { - if op_batches.is_empty() { - return Err(MastForestError::EmptyBasicBlock); - } - - // For batched operations, digest must be set - let digest = self.digest.expect("digest must be set for batched operations"); - - (op_batches, digest) - }, - }; - - // Create the node in the forest. - let node_id = forest +#[cfg(any(test, feature = "arbitrary"))] +impl BasicBlockNodeBuilder { + pub fn add_to_forest(self, forest: &mut MastForest) -> Result { + let node = self.build()?; + forest .nodes - .push(MastNode::Block(BasicBlockNode { op_batches, digest })) - .map_err(|_| MastForestError::TooManyNodes)?; - - Ok(node_id) + .push(MastNode::Block(node)) + .map_err(|_| MastForestError::TooManyNodes) } } impl MastForestContributor for BasicBlockNodeBuilder { - fn add_to_forest(self, forest: &mut MastForest) -> Result { - // Process based on operation data type - let (op_batches, digest) = match self.operation_data { - OperationData::Raw { operations } => { - if operations.is_empty() { - return Err(MastForestError::EmptyBasicBlock); - } - - // Batch operations (adds padding NOOPs) - let (op_batches, computed_digest) = batch_and_hash_ops(&operations); - - // Use the forced digest if provided, otherwise use the computed digest - let digest = self.digest.unwrap_or(computed_digest); - - (op_batches, digest) - }, - OperationData::Batched { op_batches } => { - if op_batches.is_empty() { - return Err(MastForestError::EmptyBasicBlock); - } - - let digest = self.digest.expect("digest must be set for batched operations"); - - (op_batches, digest) - }, - }; - - // Create the node in the forest. - let node_id = forest - .nodes - .push(MastNode::Block(BasicBlockNode { op_batches, digest })) - .map_err(|_| MastForestError::TooManyNodes)?; - - Ok(node_id) - } - fn fingerprint_for_node( &self, - _forest: &MastForest, + _context: &impl MastNodeContext, _hash_by_node_id: &impl LookupByIdx, ) -> Result { let (op_batches, digest) = match &self.operation_data { diff --git a/core/src/mast/node/call_node.rs b/core/src/mast/node/call_node.rs index 58d39c4d54..e9c1b8dd2f 100644 --- a/core/src/mast/node/call_node.rs +++ b/core/src/mast/node/call_node.rs @@ -8,13 +8,15 @@ use miden_formatting::{ #[cfg(feature = "serde")] use serde::{Deserialize, Serialize}; -use super::{MastForestContributor, MastNodeExt, fingerprint_with_child_fingerprints}; +use super::{ + MastForestContributor, MastNodeContext, MastNodeExt, fingerprint_with_child_fingerprints, +}; use crate::{ Felt, Word, chiplets::hasher, mast::{MastForest, MastForestError, MastNodeId}, operations::opcodes, - utils::{Idx, LookupByIdx}, + utils::LookupByIdx, }; // CALL NODE @@ -227,16 +229,16 @@ impl CallNodeBuilder { } /// Builds the CallNode. - pub fn build(self, mast_forest: &MastForest) -> Result { - if self.callee.to_usize() >= mast_forest.nodes.len() { - return Err(MastForestError::NodeIdOverflow(self.callee, mast_forest.nodes.len())); - } + pub fn build(self, context: &impl MastNodeContext) -> Result { + let callee = context + .get_node_by_id(self.callee) + .ok_or_else(|| MastForestError::NodeIdOverflow(self.callee, context.node_count()))?; // Use the forced digest if provided, otherwise compute the digest let digest = if let Some(forced_digest) = self.digest { forced_digest } else { - let callee_digest = mast_forest[self.callee].digest(); + let callee_digest = callee.digest(); let domain = if self.is_syscall { CallNode::SYSCALL_DOMAIN } else { @@ -262,52 +264,27 @@ impl CallNodeBuilder { } } -impl MastForestContributor for CallNodeBuilder { - fn add_to_forest(self, forest: &mut MastForest) -> Result { - if self.callee.to_usize() >= forest.nodes.len() { - return Err(MastForestError::NodeIdOverflow(self.callee, forest.nodes.len())); - } - - // Use the forced digest if provided, otherwise compute the digest directly - let digest = if let Some(forced_digest) = self.digest { - forced_digest - } else { - let callee_digest = forest[self.callee].digest(); - let domain = if self.is_syscall { - CallNode::SYSCALL_DOMAIN - } else { - CallNode::CALL_DOMAIN - }; - - hasher::merge_in_domain(&[callee_digest, Word::default()], domain) - }; - - // Create the node in the forest with Linked variant from the start - // Move the data directly without intermediate Owned node creation - let node_id = forest - .nodes - .push( - CallNode { - callee: self.callee, - is_syscall: self.is_syscall, - digest, - } - .into(), - ) - .map_err(|_| MastForestError::TooManyNodes)?; - - Ok(node_id) +#[cfg(any(test, feature = "arbitrary"))] +impl CallNodeBuilder { + pub fn add_to_forest(self, forest: &mut MastForest) -> Result { + let node = self.build(forest)?; + forest.nodes.push(node.into()).map_err(|_| MastForestError::TooManyNodes) } +} +impl MastForestContributor for CallNodeBuilder { fn fingerprint_for_node( &self, - forest: &MastForest, + context: &impl MastNodeContext, hash_by_node_id: &impl LookupByIdx, ) -> Result { let node_digest = if let Some(forced_digest) = self.digest { forced_digest } else { - let callee_digest = forest[self.callee].digest(); + let callee_digest = context + .get_node_by_id(self.callee) + .ok_or_else(|| MastForestError::NodeIdOverflow(self.callee, context.node_count()))? + .digest(); let domain = if self.is_syscall { CallNode::SYSCALL_DOMAIN } else { @@ -317,7 +294,7 @@ impl MastForestContributor for CallNodeBuilder { hasher::merge_in_domain(&[callee_digest, Word::default()], domain) }; - fingerprint_with_child_fingerprints(node_digest, &[self.callee], forest, hash_by_node_id) + fingerprint_with_child_fingerprints(node_digest, &[self.callee], context, hash_by_node_id) } fn remap_children(self, remapping: &impl LookupByIdx) -> Self { @@ -334,53 +311,14 @@ impl MastForestContributor for CallNodeBuilder { } } -impl CallNodeBuilder { - /// Add this node to a forest using relaxed validation. - /// - /// This method is used during deserialization where nodes may reference child nodes - /// that haven't been added to the forest yet. The child node IDs have already been - /// validated against the expected final node count during the `try_into_mast_node_builder` - /// step, so we can safely skip validation here. - /// - /// Note: This is not part of the `MastForestContributor` trait because it's only - /// intended for internal use during deserialization. - pub(in crate::mast) fn add_to_forest_relaxed( - self, - forest: &mut MastForest, - ) -> Result { - // Use the forced digest if provided, otherwise use a default digest - // The actual digest computation will be handled when the forest is complete - let Some(digest) = self.digest else { - return Err(MastForestError::DigestRequiredForDeserialization); - }; - - // Create the node in the forest with Linked variant from the start - // Move the data directly without intermediate cloning - let node_id = forest - .nodes - .push( - CallNode { - callee: self.callee, - is_syscall: self.is_syscall, - digest, - } - .into(), - ) - .map_err(|_| MastForestError::TooManyNodes)?; - - Ok(node_id) - } -} - #[cfg(any(test, feature = "arbitrary"))] impl proptest::prelude::Arbitrary for CallNodeBuilder { - type Parameters = CallNodeBuilderParams; + type Parameters = (); type Strategy = proptest::strategy::BoxedStrategy; - fn arbitrary_with(params: Self::Parameters) -> Self::Strategy { + fn arbitrary_with(_params: Self::Parameters) -> Self::Strategy { use proptest::prelude::*; - let _ = params; (any::(), any::()) .prop_map(|(callee, is_syscall)| { if is_syscall { @@ -392,8 +330,3 @@ impl proptest::prelude::Arbitrary for CallNodeBuilder { .boxed() } } - -/// Parameters for generating CallNodeBuilder instances -#[cfg(any(test, feature = "arbitrary"))] -#[derive(Clone, Debug, Default)] -pub struct CallNodeBuilderParams {} diff --git a/core/src/mast/node/dyn_node.rs b/core/src/mast/node/dyn_node.rs index 1c5e71eee3..3ec26ab00b 100644 --- a/core/src/mast/node/dyn_node.rs +++ b/core/src/mast/node/dyn_node.rs @@ -4,7 +4,7 @@ use core::fmt; #[cfg(feature = "serde")] use serde::{Deserialize, Serialize}; -use super::{MastForestContributor, MastNodeExt}; +use super::{MastForestContributor, MastNodeContext, MastNodeExt}; use crate::{ Felt, Word, mast::{MastForest, MastForestError, MastNodeId}, @@ -210,30 +210,18 @@ impl DynNodeBuilder { } } -impl MastForestContributor for DynNodeBuilder { - fn add_to_forest(self, forest: &mut MastForest) -> Result { - // Use the forced digest if provided, otherwise use the default digest - let digest = if let Some(forced_digest) = self.digest { - forced_digest - } else if self.is_dyncall { - DynNode::DYNCALL_DEFAULT_DIGEST - } else { - DynNode::DYN_DEFAULT_DIGEST - }; - - // Create the node in the forest with Linked variant from the start - // Move the data directly without intermediate cloning - let node_id = forest - .nodes - .push(DynNode { is_dyncall: self.is_dyncall, digest }.into()) - .map_err(|_| MastForestError::TooManyNodes)?; - - Ok(node_id) +#[cfg(any(test, feature = "arbitrary"))] +impl DynNodeBuilder { + pub fn add_to_forest(self, forest: &mut MastForest) -> Result { + let node = self.build(); + forest.nodes.push(node.into()).map_err(|_| MastForestError::TooManyNodes) } +} +impl MastForestContributor for DynNodeBuilder { fn fingerprint_for_node( &self, - _forest: &MastForest, + _context: &impl MastNodeContext, _hash_by_node_id: &impl LookupByIdx, ) -> Result { Ok(if let Some(forced_digest) = self.digest { @@ -256,40 +244,6 @@ impl MastForestContributor for DynNodeBuilder { } } -impl DynNodeBuilder { - /// Add this node to a forest using relaxed validation. - /// - /// This method is used during deserialization where nodes may reference child nodes - /// that haven't been added to the forest yet. The child node IDs have already been - /// validated against the expected final node count during the `try_into_mast_node_builder` - /// step, so we can safely skip validation here. - /// - /// Note: This is not part of the `MastForestContributor` trait because it's only - /// intended for internal use during deserialization. - pub(in crate::mast) fn add_to_forest_relaxed( - self, - forest: &mut MastForest, - ) -> Result { - // Use the forced digest if provided, otherwise use the default digest - let digest = if let Some(forced_digest) = self.digest { - forced_digest - } else if self.is_dyncall { - DynNode::DYNCALL_DEFAULT_DIGEST - } else { - DynNode::DYN_DEFAULT_DIGEST - }; - - // Create the node in the forest with Linked variant from the start - // Move the data directly without intermediate cloning - let node_id = forest - .nodes - .push(DynNode { is_dyncall: self.is_dyncall, digest }.into()) - .map_err(|_| MastForestError::TooManyNodes)?; - - Ok(node_id) - } -} - #[cfg(any(test, feature = "arbitrary"))] impl proptest::prelude::Arbitrary for DynNodeBuilder { type Parameters = (); @@ -320,15 +274,15 @@ mod tests { /// domain. #[test] pub fn test_dyn_node_digest() { - let mut forest = MastForest::new(); - let dyn_node_id = DynNodeBuilder::new_dyn().add_to_forest(&mut forest).unwrap(); + let mut forest = crate::mast::DenseMastForestBuilder::new(); + let dyn_node_id = forest.push_node(DynNodeBuilder::new_dyn()).unwrap(); let dyn_node = forest.get_node_by_id(dyn_node_id).unwrap().unwrap_dyn(); assert_eq!( dyn_node.digest(), Poseidon2::merge_in_domain(&[Word::default(), Word::default()], DynNode::DYN_DOMAIN) ); - let dyncall_node_id = DynNodeBuilder::new_dyncall().add_to_forest(&mut forest).unwrap(); + let dyncall_node_id = forest.push_node(DynNodeBuilder::new_dyncall()).unwrap(); let dyncall_node = forest.get_node_by_id(dyncall_node_id).unwrap().unwrap_dyn(); assert_eq!( dyncall_node.digest(), diff --git a/core/src/mast/node/external.rs b/core/src/mast/node/external.rs index 9861824c07..79e7969c9b 100644 --- a/core/src/mast/node/external.rs +++ b/core/src/mast/node/external.rs @@ -8,7 +8,7 @@ use miden_formatting::{ #[cfg(feature = "serde")] use serde::{Deserialize, Serialize}; -use super::{MastForestContributor, MastNodeExt}; +use super::{MastForestContributor, MastNodeContext, MastNodeExt}; use crate::{ Felt, Word, mast::{MastForest, MastForestError, MastNodeId}, @@ -153,19 +153,22 @@ impl ExternalNodeBuilder { } } -impl MastForestContributor for ExternalNodeBuilder { - fn add_to_forest(self, forest: &mut MastForest) -> Result { +#[cfg(any(test, feature = "arbitrary"))] +impl ExternalNodeBuilder { + pub fn add_to_forest(self, forest: &mut MastForest) -> Result { let node_id = forest .nodes - .push(ExternalNode { digest: self.digest }.into()) + .push(self.build().into()) .map_err(|_| MastForestError::TooManyNodes)?; forest.commitment = forest.compute_mast_forest_commitment(); Ok(node_id) } +} +impl MastForestContributor for ExternalNodeBuilder { fn fingerprint_for_node( &self, - _forest: &MastForest, + _context: &impl MastNodeContext, _hash_by_node_id: &impl LookupByIdx, ) -> Result { Ok(self.digest) @@ -182,38 +185,14 @@ impl MastForestContributor for ExternalNodeBuilder { } } -impl ExternalNodeBuilder { - /// Add this node to a forest using relaxed validation. - /// - /// This method is used during deserialization where nodes may reference child nodes - /// that haven't been added to the forest yet. The child node IDs have already been - /// validated against the expected final node count during the `try_into_mast_node_builder` - /// step, so we can safely skip validation here. - /// - /// Note: This is not part of the `MastForestContributor` trait because it's only - /// intended for internal use during deserialization. - pub(in crate::mast) fn add_to_forest_relaxed( - self, - forest: &mut MastForest, - ) -> Result { - let node_id = forest - .nodes - .push(ExternalNode { digest: self.digest }.into()) - .map_err(|_| MastForestError::TooManyNodes)?; - forest.commitment = forest.compute_mast_forest_commitment(); - Ok(node_id) - } -} - #[cfg(any(test, feature = "arbitrary"))] impl proptest::prelude::Arbitrary for ExternalNodeBuilder { - type Parameters = ExternalNodeBuilderParams; + type Parameters = (); type Strategy = proptest::strategy::BoxedStrategy; - fn arbitrary_with(params: Self::Parameters) -> Self::Strategy { + fn arbitrary_with(_params: Self::Parameters) -> Self::Strategy { use proptest::prelude::*; - let _ = params; any::<[u64; 4]>() .prop_map(|[a, b, c, d]| { Word::new([ @@ -227,8 +206,3 @@ impl proptest::prelude::Arbitrary for ExternalNodeBuilder { .boxed() } } - -/// Parameters for generating ExternalNodeBuilder instances -#[cfg(any(test, feature = "arbitrary"))] -#[derive(Clone, Debug, Default)] -pub struct ExternalNodeBuilderParams {} diff --git a/core/src/mast/node/join_node.rs b/core/src/mast/node/join_node.rs index 849a0fb09e..3bb8b01140 100644 --- a/core/src/mast/node/join_node.rs +++ b/core/src/mast/node/join_node.rs @@ -4,14 +4,16 @@ use core::fmt; #[cfg(feature = "serde")] use serde::{Deserialize, Serialize}; -use super::{MastForestContributor, MastNodeExt, fingerprint_with_child_fingerprints}; +use super::{ + MastForestContributor, MastNodeContext, MastNodeExt, fingerprint_with_child_fingerprints, +}; use crate::{ Felt, Word, chiplets::hasher, mast::{MastForest, MastForestError, MastNodeId}, operations::opcodes, prettier::PrettyPrint, - utils::{Idx, LookupByIdx}, + utils::LookupByIdx, }; // JOIN NODE @@ -196,20 +198,20 @@ impl JoinNodeBuilder { } /// Builds the JoinNode. - pub fn build(self, mast_forest: &MastForest) -> Result { - let forest_len = mast_forest.nodes.len(); - if self.children[0].to_usize() >= forest_len { - return Err(MastForestError::NodeIdOverflow(self.children[0], forest_len)); - } else if self.children[1].to_usize() >= forest_len { - return Err(MastForestError::NodeIdOverflow(self.children[1], forest_len)); - } + pub fn build(self, context: &impl MastNodeContext) -> Result { + let left_child = context.get_node_by_id(self.children[0]).ok_or_else(|| { + MastForestError::NodeIdOverflow(self.children[0], context.node_count()) + })?; + let right_child = context.get_node_by_id(self.children[1]).ok_or_else(|| { + MastForestError::NodeIdOverflow(self.children[1], context.node_count()) + })?; // Use the forced digest if provided, otherwise compute the digest let digest = if let Some(forced_digest) = self.digest { forced_digest } else { - let left_child_hash = mast_forest[self.children[0]].digest(); - let right_child_hash = mast_forest[self.children[1]].digest(); + let left_child_hash = left_child.digest(); + let right_child_hash = right_child.digest(); hasher::merge_in_domain(&[left_child_hash, right_child_hash], JoinNode::DOMAIN) }; @@ -225,51 +227,40 @@ impl JoinNodeBuilder { } } -impl MastForestContributor for JoinNodeBuilder { - fn add_to_forest(self, forest: &mut MastForest) -> Result { - // Validate child node IDs - let forest_len = forest.nodes.len(); - if self.children[0].to_usize() >= forest_len { - return Err(MastForestError::NodeIdOverflow(self.children[0], forest_len)); - } else if self.children[1].to_usize() >= forest_len { - return Err(MastForestError::NodeIdOverflow(self.children[1], forest_len)); - } - - // Use the forced digest if provided, otherwise compute the digest - let digest = if let Some(forced_digest) = self.digest { - forced_digest - } else { - let left_child_hash = forest[self.children[0]].digest(); - let right_child_hash = forest[self.children[1]].digest(); - - hasher::merge_in_domain(&[left_child_hash, right_child_hash], JoinNode::DOMAIN) - }; - - // Create the node in the forest with Linked variant from the start - // Move the data directly without intermediate cloning - let node_id = forest - .nodes - .push(JoinNode { children: self.children, digest }.into()) - .map_err(|_| MastForestError::TooManyNodes)?; - - Ok(node_id) +#[cfg(any(test, feature = "arbitrary"))] +impl JoinNodeBuilder { + pub fn add_to_forest(self, forest: &mut MastForest) -> Result { + let node = self.build(forest)?; + forest.nodes.push(node.into()).map_err(|_| MastForestError::TooManyNodes) } +} +impl MastForestContributor for JoinNodeBuilder { fn fingerprint_for_node( &self, - forest: &MastForest, + context: &impl MastNodeContext, hash_by_node_id: &impl LookupByIdx, ) -> Result { let node_digest = if let Some(forced_digest) = self.digest { forced_digest } else { - let left_child_hash = forest[self.children[0]].digest(); - let right_child_hash = forest[self.children[1]].digest(); + let left_child_hash = context + .get_node_by_id(self.children[0]) + .ok_or_else(|| { + MastForestError::NodeIdOverflow(self.children[0], context.node_count()) + })? + .digest(); + let right_child_hash = context + .get_node_by_id(self.children[1]) + .ok_or_else(|| { + MastForestError::NodeIdOverflow(self.children[1], context.node_count()) + })? + .digest(); hasher::merge_in_domain(&[left_child_hash, right_child_hash], JoinNode::DOMAIN) }; - fingerprint_with_child_fingerprints(node_digest, &self.children, forest, hash_by_node_id) + fingerprint_with_child_fingerprints(node_digest, &self.children, context, hash_by_node_id) } fn remap_children(self, remapping: &impl LookupByIdx) -> Self { @@ -288,51 +279,14 @@ impl MastForestContributor for JoinNodeBuilder { } } -impl JoinNodeBuilder { - /// Add this node to a forest using relaxed validation. - /// - /// This method is used during deserialization where nodes may reference child nodes - /// that haven't been added to the forest yet. The child node IDs have already been - /// validated against the expected final node count during the `try_into_mast_node_builder` - /// step, so we can safely skip validation here. - /// - /// Note: This is not part of the `MastForestContributor` trait because it's only - /// intended for internal use during deserialization. - pub(in crate::mast) fn add_to_forest_relaxed( - self, - forest: &mut MastForest, - ) -> Result { - // Use the forced digest if provided, otherwise use a default digest - // The actual digest computation will be handled when the forest is complete - let Some(digest) = self.digest else { - return Err(MastForestError::DigestRequiredForDeserialization); - }; - - // Create the node in the forest with Linked variant from the start - // Move the data directly without intermediate cloning - let node_id = forest - .nodes - .push(JoinNode { children: self.children, digest }.into()) - .map_err(|_| MastForestError::TooManyNodes)?; - - Ok(node_id) - } -} - #[cfg(any(test, feature = "arbitrary"))] impl proptest::prelude::Arbitrary for JoinNodeBuilder { - type Parameters = JoinNodeBuilderParams; + type Parameters = (); type Strategy = proptest::strategy::BoxedStrategy; - fn arbitrary_with(params: Self::Parameters) -> Self::Strategy { + fn arbitrary_with(_params: Self::Parameters) -> Self::Strategy { use proptest::prelude::*; - let _ = params; any::<[MastNodeId; 2]>().prop_map(Self::new).boxed() } } - -/// Parameters for generating JoinNodeBuilder instances -#[cfg(any(test, feature = "arbitrary"))] -#[derive(Clone, Debug, Default)] -pub struct JoinNodeBuilderParams {} diff --git a/core/src/mast/node/loop_node.rs b/core/src/mast/node/loop_node.rs index 7bd6add74b..448477b0bc 100644 --- a/core/src/mast/node/loop_node.rs +++ b/core/src/mast/node/loop_node.rs @@ -4,14 +4,16 @@ use core::fmt; #[cfg(feature = "serde")] use serde::{Deserialize, Serialize}; -use super::{MastForestContributor, MastNodeExt, fingerprint_with_child_fingerprints}; +use super::{ + MastForestContributor, MastNodeContext, MastNodeExt, fingerprint_with_child_fingerprints, +}; use crate::{ Felt, Word, chiplets::hasher, mast::{MastForest, MastForestError, MastNodeId}, operations::opcodes, prettier::PrettyPrint, - utils::{Idx, LookupByIdx}, + utils::LookupByIdx, }; // LOOP NODE @@ -180,16 +182,16 @@ impl LoopNodeBuilder { } /// Builds the LoopNode. - pub fn build(self, mast_forest: &MastForest) -> Result { - if self.body.to_usize() >= mast_forest.nodes.len() { - return Err(MastForestError::NodeIdOverflow(self.body, mast_forest.nodes.len())); - } + pub fn build(self, context: &impl MastNodeContext) -> Result { + let body = context + .get_node_by_id(self.body) + .ok_or_else(|| MastForestError::NodeIdOverflow(self.body, context.node_count()))?; // Use the forced digest if provided, otherwise compute the digest let digest = if let Some(forced_digest) = self.digest { forced_digest } else { - let body_hash = mast_forest[self.body].digest(); + let body_hash = body.digest(); hasher::merge_in_domain(&[body_hash, Word::default()], LoopNode::DOMAIN) }; @@ -205,31 +207,32 @@ impl LoopNodeBuilder { } } -impl MastForestContributor for LoopNodeBuilder { - fn add_to_forest(self, forest: &mut MastForest) -> Result { +#[cfg(any(test, feature = "arbitrary"))] +impl LoopNodeBuilder { + pub fn add_to_forest(self, forest: &mut MastForest) -> Result { let node = self.build(forest)?; - - // Create the node in the forest with Linked variant from the start - // Move the data directly without intermediate cloning - let node_id = forest.nodes.push(node.into()).map_err(|_| MastForestError::TooManyNodes)?; - - Ok(node_id) + forest.nodes.push(node.into()).map_err(|_| MastForestError::TooManyNodes) } +} +impl MastForestContributor for LoopNodeBuilder { fn fingerprint_for_node( &self, - forest: &MastForest, + context: &impl MastNodeContext, hash_by_node_id: &impl LookupByIdx, ) -> Result { let node_digest = if let Some(forced_digest) = self.digest { forced_digest } else { - let body_hash = forest[self.body].digest(); + let body_hash = context + .get_node_by_id(self.body) + .ok_or_else(|| MastForestError::NodeIdOverflow(self.body, context.node_count()))? + .digest(); hasher::merge_in_domain(&[body_hash, Word::default()], LoopNode::DOMAIN) }; - fingerprint_with_child_fingerprints(node_digest, &[self.body], forest, hash_by_node_id) + fingerprint_with_child_fingerprints(node_digest, &[self.body], context, hash_by_node_id) } fn remap_children(self, remapping: &impl LookupByIdx) -> Self { @@ -245,51 +248,14 @@ impl MastForestContributor for LoopNodeBuilder { } } -impl LoopNodeBuilder { - /// Add this node to a forest using relaxed validation. - /// - /// This method is used during deserialization where nodes may reference child nodes - /// that haven't been added to the forest yet. The child node IDs have already been - /// validated against the expected final node count during the `try_into_mast_node_builder` - /// step, so we can safely skip validation here. - /// - /// Note: This is not part of the `MastForestContributor` trait because it's only - /// intended for internal use during deserialization. - pub(in crate::mast) fn add_to_forest_relaxed( - self, - forest: &mut MastForest, - ) -> Result { - // Use the forced digest if provided, otherwise use a default digest - // The actual digest computation will be handled when the forest is complete - let Some(digest) = self.digest else { - return Err(MastForestError::DigestRequiredForDeserialization); - }; - - // Create the node in the forest with Linked variant from the start - // Move the data directly without intermediate cloning - let node_id = forest - .nodes - .push(LoopNode { body: self.body, digest }.into()) - .map_err(|_| MastForestError::TooManyNodes)?; - - Ok(node_id) - } -} - #[cfg(any(test, feature = "arbitrary"))] impl proptest::prelude::Arbitrary for LoopNodeBuilder { - type Parameters = LoopNodeBuilderParams; + type Parameters = (); type Strategy = proptest::strategy::BoxedStrategy; - fn arbitrary_with(params: Self::Parameters) -> Self::Strategy { + fn arbitrary_with(_params: Self::Parameters) -> Self::Strategy { use proptest::prelude::*; - let _ = params; any::().prop_map(Self::new).boxed() } } - -/// Parameters for generating LoopNodeBuilder instances -#[cfg(any(test, feature = "arbitrary"))] -#[derive(Clone, Debug, Default)] -pub struct LoopNodeBuilderParams {} diff --git a/core/src/mast/node/mast_forest_contributor.rs b/core/src/mast/node/mast_forest_contributor.rs index f15e76c6cd..3f296f4935 100644 --- a/core/src/mast/node/mast_forest_contributor.rs +++ b/core/src/mast/node/mast_forest_contributor.rs @@ -10,7 +10,7 @@ use crate::{ Felt, Word, chiplets::hasher, mast::{MastForest, MastForestError, MastNode, MastNodeId}, - utils::{Idx, LookupByIdx}, + utils::LookupByIdx, }; const CHILD_FINGERPRINT_DOMAIN: Felt = Felt::new_unchecked(0x2473_0002); @@ -18,7 +18,7 @@ const CHILD_FINGERPRINT_DOMAIN: Felt = Felt::new_unchecked(0x2473_0002); pub(crate) fn fingerprint_with_child_fingerprints( node_digest: Word, child_ids: &[MastNodeId], - forest: &MastForest, + context: &impl MastNodeContext, fingerprint_by_node_id: &impl LookupByIdx, ) -> Result { let mut has_non_digest_child = false; @@ -27,14 +27,13 @@ pub(crate) fn fingerprint_with_child_fingerprints( elements.extend_from_slice(node_digest.as_elements()); for &child_id in child_ids { - if child_id.to_usize() >= forest.nodes().len() { - return Err(MastForestError::NodeIdOverflow(child_id, forest.nodes().len())); - } - - let child_digest = forest[child_id].digest(); + let child = context + .get_node_by_id(child_id) + .ok_or_else(|| MastForestError::NodeIdOverflow(child_id, context.node_count()))?; + let child_digest = child.digest(); let child_fingerprint = *fingerprint_by_node_id .get(child_id) - .ok_or(MastForestError::NodeIdOverflow(child_id, forest.nodes().len()))?; + .ok_or_else(|| MastForestError::NodeIdOverflow(child_id, context.node_count()))?; has_non_digest_child |= child_fingerprint != child_digest; elements.extend_from_slice(child_fingerprint.as_elements()); } @@ -46,9 +45,23 @@ pub(crate) fn fingerprint_with_child_fingerprints( } } -pub trait MastForestContributor { - fn add_to_forest(self, forest: &mut MastForest) -> Result; +pub trait MastNodeContext { + fn node_count(&self) -> usize; + fn get_node_by_id(&self, node_id: MastNodeId) -> Option<&MastNode>; +} + +impl MastNodeContext for MastForest { + fn node_count(&self) -> usize { + self.nodes().len() + } + + fn get_node_by_id(&self, node_id: MastNodeId) -> Option<&MastNode> { + self.get_node_by_id(node_id) + } +} + +pub trait MastForestContributor { /// Returns the fingerprint for this builder without constructing a MastNode. /// /// This method computes the fingerprint for a node directly from the builder data @@ -56,7 +69,7 @@ pub trait MastForestContributor { /// traditional fingerprint computation approach. fn fingerprint_for_node( &self, - forest: &MastForest, + context: &impl MastNodeContext, hash_by_node_id: &impl LookupByIdx, ) -> Result; @@ -84,21 +97,76 @@ pub enum MastNodeBuilder { Split(SplitNodeBuilder), } +impl From for MastNodeBuilder { + fn from(builder: BasicBlockNodeBuilder) -> Self { + Self::BasicBlock(builder) + } +} + +impl From for MastNodeBuilder { + fn from(builder: CallNodeBuilder) -> Self { + Self::Call(builder) + } +} + +impl From for MastNodeBuilder { + fn from(builder: DynNodeBuilder) -> Self { + Self::Dyn(builder) + } +} + +impl From for MastNodeBuilder { + fn from(builder: ExternalNodeBuilder) -> Self { + Self::External(builder) + } +} + +impl From for MastNodeBuilder { + fn from(builder: JoinNodeBuilder) -> Self { + Self::Join(builder) + } +} + +impl From for MastNodeBuilder { + fn from(builder: LoopNodeBuilder) -> Self { + Self::Loop(builder) + } +} + +impl From for MastNodeBuilder { + fn from(builder: SplitNodeBuilder) -> Self { + Self::Split(builder) + } +} + impl MastNodeBuilder { + #[cfg(any(test, feature = "arbitrary"))] + pub fn add_to_forest(self, forest: &mut MastForest) -> Result { + match self { + MastNodeBuilder::BasicBlock(builder) => builder.add_to_forest(forest), + MastNodeBuilder::Call(builder) => builder.add_to_forest(forest), + MastNodeBuilder::Dyn(builder) => builder.add_to_forest(forest), + MastNodeBuilder::External(builder) => builder.add_to_forest(forest), + MastNodeBuilder::Join(builder) => builder.add_to_forest(forest), + MastNodeBuilder::Loop(builder) => builder.add_to_forest(forest), + MastNodeBuilder::Split(builder) => builder.add_to_forest(forest), + } + } + /// Build the node from this builder. /// /// For nodes that depend on a MastForest (Call, Join, Loop, Split), the forest is required. /// For nodes that don't depend on a MastForest (BasicBlock, Dyn, External), the forest is /// ignored. - pub fn build(self, mast_forest: &MastForest) -> Result { + pub fn build(self, context: &impl MastNodeContext) -> Result { match self { MastNodeBuilder::BasicBlock(builder) => Ok(builder.build()?.into()), - MastNodeBuilder::Call(builder) => Ok(builder.build(mast_forest)?.into()), + MastNodeBuilder::Call(builder) => Ok(builder.build(context)?.into()), MastNodeBuilder::Dyn(builder) => Ok(builder.build().into()), MastNodeBuilder::External(builder) => Ok(builder.build().into()), - MastNodeBuilder::Join(builder) => Ok(builder.build(mast_forest)?.into()), - MastNodeBuilder::Loop(builder) => Ok(builder.build(mast_forest)?.into()), - MastNodeBuilder::Split(builder) => Ok(builder.build(mast_forest)?.into()), + MastNodeBuilder::Join(builder) => Ok(builder.build(context)?.into()), + MastNodeBuilder::Loop(builder) => Ok(builder.build(context)?.into()), + MastNodeBuilder::Split(builder) => Ok(builder.build(context)?.into()), } } @@ -118,26 +186,6 @@ impl MastNodeBuilder { MastNodeBuilder::Split(builder) => Ok(builder.build_linked()?.into()), } } - - /// Adds the node from this builder to the forest without validation, used during - /// deserialization. - /// - /// This method bypasses normal validation. It should only be used during deserialization where - /// the forest structure is being reconstructed. - pub(in crate::mast) fn add_to_forest_relaxed( - self, - mast_forest: &mut MastForest, - ) -> Result { - match self { - MastNodeBuilder::BasicBlock(builder) => builder.add_to_forest_relaxed(mast_forest), - MastNodeBuilder::Call(builder) => builder.add_to_forest_relaxed(mast_forest), - MastNodeBuilder::Dyn(builder) => builder.add_to_forest_relaxed(mast_forest), - MastNodeBuilder::External(builder) => builder.add_to_forest_relaxed(mast_forest), - MastNodeBuilder::Join(builder) => builder.add_to_forest_relaxed(mast_forest), - MastNodeBuilder::Loop(builder) => builder.add_to_forest_relaxed(mast_forest), - MastNodeBuilder::Split(builder) => builder.add_to_forest_relaxed(mast_forest), - } - } } #[cfg(any(test, feature = "arbitrary"))] @@ -203,7 +251,7 @@ mod round_trip_tests { use crate::{ Word, mast::{ - BasicBlockNodeBuilder, MastForest, MastNodeBuilder, MastNodeExt, + BasicBlockNodeBuilder, DenseMastForestBuilder, MastNodeBuilder, MastNodeExt, node::mast_forest_contributor::MastForestContributor, }, operations::Operation, @@ -211,15 +259,13 @@ mod round_trip_tests { #[test] fn test_mast_node_builder_enum_digest_forcing() { - let mut forest = MastForest::new(); + let mut forest = DenseMastForestBuilder::new(); let mast_builder1 = MastNodeBuilder::BasicBlock(BasicBlockNodeBuilder::new(vec![Operation::Push( Felt::new_unchecked(10), )])); - let mast_node_id1 = mast_builder1 - .add_to_forest(&mut forest) - .expect("Failed to add mast node1 to forest"); + let mast_node_id1 = forest.push_node(mast_builder1).expect("failed to add mast node1"); let mast_node1 = forest.get_node_by_id(mast_node_id1).unwrap().unwrap_basic_block(); let mast_normal_digest = mast_node1.digest(); @@ -233,9 +279,9 @@ mod round_trip_tests { BasicBlockNodeBuilder::new(vec![Operation::Push(Felt::new_unchecked(10))]) .with_digest(forced_mast_digest), ); - let mast_node_id2 = mast_builder2 - .add_to_forest(&mut forest) - .expect("Failed to add mast node with forced digest to forest"); + let mast_node_id2 = forest + .push_node(mast_builder2) + .expect("failed to add mast node with forced digest"); let mast_node2 = forest.get_node_by_id(mast_node_id2).unwrap().unwrap_basic_block(); assert_ne!( diff --git a/core/src/mast/node/mod.rs b/core/src/mast/node/mod.rs index 3e3fbb7f0e..73d7924fb0 100644 --- a/core/src/mast/node/mod.rs +++ b/core/src/mast/node/mod.rs @@ -34,7 +34,7 @@ pub use loop_node::{LoopNode, LoopNodeBuilder}; mod mast_forest_contributor; pub(super) use mast_forest_contributor::fingerprint_with_child_fingerprints; -pub use mast_forest_contributor::{MastForestContributor, MastNodeBuilder}; +pub use mast_forest_contributor::{MastForestContributor, MastNodeBuilder, MastNodeContext}; use crate::mast::{MastForest, MastNodeId}; diff --git a/core/src/mast/node/split_node.rs b/core/src/mast/node/split_node.rs index dffe8248f0..1369f62cef 100644 --- a/core/src/mast/node/split_node.rs +++ b/core/src/mast/node/split_node.rs @@ -4,14 +4,16 @@ use core::fmt; #[cfg(feature = "serde")] use serde::{Deserialize, Serialize}; -use super::{MastForestContributor, MastNodeExt, fingerprint_with_child_fingerprints}; +use super::{ + MastForestContributor, MastNodeContext, MastNodeExt, fingerprint_with_child_fingerprints, +}; use crate::{ Felt, Word, chiplets::hasher, mast::{MastForest, MastForestError, MastNodeId}, operations::opcodes, prettier::PrettyPrint, - utils::{Idx, LookupByIdx}, + utils::LookupByIdx, }; // SPLIT NODE @@ -194,20 +196,20 @@ impl SplitNodeBuilder { } /// Builds the SplitNode. - pub fn build(self, mast_forest: &MastForest) -> Result { - let forest_len = mast_forest.nodes.len(); - if self.branches[0].to_usize() >= forest_len { - return Err(MastForestError::NodeIdOverflow(self.branches[0], forest_len)); - } else if self.branches[1].to_usize() >= forest_len { - return Err(MastForestError::NodeIdOverflow(self.branches[1], forest_len)); - } + pub fn build(self, context: &impl MastNodeContext) -> Result { + let true_branch = context.get_node_by_id(self.branches[0]).ok_or_else(|| { + MastForestError::NodeIdOverflow(self.branches[0], context.node_count()) + })?; + let false_branch = context.get_node_by_id(self.branches[1]).ok_or_else(|| { + MastForestError::NodeIdOverflow(self.branches[1], context.node_count()) + })?; // Use the forced digest if provided, otherwise compute the digest let digest = if let Some(forced_digest) = self.digest { forced_digest } else { - let true_branch_hash = mast_forest[self.branches[0]].digest(); - let false_branch_hash = mast_forest[self.branches[1]].digest(); + let true_branch_hash = true_branch.digest(); + let false_branch_hash = false_branch.digest(); hasher::merge_in_domain(&[true_branch_hash, false_branch_hash], SplitNode::DOMAIN) }; @@ -223,51 +225,40 @@ impl SplitNodeBuilder { } } -impl MastForestContributor for SplitNodeBuilder { - fn add_to_forest(self, forest: &mut MastForest) -> Result { - // Validate branch node IDs - let forest_len = forest.nodes.len(); - if self.branches[0].to_usize() >= forest_len { - return Err(MastForestError::NodeIdOverflow(self.branches[0], forest_len)); - } else if self.branches[1].to_usize() >= forest_len { - return Err(MastForestError::NodeIdOverflow(self.branches[1], forest_len)); - } - - // Use the forced digest if provided, otherwise compute the digest - let digest = if let Some(forced_digest) = self.digest { - forced_digest - } else { - let true_branch_hash = forest[self.branches[0]].digest(); - let false_branch_hash = forest[self.branches[1]].digest(); - - hasher::merge_in_domain(&[true_branch_hash, false_branch_hash], SplitNode::DOMAIN) - }; - - // Create the node in the forest with Linked variant from the start - // Move the data directly without intermediate cloning - let node_id = forest - .nodes - .push(SplitNode { branches: self.branches, digest }.into()) - .map_err(|_| MastForestError::TooManyNodes)?; - - Ok(node_id) +#[cfg(any(test, feature = "arbitrary"))] +impl SplitNodeBuilder { + pub fn add_to_forest(self, forest: &mut MastForest) -> Result { + let node = self.build(forest)?; + forest.nodes.push(node.into()).map_err(|_| MastForestError::TooManyNodes) } +} +impl MastForestContributor for SplitNodeBuilder { fn fingerprint_for_node( &self, - forest: &MastForest, + context: &impl MastNodeContext, hash_by_node_id: &impl LookupByIdx, ) -> Result { let node_digest = if let Some(forced_digest) = self.digest { forced_digest } else { - let if_branch_hash = forest[self.branches[0]].digest(); - let else_branch_hash = forest[self.branches[1]].digest(); + let if_branch_hash = context + .get_node_by_id(self.branches[0]) + .ok_or_else(|| { + MastForestError::NodeIdOverflow(self.branches[0], context.node_count()) + })? + .digest(); + let else_branch_hash = context + .get_node_by_id(self.branches[1]) + .ok_or_else(|| { + MastForestError::NodeIdOverflow(self.branches[1], context.node_count()) + })? + .digest(); hasher::merge_in_domain(&[if_branch_hash, else_branch_hash], SplitNode::DOMAIN) }; - fingerprint_with_child_fingerprints(node_digest, &self.branches, forest, hash_by_node_id) + fingerprint_with_child_fingerprints(node_digest, &self.branches, context, hash_by_node_id) } fn remap_children(self, remapping: &impl LookupByIdx) -> Self { @@ -286,51 +277,14 @@ impl MastForestContributor for SplitNodeBuilder { } } -impl SplitNodeBuilder { - /// Add this node to a forest using relaxed validation. - /// - /// This method is used during deserialization where nodes may reference child nodes - /// that haven't been added to the forest yet. The child node IDs have already been - /// validated against the expected final node count during the `try_into_mast_node_builder` - /// step, so we can safely skip validation here. - /// - /// Note: This is not part of the `MastForestContributor` trait because it's only - /// intended for internal use during deserialization. - pub(in crate::mast) fn add_to_forest_relaxed( - self, - forest: &mut MastForest, - ) -> Result { - // Use the forced digest if provided, otherwise use a default digest - // The actual digest computation will be handled when the forest is complete - let Some(digest) = self.digest else { - return Err(MastForestError::DigestRequiredForDeserialization); - }; - - // Create the node in the forest with Linked variant from the start - // Move the data directly without intermediate cloning - let node_id = forest - .nodes - .push(SplitNode { branches: self.branches, digest }.into()) - .map_err(|_| MastForestError::TooManyNodes)?; - - Ok(node_id) - } -} - #[cfg(any(test, feature = "arbitrary"))] impl proptest::prelude::Arbitrary for SplitNodeBuilder { - type Parameters = SplitNodeBuilderParams; + type Parameters = (); type Strategy = proptest::strategy::BoxedStrategy; - fn arbitrary_with(params: Self::Parameters) -> Self::Strategy { + fn arbitrary_with(_params: Self::Parameters) -> Self::Strategy { use proptest::prelude::*; - let _ = params; any::<[MastNodeId; 2]>().prop_map(Self::new).boxed() } } - -/// Parameters for generating SplitNodeBuilder instances -#[cfg(any(test, feature = "arbitrary"))] -#[derive(Clone, Debug, Default)] -pub struct SplitNodeBuilderParams {} diff --git a/core/src/mast/serialization/mod.rs b/core/src/mast/serialization/mod.rs index 15b5abbf90..5cad3c63df 100644 --- a/core/src/mast/serialization/mod.rs +++ b/core/src/mast/serialization/mod.rs @@ -239,19 +239,8 @@ impl MastForest { /// /// Current writers encode normal execution payloads or hashless validation payloads. fn write_into_with_options(&self, target: &mut W, hashless: bool) { - if self.validate_dense_node_order().is_err() { - // Construction-phase forests may still be append-ordered. Finalized forests take the - // direct path below; this fallback prevents public serialization from panicking before - // callers have moved through a finalization path. - let canonical = MastForest::from_raw_parts( - self.nodes.clone(), - self.roots.clone(), - self.advice_map.clone(), - ) - .expect("dense MAST forest must be valid before serialization"); - canonical.write_into_with_options(target, hashless); - return; - } + self.validate_dense_node_order() + .expect("dense MAST forest must be canonical before serialization"); let mut basic_block_data_builder = BasicBlockDataBuilder::new(); @@ -366,16 +355,15 @@ pub enum MastForestReadView<'a> { /// /// ``` /// use miden_core::{ -/// mast::{BasicBlockNodeBuilder, MastForest, MastForestContributor, MastForestWireView}, +/// mast::{BasicBlockNodeBuilder, DenseMastForestBuilder, MastForestWireView}, /// operations::Operation, /// serde::Serializable, /// }; /// -/// let mut forest = MastForest::new(); -/// let block_id = BasicBlockNodeBuilder::new(vec![Operation::Add]) -/// .add_to_forest(&mut forest) -/// .unwrap(); -/// forest.make_root(block_id); +/// let mut builder = DenseMastForestBuilder::new(); +/// let block_id = builder.push_node(BasicBlockNodeBuilder::new(vec![Operation::Add])).unwrap(); +/// builder.mark_root(block_id); +/// let forest = builder.finish().unwrap(); /// /// let mut bytes = Vec::new(); /// forest.write_into(&mut bytes); @@ -423,16 +411,15 @@ impl<'a> MastForestWireView<'a> { /// /// ``` /// use miden_core::{ - /// mast::{BasicBlockNodeBuilder, MastForest, MastForestContributor, MastForestWireView}, + /// mast::{BasicBlockNodeBuilder, DenseMastForestBuilder, MastForestWireView}, /// operations::Operation, /// serde::Serializable, /// }; /// - /// let mut forest = MastForest::new(); - /// let block_id = BasicBlockNodeBuilder::new(vec![Operation::Add]) - /// .add_to_forest(&mut forest) - /// .unwrap(); - /// forest.make_root(block_id); + /// let mut builder = DenseMastForestBuilder::new(); + /// let block_id = builder.push_node(BasicBlockNodeBuilder::new(vec![Operation::Add])).unwrap(); + /// builder.mark_root(block_id); + /// let forest = builder.finish().unwrap(); /// /// let mut bytes = Vec::new(); /// forest.write_into(&mut bytes); @@ -481,16 +468,15 @@ impl<'a> MastForestWireView<'a> { /// /// ``` /// use miden_core::{ - /// mast::{BasicBlockNodeBuilder, MastForest, MastForestContributor, MastForestWireView}, + /// mast::{BasicBlockNodeBuilder, DenseMastForestBuilder, MastForestWireView}, /// operations::Operation, /// serde::Serializable, /// }; /// - /// let mut forest = MastForest::new(); - /// let block_id = BasicBlockNodeBuilder::new(vec![Operation::Add]) - /// .add_to_forest(&mut forest) - /// .unwrap(); - /// forest.make_root(block_id); + /// let mut builder = DenseMastForestBuilder::new(); + /// let block_id = builder.push_node(BasicBlockNodeBuilder::new(vec![Operation::Add])).unwrap(); + /// builder.mark_root(block_id); + /// let forest = builder.finish().unwrap(); /// /// let mut bytes = Vec::new(); /// forest.write_into(&mut bytes); diff --git a/core/src/mast/serialization/resolved.rs b/core/src/mast/serialization/resolved.rs index 64f71d4647..87453d279a 100644 --- a/core/src/mast/serialization/resolved.rs +++ b/core/src/mast/serialization/resolved.rs @@ -8,11 +8,11 @@ use crate::{ Felt, chiplets::hasher, mast::{ - CallNode, DynNode, JoinNode, LoopNode, MastForestParts, SplitNode, + CallNode, DynNode, JoinNode, LoopNode, MastForestParts, MastNode, SplitNode, serialization::{basic_blocks::BasicBlockDataDecoder, layout::read_fixed_section_entry}, }, serde::{Deserializable, DeserializationError, SliceReader}, - utils::Idx, + utils::{Idx, IndexVec}, }; /// Digest sources for a parsed serialized forest. @@ -137,7 +137,7 @@ impl<'a> ResolvedSerializedForest<'a> { self.layout.basic_block_offset(), self.layout.basic_block_len(), )?); - let mut mast_forest = MastForest::new(); + let mut nodes = IndexVec::::with_capacity(self.node_count()); for index in 0..self.node_count() { let entry = self.node_entry_at(index)?; @@ -148,11 +148,16 @@ impl<'a> ResolvedSerializedForest<'a> { &basic_block_data_decoder, digest, )?; - mast_node_builder.add_to_forest_relaxed(&mut mast_forest).map_err(|e| { + let node = mast_node_builder.build_linked().map_err(|e| { DeserializationError::InvalidValue(format!( - "failed to add node to MAST forest while deserializing: {e}", + "failed to build node while deserializing MAST forest: {e}", )) })?; + nodes.push(node).map_err(|_| { + DeserializationError::InvalidValue( + "too many nodes while deserializing MAST forest".into(), + ) + })?; } let mut roots = Vec::with_capacity(self.procedure_root_count()); @@ -160,16 +165,12 @@ impl<'a> ResolvedSerializedForest<'a> { roots.push(self.procedure_root_at(index)?); } - MastForest::from_trusted_deserialization_parts(MastForestParts { - nodes: mast_forest.nodes, - roots, - advice_map, - }) - .map_err(|e| { - DeserializationError::InvalidValue(format!( - "failed to construct trusted deserialized MAST forest: {e}", - )) - }) + MastForest::from_trusted_deserialization_parts(MastForestParts { nodes, roots, advice_map }) + .map_err(|e| { + DeserializationError::InvalidValue(format!( + "failed to construct trusted deserialized MAST forest: {e}", + )) + }) } pub(super) fn node_count(&self) -> usize { @@ -474,12 +475,12 @@ fn checked_child_index( } pub(super) fn basic_block_offset_for_node_index( - nodes: &[super::MastNode], + nodes: &[MastNode], node_index: usize, ) -> Result { let mut offset = 0usize; for node in nodes.iter().take(node_index) { - if let super::MastNode::Block(block) = node { + if let MastNode::Block(block) = node { offset = offset.checked_add(basic_block_data_len(block)).ok_or_else(|| { DeserializationError::InvalidValue("basic-block data offset overflow".to_string()) })?; diff --git a/core/src/mast/serialization/seed_gen.rs b/core/src/mast/serialization/seed_gen.rs index d991d18abe..8ced05172c 100644 --- a/core/src/mast/serialization/seed_gen.rs +++ b/core/src/mast/serialization/seed_gen.rs @@ -9,7 +9,7 @@ use crate::{ Felt, Word, advice::{AdviceInputs, AdviceMap}, events::EventId, - mast::{BasicBlockNodeBuilder, JoinNodeBuilder, MastForest, MastForestContributor}, + mast::{BasicBlockNodeBuilder, JoinNodeBuilder, MastForest}, operations::Operation, precompile::PrecompileRequest, program::{Kernel, Program, StackInputs, StackOutputs}, diff --git a/core/src/mast/serialization/tests.rs b/core/src/mast/serialization/tests.rs index 13aa308b61..46ae2323f9 100644 --- a/core/src/mast/serialization/tests.rs +++ b/core/src/mast/serialization/tests.rs @@ -10,11 +10,10 @@ use crate::{ chiplets::hasher, mast::{ BasicBlockNodeBuilder, CallNodeBuilder, DynNodeBuilder, ExecutableMastForest, - ExternalNodeBuilder, JoinNodeBuilder, LoopNodeBuilder, MastForestContributor, - MastForestError, MastForestView, MastNodeExt, MastNodeId, OP_BATCH_SIZE, OpBatch, - SparseMastForest, SparseMastForestBuilder, SparseMastForestReadOptions, SplitNodeBuilder, - UntrustedMastForest, UntrustedMastForestReadOptions, VisitKind, - compute_mast_forest_commitment_from_parts, + ExternalNodeBuilder, JoinNodeBuilder, LoopNodeBuilder, MastForestError, MastForestView, + MastNodeExt, MastNodeId, OP_BATCH_SIZE, OpBatch, SparseMastForest, SparseMastForestBuilder, + SparseMastForestReadOptions, SplitNodeBuilder, UntrustedMastForest, + UntrustedMastForestReadOptions, VisitKind, compute_mast_forest_commitment_from_parts, }, operations::Operation, serde::{ByteReader, Deserializable, DeserializationError, Serializable, SliceReader}, @@ -902,6 +901,11 @@ fn sparse_mast_round_trip_preserves_external_full_node() { let external = ExternalNodeBuilder::new(external_digest).add_to_forest(&mut forest).unwrap(); forest.make_root(external); + let (forest, remapping) = + MastForest::from_raw_parts_with_id_map(forest.nodes, forest.roots, forest.advice_map) + .unwrap(); + let unvisited = remapping.get(unvisited).unwrap(); + let external = remapping.get(external).unwrap(); let forest = Arc::new(forest); let mut builder = SparseMastForestBuilder::new(Arc::clone(&forest)); builder.record_visit(external, VisitKind::FullVisit); @@ -995,6 +999,12 @@ fn sparse_mast_round_trip_preserves_sorted_commitment_inputs() { forest.make_root(low_id); forest.make_root(middle_id); + let (forest, remapping) = + MastForest::from_raw_parts_with_id_map(forest.nodes, forest.roots, forest.advice_map) + .unwrap(); + let high_id = remapping.get(high_id).unwrap(); + let low_id = remapping.get(low_id).unwrap(); + let middle_id = remapping.get(middle_id).unwrap(); let forest = Arc::new(forest); let mut builder = SparseMastForestBuilder::new(Arc::clone(&forest)); builder.record_visit(high_id, VisitKind::FullVisit); diff --git a/core/src/mast/tests.rs b/core/src/mast/tests.rs index f7c614d80d..50c9825140 100644 --- a/core/src/mast/tests.rs +++ b/core/src/mast/tests.rs @@ -9,8 +9,7 @@ use crate::{ chiplets::hasher, mast::{ BasicBlockNodeBuilder, CallNodeBuilder, DynNode, DynNodeBuilder, ExternalNodeBuilder, - JoinNodeBuilder, MastForest, MastForestContributor, MastForestError, MastNodeExt, - MastNodeId, + JoinNodeBuilder, MastForest, MastForestError, MastNodeExt, MastNodeId, }, operations::Operation, program::{Kernel, ProgramInfo}, @@ -166,6 +165,9 @@ fn mast_forest_commitment_separates_interface_and_dependencies() { .add_to_forest(&mut first) .unwrap(); + let (first, _) = + MastForest::from_raw_parts_with_id_map(first.nodes, first.roots, first.advice_map).unwrap(); + let mut second = MastForest::new(); let second_root = BasicBlockNodeBuilder::new(vec![Operation::Add]) .add_to_forest(&mut second) @@ -179,6 +181,9 @@ fn mast_forest_commitment_separates_interface_and_dependencies() { ])) .add_to_forest(&mut second) .unwrap(); + let (second, _) = + MastForest::from_raw_parts_with_id_map(second.nodes, second.roots, second.advice_map) + .unwrap(); assert_eq!(first.interface_commitment(), second.interface_commitment()); assert_ne!(first.dependency_commitment(), second.dependency_commitment()); diff --git a/crates/assembly/src/mast_forest_builder/finalizer.rs b/crates/assembly/src/mast_forest_builder/finalizer.rs index 53e7c19025..99bff7eaba 100644 --- a/crates/assembly/src/mast_forest_builder/finalizer.rs +++ b/crates/assembly/src/mast_forest_builder/finalizer.rs @@ -182,33 +182,48 @@ impl MastForestFinalizer { debug_vars: &IndexVec, advice_map: AdviceMap, ) -> Result { + let MastForestFinalizer { nodes, mut node_id_by_ref } = self; + let mut roots = Vec::with_capacity(procedure_root_refs.len()); for &root_ref in procedure_root_refs { - let root_id = *self.node_id_by_ref.get(&root_ref).ok_or_else(|| { + let root_id = *node_id_by_ref.get(&root_ref).ok_or_else(|| { Report::new(MastForestBuilderError::MissingProcedureRoot { root_ref }) })?; roots.push(root_id); } - let (source_graph, source_id_by_ref) = self.finalize_source_graph( + let (mast_forest, final_id_remapping) = + MastForest::from_raw_parts_with_id_map(nodes, roots, advice_map) + .map_err(|source| Report::new(MastForestBuilderError::FinalizeForest { source }))?; + + for node_id in node_id_by_ref.values_mut() { + *node_id = final_id_remapping.get(*node_id).ok_or_else(|| { + Report::new(MastForestBuilderError::FinalizeForest { + source: MastForestError::NodeIdOverflow(*node_id, final_id_remapping.len()), + }) + })?; + } + + let (source_graph, source_id_by_ref) = Self::finalize_source_graph( + &mast_forest, + &node_id_by_ref, procedure_source_root_refs, source_nodes, asm_op_by_ref, debug_vars, )?; - let mast_forest = MastForest::from_raw_parts(self.nodes, roots, advice_map) - .map_err(|source| Report::new(MastForestBuilderError::FinalizeForest { source }))?; Ok(BuiltMastForest { mast_forest, source_graph, - node_id_by_ref: self.node_id_by_ref, + node_id_by_ref, source_id_by_ref, }) } fn finalize_source_graph( - &self, + mast_forest: &MastForest, + node_id_by_ref: &BTreeMap, procedure_source_root_refs: &[SourceNodeRef], source_nodes: &IndexVec, asm_op_by_ref: &IndexVec, @@ -219,7 +234,7 @@ impl MastForestFinalizer { .iter() .enumerate() .filter_map(|(index, source_node)| { - self.node_id_by_ref + node_id_by_ref .contains_key(&source_node.exec_ref) .then_some(SourceNodeRef::from(index as u32)) }) @@ -234,7 +249,7 @@ impl MastForestFinalizer { for &source_ref in &live_source_refs { let pending_source_node = &source_nodes[source_ref]; let exec_node = - *self.node_id_by_ref.get(&pending_source_node.exec_ref).ok_or_else(|| { + *node_id_by_ref.get(&pending_source_node.exec_ref).ok_or_else(|| { Report::new(MastForestBuilderError::MissingFinalSourceExec { source_ref, exec_ref: pending_source_node.exec_ref, @@ -252,7 +267,7 @@ impl MastForestFinalizer { }) }) .collect::, Report>>()?; - let node = &self.nodes[exec_node]; + let node = &mast_forest[exec_node]; let (_, asm_op_refs) = compute_operations_and_adjust_mappings(node, pending_source_node.asm_ops.clone()); let asm_ops = asm_op_refs diff --git a/crates/assembly/src/project/tests.rs b/crates/assembly/src/project/tests.rs index 6387d4bf91..3b9b005f71 100644 --- a/crates/assembly/src/project/tests.rs +++ b/crates/assembly/src/project/tests.rs @@ -2,7 +2,7 @@ use std::{path::Path, process::Command, string::String, sync::Arc}; use miden_assembly_syntax::source_file; use miden_core::{ - mast::{BasicBlockNodeBuilder, MastForest, MastForestContributor, MastNodeExt}, + mast::{BasicBlockNodeBuilder, MastForest, MastNodeExt}, operations::{DebugVarInfo, DebugVarLocation, Operation}, serde::{Deserializable, Serializable, SliceReader}, utils::hash_string_to_word, diff --git a/crates/mast-package/src/debug_info/types.rs b/crates/mast-package/src/debug_info/types.rs index e15af6b0f4..ef4d910127 100644 --- a/crates/mast-package/src/debug_info/types.rs +++ b/crates/mast-package/src/debug_info/types.rs @@ -1854,16 +1854,18 @@ mod tests { #[test] fn test_package_source_debug_merge_remaps_execution_nodes_without_collapsing_sources() { use miden_core::{ - mast::{BasicBlockNodeBuilder, MastForest, MastForestContributor}, + mast::{BasicBlockNodeBuilder, DenseMastForestBuilder, MastForest}, operations::{DebugVarInfo, DebugVarLocation, Operation}, }; fn forest_with_add_block() -> (MastForest, MastNodeId) { - let mut forest = MastForest::new(); - let block = BasicBlockNodeBuilder::new(alloc::vec![Operation::Add]) - .add_to_forest(&mut forest) + let mut builder = DenseMastForestBuilder::new(); + let block = builder + .push_node(BasicBlockNodeBuilder::new(alloc::vec![Operation::Add])) .unwrap(); - forest.make_root(block); + builder.mark_root(block); + let (forest, remapping) = builder.finish_with_id_map().unwrap(); + let block = remapping.get(block).unwrap(); (forest, block) } @@ -1989,16 +1991,19 @@ mod tests { #[test] fn test_package_source_debug_merge_remaps_non_root_execution_nodes() { use miden_core::{ - mast::{BasicBlockNodeBuilder, CallNodeBuilder, MastForest, MastForestContributor}, + mast::{BasicBlockNodeBuilder, CallNodeBuilder, DenseMastForestBuilder, MastForest}, operations::Operation, }; - let mut forest = MastForest::new(); - let callee = BasicBlockNodeBuilder::new(alloc::vec![Operation::Add]) - .add_to_forest(&mut forest) + let mut builder = DenseMastForestBuilder::new(); + let callee = builder + .push_node(BasicBlockNodeBuilder::new(alloc::vec![Operation::Add])) .unwrap(); - let call = CallNodeBuilder::new(callee).add_to_forest(&mut forest).unwrap(); - forest.make_root(call); + let call = builder.push_node(CallNodeBuilder::new(callee)).unwrap(); + builder.mark_root(call); + let (forest, remapping) = builder.finish_with_id_map().unwrap(); + let callee = remapping.get(callee).unwrap(); + let call = remapping.get(call).unwrap(); let root_source = DebugSourceNodeId::from(0); let child_source = DebugSourceNodeId::from(1); diff --git a/crates/mast-package/src/package/arbitrary.rs b/crates/mast-package/src/package/arbitrary.rs index fd469699e4..6f50d4a2f3 100644 --- a/crates/mast-package/src/package/arbitrary.rs +++ b/crates/mast-package/src/package/arbitrary.rs @@ -50,7 +50,7 @@ impl proptest::arbitrary::Arbitrary for Package { fn arbitrary_with(params: Self::Parameters) -> Self::Strategy { use miden_core::{ Felt, - mast::{BasicBlockNodeBuilder, MastForestContributor, MastNodeExt}, + mast::{BasicBlockNodeBuilder, DenseMastForestBuilder, MastNodeExt}, operations::Operation, }; use proptest::prelude::*; @@ -72,50 +72,72 @@ impl proptest::arbitrary::Arbitrary for Package { } // Create a MastForest with actual nodes for the exports - let mut mast_forest = Box::new(MastForest::new()); + let mut mast_forest_builder = DenseMastForestBuilder::new(); let mut nodes = Vec::with_capacity(exports.len()); for export in exports.iter_mut() { if let PackageExport::Procedure(export) = export { let procedure_index = nodes.len() as u64; - let node_id = BasicBlockNodeBuilder::new(vec![ - Operation::Push(Felt::new_unchecked(procedure_index)), - Operation::Add, - Operation::Mul, - ]) - .add_to_forest(&mut mast_forest) - .unwrap(); + let node_id = mast_forest_builder + .push_node_builder( + BasicBlockNodeBuilder::new(vec![ + Operation::Push(Felt::new_unchecked(procedure_index)), + Operation::Add, + Operation::Mul, + ]) + .into(), + ) + .unwrap(); // Add the node to the forest roots if it's not already there - mast_forest.make_root(node_id); + mast_forest_builder.mark_root(node_id); nodes.push(node_id); export.node = Some(node_id); - export.digest = mast_forest[node_id].digest(); + export.digest = + mast_forest_builder.get_node_by_id(node_id).unwrap().digest(); } } // Generate an entrypoint export if needed if kind.is_executable() { let procedure_index = nodes.len() as u64; - let node_id = BasicBlockNodeBuilder::new(vec![ - Operation::Push(Felt::new_unchecked(procedure_index)), - Operation::Add, - Operation::Mul, - ]) - .add_to_forest(&mut mast_forest) - .unwrap(); + let node_id = mast_forest_builder + .push_node_builder( + BasicBlockNodeBuilder::new(vec![ + Operation::Push(Felt::new_unchecked(procedure_index)), + Operation::Add, + Operation::Mul, + ]) + .into(), + ) + .unwrap(); // Add the node to the forest roots if it's not already there - mast_forest.make_root(node_id); + mast_forest_builder.mark_root(node_id); nodes.push(node_id); let path: Arc = Path::EXEC.join(ast::ProcedureName::MAIN_PROC_NAME).into(); exports.push(PackageExport::Procedure(ProcedureExport::new( path, Some(node_id), - mast_forest[node_id].digest(), + mast_forest_builder.get_node_by_id(node_id).unwrap().digest(), None, ))); } - let mast_forest = Arc::from(mast_forest); + let (mast_forest, id_remapping) = mast_forest_builder + .finish_with_id_map() + .expect("generated MAST forest should be valid"); + for export in exports.iter_mut() { + if let PackageExport::Procedure(export) = export + && let Some(builder_node_id) = export.node + { + let node_id = id_remapping + .get(builder_node_id) + .expect("procedure export should map to a finalized node"); + export.node = Some(node_id); + export.digest = mast_forest[node_id].digest(); + } + } + + let mast_forest = Arc::new(mast_forest); Package::create( name.clone(), version.clone(), diff --git a/crates/mast-package/src/package/mod.rs b/crates/mast-package/src/package/mod.rs index edb95dd21a..61666bc4f4 100644 --- a/crates/mast-package/src/package/mod.rs +++ b/crates/mast-package/src/package/mod.rs @@ -1323,7 +1323,7 @@ mod tests { use miden_core::{ advice::AdviceMap, mast::{ - BasicBlockNodeBuilder, ExternalNodeBuilder, MastForest, MastForestContributor, + BasicBlockNodeBuilder, DenseMastForestBuilder, ExternalNodeBuilder, MastForest, MastNode, MastNodeExt, MastNodeId, SplitNodeBuilder, }, operations::Operation, @@ -1343,26 +1343,32 @@ mod tests { }; fn build_forest() -> (MastForest, MastNodeId) { - let mut forest = MastForest::new(); - let node_id = BasicBlockNodeBuilder::new(vec![Operation::Add]) - .add_to_forest(&mut forest) + let mut builder = DenseMastForestBuilder::new(); + let node_id = builder + .push_node(BasicBlockNodeBuilder::new(vec![Operation::Add])) .expect("failed to build basic block"); - forest.make_root(node_id); + builder.mark_root(node_id); + let (forest, remapping) = builder.finish_with_id_map().expect("failed to build forest"); + let node_id = remapping.get(node_id).expect("root node should be retained"); (forest, node_id) } fn build_split_forest() -> (MastForest, MastNodeId, MastNodeId, MastNodeId) { - let mut forest = MastForest::new(); - let left_id = BasicBlockNodeBuilder::new(vec![Operation::Add]) - .add_to_forest(&mut forest) + let mut builder = DenseMastForestBuilder::new(); + let left_id = builder + .push_node(BasicBlockNodeBuilder::new(vec![Operation::Add])) .expect("failed to build left basic block"); - let right_id = BasicBlockNodeBuilder::new(vec![Operation::Mul]) - .add_to_forest(&mut forest) + let right_id = builder + .push_node(BasicBlockNodeBuilder::new(vec![Operation::Mul])) .expect("failed to build right basic block"); - let root_id = SplitNodeBuilder::new([left_id, right_id]) - .add_to_forest(&mut forest) + let root_id = builder + .push_node(SplitNodeBuilder::new([left_id, right_id])) .expect("failed to build split node"); - forest.make_root(root_id); + builder.mark_root(root_id); + let (forest, remapping) = builder.finish_with_id_map().expect("failed to build forest"); + let root_id = remapping.get(root_id).expect("root node should be retained"); + let left_id = remapping.get(left_id).expect("left node should be retained"); + let right_id = remapping.get(right_id).expect("right node should be retained"); (forest, root_id, left_id, right_id) } @@ -2109,18 +2115,23 @@ mod tests { } } - let mut concrete_forest = MastForest::new(); - let concrete_root = BasicBlockNodeBuilder::new(vec![Operation::Add]) - .add_to_forest(&mut concrete_forest) + let mut concrete_builder = DenseMastForestBuilder::new(); + let concrete_root = concrete_builder + .push_node(BasicBlockNodeBuilder::new(vec![Operation::Add])) .unwrap(); - concrete_forest.make_root(concrete_root); + concrete_builder.mark_root(concrete_root); + let (concrete_forest, concrete_remapping) = concrete_builder.finish_with_id_map().unwrap(); + let concrete_root = concrete_remapping.get(concrete_root).unwrap(); let concrete_digest = concrete_forest[concrete_root].digest(); - let mut placeholder_forest = MastForest::new(); - let placeholder_root = ExternalNodeBuilder::new(concrete_digest) - .add_to_forest(&mut placeholder_forest) + let mut placeholder_builder = DenseMastForestBuilder::new(); + let placeholder_root = placeholder_builder + .push_node(ExternalNodeBuilder::new(concrete_digest)) .unwrap(); - placeholder_forest.make_root(placeholder_root); + placeholder_builder.mark_root(placeholder_root); + let (placeholder_forest, placeholder_remapping) = + placeholder_builder.finish_with_id_map().unwrap(); + let placeholder_root = placeholder_remapping.get(placeholder_root).unwrap(); let placeholder_debug = debug_info_for_root(placeholder_root, "placeholder"); let concrete_debug = debug_info_for_root(concrete_root, "concrete"); diff --git a/crates/mast-package/src/package/seed_gen.rs b/crates/mast-package/src/package/seed_gen.rs index c89e90cc9f..1c22aa693b 100644 --- a/crates/mast-package/src/package/seed_gen.rs +++ b/crates/mast-package/src/package/seed_gen.rs @@ -9,7 +9,7 @@ use miden_assembly_syntax::{ semver::Version, }; use miden_core::{ - mast::{BasicBlockNodeBuilder, MastForest, MastForestContributor, MastNodeExt, MastNodeId}, + mast::{BasicBlockNodeBuilder, DenseMastForestBuilder, MastForest, MastNodeExt, MastNodeId}, operations::Operation, serde::Serializable, }; @@ -18,11 +18,15 @@ use super::{PackageId, TargetType}; use crate::{Package, PackageExport, ProcedureExport}; fn build_forest() -> (MastForest, MastNodeId) { - let mut forest = MastForest::new(); - let node_id = BasicBlockNodeBuilder::new(vec![Operation::Add]) - .add_to_forest(&mut forest) + let mut builder = DenseMastForestBuilder::new(); + let builder_node_id = builder + .push_node_builder(BasicBlockNodeBuilder::new(vec![Operation::Add]).into()) .expect("failed to build basic block"); - forest.make_root(node_id); + builder.mark_root(builder_node_id); + let (forest, id_remapping) = builder.finish_with_id_map().expect("seed forest should be valid"); + let node_id = id_remapping + .get(builder_node_id) + .expect("seed root should map to a finalized node"); (forest, node_id) } diff --git a/crates/mast-package/src/package/serialization.rs b/crates/mast-package/src/package/serialization.rs index 8e99a4472e..2396fc71c7 100644 --- a/crates/mast-package/src/package/serialization.rs +++ b/crates/mast-package/src/package/serialization.rs @@ -1190,7 +1190,7 @@ mod tests { Felt, Word, advice::AdviceMap, mast::{ - BasicBlockNodeBuilder, MastForest, MastForestContributor, MastNode, MastNodeExt, + BasicBlockNodeBuilder, DenseMastForestBuilder, MastForest, MastNode, MastNodeExt, MastNodeId, }, operations::Operation, @@ -1214,15 +1214,26 @@ mod tests { }, }; - fn build_forest() -> (MastForest, MastNodeId) { - let mut forest = MastForest::new(); - let node_id = BasicBlockNodeBuilder::new(vec![Operation::Add]) - .add_to_forest(&mut forest) + fn build_single_node_forest( + operations: Vec, + make_root: bool, + ) -> (MastForest, MastNodeId) { + let mut builder = DenseMastForestBuilder::new(); + let node_id = builder + .push_node(BasicBlockNodeBuilder::new(operations)) .expect("failed to build basic block"); - forest.make_root(node_id); + if make_root { + builder.mark_root(node_id); + } + let (forest, remapping) = builder.finish_with_id_map().expect("forest should be valid"); + let node_id = remapping.get(node_id).expect("node should be retained"); (forest, node_id) } + fn build_forest() -> (MastForest, MastNodeId) { + build_single_node_forest(vec![Operation::Add], true) + } + fn absolute_path(name: &str) -> Arc { let path = PathBuf::new(name).expect("invalid path"); let path = path.as_path().to_absolute().unwrap().into_owned(); @@ -1777,10 +1788,7 @@ mod tests { /// procedure root in the underlying MAST forest (issue #2831). #[test] fn package_rejects_non_root_export() { - let mut forest = MastForest::new(); - let node_id = BasicBlockNodeBuilder::new(vec![Operation::Add]) - .add_to_forest(&mut forest) - .expect("failed to build basic block"); + let (forest, node_id) = build_single_node_forest(vec![Operation::Add], false); let digest = forest[node_id].digest(); let path = absolute_path("test::proc"); @@ -2021,10 +2029,8 @@ mod tests { // pub proc p // push.1 // end - let mut forest = MastForest::new(); - let node_id = BasicBlockNodeBuilder::new(vec![Operation::Push(Felt::from_u32(1))]) - .add_to_forest(&mut forest) - .expect("failed to build basic block"); + let (forest, node_id) = + build_single_node_forest(vec![Operation::Push(Felt::from_u32(1))], false); let digest = forest[node_id].digest(); let path = absolute_path("lib::p"); @@ -2068,10 +2074,8 @@ mod tests { // pub proc p // push.1 // end - let mut forest = MastForest::new(); - let node_id = BasicBlockNodeBuilder::new(vec![Operation::Push(Felt::from_u32(1))]) - .add_to_forest(&mut forest) - .expect("failed to build basic block"); + let (forest, node_id) = + build_single_node_forest(vec![Operation::Push(Felt::from_u32(1))], false); let digest = forest[node_id].digest(); let path = absolute_path("lib::p"); @@ -2112,10 +2116,8 @@ mod tests { // pub proc k1 // push.1 // end - let mut forest = MastForest::new(); - let node_id = BasicBlockNodeBuilder::new(vec![Operation::Push(Felt::from_u32(1))]) - .add_to_forest(&mut forest) - .expect("failed to build basic block"); + let (forest, node_id) = + build_single_node_forest(vec![Operation::Push(Felt::from_u32(1))], false); let digest = forest[node_id].digest(); let path = absolute_path("$kernel::k1"); @@ -2161,10 +2163,8 @@ mod tests { // pub proc k1 // push.1 // end - let mut forest = MastForest::new(); - let node_id = BasicBlockNodeBuilder::new(vec![Operation::Push(Felt::from_u32(1))]) - .add_to_forest(&mut forest) - .expect("failed to build basic block"); + let (forest, node_id) = + build_single_node_forest(vec![Operation::Push(Felt::from_u32(1))], false); let digest = forest[node_id].digest(); let path = absolute_path("$kernel::k1"); @@ -2217,10 +2217,8 @@ mod tests { // pub proc k1 // push.1 // end - let mut forest = MastForest::new(); - let node_id = BasicBlockNodeBuilder::new(vec![Operation::Push(Felt::from_u32(1))]) - .add_to_forest(&mut forest) - .expect("failed to build basic block"); + let (forest, node_id) = + build_single_node_forest(vec![Operation::Push(Felt::from_u32(1))], false); let digest = forest[node_id].digest(); let path = absolute_path("$kernel::k1"); diff --git a/crates/package-registry/src/lib.rs b/crates/package-registry/src/lib.rs index 39199c463c..809c19493d 100644 --- a/crates/package-registry/src/lib.rs +++ b/crates/package-registry/src/lib.rs @@ -286,7 +286,9 @@ mod tests { use miden_assembly_syntax::ast::{Path as AstPath, PathBuf}; use miden_core::{ - mast::{BasicBlockNodeBuilder, MastForest, MastForestContributor, MastNodeExt, MastNodeId}, + mast::{ + BasicBlockNodeBuilder, DenseMastForestBuilder, MastForest, MastNodeExt, MastNodeId, + }, operations::Operation, }; use miden_mast_package::{Package, PackageExport, ProcedureExport, TargetType}; @@ -294,11 +296,13 @@ mod tests { use super::*; fn build_forest() -> (MastForest, MastNodeId) { - let mut forest = MastForest::new(); - let node_id = BasicBlockNodeBuilder::new(vec![Operation::Add]) - .add_to_forest(&mut forest) + let mut builder = DenseMastForestBuilder::new(); + let node_id = builder + .push_node(BasicBlockNodeBuilder::new(vec![Operation::Add])) .expect("failed to build basic block"); - forest.make_root(node_id); + builder.mark_root(node_id); + let (forest, remapping) = builder.finish_with_id_map().expect("failed to build forest"); + let node_id = remapping.get(node_id).expect("root node should be retained"); (forest, node_id) } diff --git a/crates/package-registry/src/resolver/index.rs b/crates/package-registry/src/resolver/index.rs index d906b954f9..b3d01d03df 100644 --- a/crates/package-registry/src/resolver/index.rs +++ b/crates/package-registry/src/resolver/index.rs @@ -240,7 +240,9 @@ mod tests { use miden_assembly_syntax::ast::{Path as AstPath, PathBuf}; use miden_core::{ - mast::{BasicBlockNodeBuilder, MastForest, MastForestContributor, MastNodeExt, MastNodeId}, + mast::{ + BasicBlockNodeBuilder, DenseMastForestBuilder, MastForest, MastNodeExt, MastNodeId, + }, operations::Operation, }; use miden_mast_package::{ @@ -250,11 +252,13 @@ mod tests { use super::*; fn build_forest() -> (MastForest, MastNodeId) { - let mut forest = MastForest::new(); - let node_id = BasicBlockNodeBuilder::new(vec![Operation::Add]) - .add_to_forest(&mut forest) + let mut builder = DenseMastForestBuilder::new(); + let node_id = builder + .push_node(BasicBlockNodeBuilder::new(vec![Operation::Add])) .expect("failed to build basic block"); - forest.make_root(node_id); + builder.mark_root(node_id); + let (forest, remapping) = builder.finish_with_id_map().expect("failed to build forest"); + let node_id = remapping.get(node_id).expect("root node should be retained"); (forest, node_id) } diff --git a/crates/utils-core-derive/src/lib.rs b/crates/utils-core-derive/src/lib.rs index 1d93bc7e7c..09e1037572 100644 --- a/crates/utils-core-derive/src/lib.rs +++ b/crates/utils-core-derive/src/lib.rs @@ -246,11 +246,9 @@ pub fn derive_mast_forest_contributor(input: TokenStream) -> TokenStream { // Extract variant information let variants: Vec<_> = enum_data.variants.iter().collect(); let variant_names: Vec<_> = variants.iter().map(|v| &v.ident).collect(); - let variant_fields: Vec<_> = variants.iter().map(|v| extract_single_field(v)).collect(); // Generate trait implementation by reading the trait definition - let trait_impl = - generate_mast_forest_contributor_impl(enum_name, generics, &variant_names, &variant_fields); + let trait_impl = generate_mast_forest_contributor_impl(enum_name, generics, &variant_names); TokenStream::from(trait_impl) } @@ -260,31 +258,16 @@ fn generate_mast_forest_contributor_impl( enum_name: &Ident, generics: &syn::Generics, variant_names: &[&Ident], - variant_fields: &[Ident], ) -> proc_macro2::TokenStream { - // For now, let's generate a simple implementation to test the macro - let add_to_forest_arms = - variant_names.iter().zip(variant_fields.iter()).map(|(variant, field)| { - quote! { - #enum_name::#variant(#field) => #field.add_to_forest(forest) - } - }); - quote! { impl #generics crate::mast::MastForestContributor for #enum_name #generics { - fn add_to_forest(self, forest: &mut crate::mast::MastForest) -> Result { - match self { - #(#add_to_forest_arms),* - } - } - fn fingerprint_for_node( &self, - forest: &crate::mast::MastForest, + context: &impl crate::mast::MastNodeContext, hash_by_node_id: &impl crate::utils::LookupByIdx, ) -> Result { match self { - #(#enum_name::#variant_names(field) => field.fingerprint_for_node(forest, hash_by_node_id)),* + #(#enum_name::#variant_names(field) => field.fingerprint_for_node(context, hash_by_node_id)),* } } diff --git a/miden-vm/tests/integration/operations/io_ops/env_ops.rs b/miden-vm/tests/integration/operations/io_ops/env_ops.rs index 705d2e1791..112b84e2a5 100644 --- a/miden-vm/tests/integration/operations/io_ops/env_ops.rs +++ b/miden-vm/tests/integration/operations/io_ops/env_ops.rs @@ -1,9 +1,6 @@ use miden_core::{ FMP_INIT_VALUE, - mast::{ - BasicBlockNodeBuilder, CallNodeBuilder, MastForest, MastForestContributor, MastNode, - MastNodeExt, - }, + mast::{BasicBlockNodeBuilder, CallNodeBuilder, MastForest, MastNode, MastNodeExt}, operations::Operation, }; use miden_utils_testing::{MIN_STACK_DEPTH, Word, build_op_test, build_test}; diff --git a/processor/src/execution/external.rs b/processor/src/execution/external.rs index aac0afa498..adc66abeef 100644 --- a/processor/src/execution/external.rs +++ b/processor/src/execution/external.rs @@ -115,7 +115,7 @@ mod tests { use miden_core::{ Felt, - mast::{BasicBlockNodeBuilder, ExternalNodeBuilder, MastForest, MastForestContributor}, + mast::{BasicBlockNodeBuilder, ExternalNodeBuilder, MastForest}, operations::Operation, program::Program, }; diff --git a/processor/src/fast/tests/mod.rs b/processor/src/fast/tests/mod.rs index 5956057f5b..0183f27abe 100644 --- a/processor/src/fast/tests/mod.rs +++ b/processor/src/fast/tests/mod.rs @@ -10,8 +10,8 @@ use miden_core::{ ONE, Word, events::SystemEvent, mast::{ - BasicBlockNodeBuilder, CallNodeBuilder, ExternalNodeBuilder, JoinNodeBuilder, - MastForestContributor, MastNodeExt, SplitNodeBuilder, + BasicBlockNodeBuilder, CallNodeBuilder, ExternalNodeBuilder, JoinNodeBuilder, MastNodeExt, + SplitNodeBuilder, }, operations::Operation, program::StackInputs, diff --git a/processor/src/tests/mod.rs b/processor/src/tests/mod.rs index a19c959604..02014b9d69 100644 --- a/processor/src/tests/mod.rs +++ b/processor/src/tests/mod.rs @@ -7,7 +7,7 @@ use miden_assembly::{ }; use miden_core::{ crypto::merkle::{MerkleStore, MerkleTree}, - mast::{BasicBlockNodeBuilder, MastForest, MastForestContributor, error_code_from_msg}, + mast::{BasicBlockNodeBuilder, MastForest, error_code_from_msg}, }; use miden_debug_types::{Location, SourceFile, SourceManager, SourceSpan}; use miden_utils_testing::crypto::{init_merkle_leaves, init_merkle_store}; diff --git a/processor/src/trace/chiplets/tests.rs b/processor/src/trace/chiplets/tests.rs index 0636169c6c..7aded0f46e 100644 --- a/processor/src/trace/chiplets/tests.rs +++ b/processor/src/trace/chiplets/tests.rs @@ -11,7 +11,7 @@ use miden_air::trace::{ }; use miden_core::{ Felt, ONE, Word, ZERO, - mast::{BasicBlockNodeBuilder, CallNodeBuilder, MastForest, MastForestContributor}, + mast::{BasicBlockNodeBuilder, CallNodeBuilder, MastForest}, program::{Program, StackInputs}, }; diff --git a/processor/src/trace/parallel/core_trace_fragment/tests.rs b/processor/src/trace/parallel/core_trace_fragment/tests.rs index ff911a5302..408e1bc6b9 100644 --- a/processor/src/trace/parallel/core_trace_fragment/tests.rs +++ b/processor/src/trace/parallel/core_trace_fragment/tests.rs @@ -27,7 +27,7 @@ use miden_core::{ events::EventName, mast::{ BasicBlockNodeBuilder, CallNodeBuilder, DynNodeBuilder, JoinNodeBuilder, LoopNodeBuilder, - MastForest, MastForestContributor, MastNodeExt, OP_BATCH_SIZE, SplitNodeBuilder, + MastForest, MastNodeExt, OP_BATCH_SIZE, SplitNodeBuilder, }, operations::{Operation, opcodes}, program::{Kernel, Program, StackInputs}, diff --git a/processor/src/trace/parallel/tests.rs b/processor/src/trace/parallel/tests.rs index 15707683f2..52535ba26b 100644 --- a/processor/src/trace/parallel/tests.rs +++ b/processor/src/trace/parallel/tests.rs @@ -11,8 +11,8 @@ use miden_core::{ field::QuadFelt, mast::{ BasicBlockNodeBuilder, CallNodeBuilder, DynNodeBuilder, ExternalNodeBuilder, - JoinNodeBuilder, LoopNodeBuilder, MastForest, MastForestContributor, MastForestId, - MastNodeExt, MastNodeId, SplitNodeBuilder, + JoinNodeBuilder, LoopNodeBuilder, MastForest, MastForestId, MastNodeExt, MastNodeId, + SplitNodeBuilder, }, operations::{Operation, opcodes}, precompile::PrecompileRequest, diff --git a/processor/src/trace/parallel/tracer/mod.rs b/processor/src/trace/parallel/tracer/mod.rs index 143e90a5ca..a09424f2d6 100644 --- a/processor/src/trace/parallel/tracer/mod.rs +++ b/processor/src/trace/parallel/tracer/mod.rs @@ -660,10 +660,7 @@ where mod tests { use alloc::vec; - use miden_core::{ - mast::{DynNodeBuilder, MastForestContributor}, - precompile::PrecompileTranscriptState, - }; + use miden_core::{mast::DynNodeBuilder, precompile::PrecompileTranscriptState}; use super::*; use crate::{ diff --git a/processor/src/trace/tests/chiplets/hasher.rs b/processor/src/trace/tests/chiplets/hasher.rs index 6a0709cec8..08571bfadc 100644 --- a/processor/src/trace/tests/chiplets/hasher.rs +++ b/processor/src/trace/tests/chiplets/hasher.rs @@ -28,7 +28,7 @@ use miden_air::{ use miden_core::{ Felt, ONE, Word, ZERO, crypto::merkle::{MerkleStore, MerkleTree, NodeIndex}, - mast::{BasicBlockNodeBuilder, MastForest, MastForestContributor, SplitNodeBuilder}, + mast::{BasicBlockNodeBuilder, MastForest, SplitNodeBuilder}, operations::{Operation, opcodes}, program::Program, }; diff --git a/processor/src/trace/tests/decoder.rs b/processor/src/trace/tests/decoder.rs index dbfdda5dc2..d88595d596 100644 --- a/processor/src/trace/tests/decoder.rs +++ b/processor/src/trace/tests/decoder.rs @@ -19,7 +19,7 @@ use miden_core::{ Felt, ONE, ZERO, mast::{ BasicBlockNodeBuilder, CallNodeBuilder, JoinNodeBuilder, LoopNodeBuilder, MastForest, - MastForestContributor, MastNodeExt, SplitNodeBuilder, + MastNodeExt, SplitNodeBuilder, }, operations::{Operation, opcodes}, program::Program, @@ -824,11 +824,7 @@ fn op_group_span_two_batch_transition_inserts( fn decoder_dyncall_at_min_stack_depth_records_post_drop_ctx_info() { use std::sync::Arc; - use crate::{ - MIN_STACK_DEPTH, - mast::{DynNodeBuilder, MastForestContributor}, - operation::opcodes, - }; + use crate::{MIN_STACK_DEPTH, mast::DynNodeBuilder, operation::opcodes}; // Build exactly the same program shape as `dyncall_program()` in parallel/tests.rs: // join( @@ -913,10 +909,7 @@ fn decoder_dyncall_with_multiple_overflow_entries_records_correct_overflow_addr( // the second-to-last entry), not the pre-pop address (the clock of the top entry). use std::sync::Arc; - use crate::{ - mast::{DynNodeBuilder, MastForestContributor}, - operation::opcodes, - }; + use crate::{mast::DynNodeBuilder, operation::opcodes}; const HASH_ADDR: Felt = Felt::new_unchecked(40); diff --git a/processor/src/trace/tests/mod.rs b/processor/src/trace/tests/mod.rs index 92f0ac08b8..b21d073383 100644 --- a/processor/src/trace/tests/mod.rs +++ b/processor/src/trace/tests/mod.rs @@ -1,7 +1,7 @@ use alloc::vec::Vec; use miden_core::{ - mast::{BasicBlockNodeBuilder, MastForest, MastForestContributor}, + mast::{BasicBlockNodeBuilder, MastForest}, operations::Operation, program::Program, }; diff --git a/processor/tests/eval_circuit_overflow.rs b/processor/tests/eval_circuit_overflow.rs index 67d3b3ed4d..70697d5768 100644 --- a/processor/tests/eval_circuit_overflow.rs +++ b/processor/tests/eval_circuit_overflow.rs @@ -1,6 +1,6 @@ use miden_core::{ field::PrimeField64, - mast::{BasicBlockNodeBuilder, MastForest, MastForestContributor}, + mast::{BasicBlockNodeBuilder, MastForest}, }; use miden_processor::{ AceError, DefaultHost, ExecutionError, FastProcessor, Felt, Program, StackInputs,