miden_core/mast/
mod.rs

1use alloc::{
2    collections::{BTreeMap, BTreeSet},
3    sync::Arc,
4    vec::Vec,
5};
6use core::{
7    fmt,
8    ops::{Index, IndexMut},
9};
10
11pub use miden_utils_indexing::{IndexVec, IndexedVecError};
12#[cfg(feature = "serde")]
13use serde::{Deserialize, Serialize};
14
15mod node;
16#[cfg(any(test, feature = "arbitrary"))]
17pub use node::arbitrary;
18pub use node::{
19    BasicBlockNode, BasicBlockNodeBuilder, CallNode, CallNodeBuilder, DecoratedOpLink,
20    DecoratorOpLinkIterator, DecoratorStore, DynNode, DynNodeBuilder, ExternalNode,
21    ExternalNodeBuilder, JoinNode, JoinNodeBuilder, LoopNode, LoopNodeBuilder,
22    MastForestContributor, MastNode, MastNodeBuilder, MastNodeErrorContext, MastNodeExt,
23    OP_BATCH_SIZE, OP_GROUP_SIZE, OpBatch, OperationOrDecorator, SplitNode, SplitNodeBuilder,
24};
25
26use crate::{
27    AdviceMap, Decorator, Felt, Idx, LexicographicWord, Word,
28    crypto::hash::Hasher,
29    utils::{ByteWriter, DeserializationError, Serializable, hash_string_to_word},
30};
31
32mod debuginfo;
33pub use debuginfo::{
34    DebugInfo, DecoratedLinks, DecoratedLinksIter, DecoratorIndexError, NodeToDecoratorIds,
35    OpToDecoratorIds,
36};
37
38mod serialization;
39
40mod merger;
41pub(crate) use merger::MastForestMerger;
42pub use merger::MastForestRootMap;
43
44mod multi_forest_node_iterator;
45pub(crate) use multi_forest_node_iterator::*;
46
47mod node_fingerprint;
48pub use node_fingerprint::{DecoratorFingerprint, MastNodeFingerprint};
49
50mod node_builder_utils;
51pub use node_builder_utils::build_node_with_remapped_ids;
52
53#[cfg(test)]
54mod tests;
55
56// MAST FOREST
57// ================================================================================================
58
59/// Represents one or more procedures, represented as a collection of [`MastNode`]s.
60///
61/// A [`MastForest`] does not have an entrypoint, and hence is not executable. A [`crate::Program`]
62/// can be built from a [`MastForest`] to specify an entrypoint.
63#[derive(Clone, Debug, Default, PartialEq, Eq)]
64pub struct MastForest {
65    /// All of the nodes local to the trees comprising the MAST forest.
66    nodes: IndexVec<MastNodeId, MastNode>,
67
68    /// Roots of procedures defined within this MAST forest.
69    roots: Vec<MastNodeId>,
70
71    /// Advice map to be loaded into the VM prior to executing procedures from this MAST forest.
72    advice_map: AdviceMap,
73
74    /// Debug information including decorators and error codes.
75    /// Always present (as per issue #1821), but can be empty for stripped builds.
76    debug_info: DebugInfo,
77}
78
79// ------------------------------------------------------------------------------------------------
80/// Constructors
81impl MastForest {
82    /// Creates a new empty [`MastForest`].
83    pub fn new() -> Self {
84        Self {
85            nodes: IndexVec::new(),
86            roots: Vec::new(),
87            advice_map: AdviceMap::default(),
88            debug_info: DebugInfo::new(),
89        }
90    }
91}
92
93// ------------------------------------------------------------------------------------------------
94/// State mutators
95impl MastForest {
96    /// The maximum number of nodes that can be stored in a single MAST forest.
97    const MAX_NODES: usize = (1 << 30) - 1;
98
99    /// Marks the given [`MastNodeId`] as being the root of a procedure.
100    ///
101    /// If the specified node is already marked as a root, this will have no effect.
102    ///
103    /// # Panics
104    /// - if `new_root_id`'s internal index is larger than the number of nodes in this forest (i.e.
105    ///   clearly doesn't belong to this MAST forest).
106    pub fn make_root(&mut self, new_root_id: MastNodeId) {
107        assert!(new_root_id.to_usize() < self.nodes.len());
108
109        if !self.roots.contains(&new_root_id) {
110            self.roots.push(new_root_id);
111        }
112    }
113
114    /// Removes all nodes in the provided set from the MAST forest. The nodes MUST be orphaned (i.e.
115    /// have no parent). Otherwise, this parent's reference is considered "dangling" after the
116    /// removal (i.e. will point to an incorrect node after the removal), and this removal operation
117    /// would result in an invalid [`MastForest`].
118    ///
119    /// It also returns the map from old node IDs to new node IDs. Any [`MastNodeId`] used in
120    /// reference to the old [`MastForest`] should be remapped using this map.
121    pub fn remove_nodes(
122        &mut self,
123        nodes_to_remove: &BTreeSet<MastNodeId>,
124    ) -> BTreeMap<MastNodeId, MastNodeId> {
125        if nodes_to_remove.is_empty() {
126            return BTreeMap::new();
127        }
128
129        let old_nodes = core::mem::replace(&mut self.nodes, IndexVec::new());
130        let old_root_ids = core::mem::take(&mut self.roots);
131        let (retained_nodes, id_remappings) = remove_nodes(old_nodes.into_inner(), nodes_to_remove);
132
133        self.remap_and_add_nodes(retained_nodes, &id_remappings);
134        self.remap_and_add_roots(old_root_ids, &id_remappings);
135        id_remappings
136    }
137
138    /// Removes all decorators from this MAST forest.
139    ///
140    /// This method modifies the forest in-place, removing all decorator information
141    /// including operation-indexed decorators, before-enter decorators, after-exit
142    /// decorators, and error codes while keeping the CSR structure valid.
143    ///
144    /// # Examples
145    ///
146    /// ```rust
147    /// use miden_core::mast::MastForest;
148    ///
149    /// let mut forest = MastForest::new();
150    /// // Add decorators and nodes to the forest
151    /// forest.strip_decorators(); // forest is now stripped
152    /// ```
153    pub fn strip_decorators(&mut self) {
154        self.debug_info = DebugInfo::empty_for_nodes(self.nodes.len());
155    }
156
157    /// Compacts the forest by merging duplicate nodes.
158    ///
159    /// This operation performs node deduplication by merging the forest with itself.
160    /// The method assumes that decorators have already been stripped if that is desired.
161    /// The operation modifies the forest in-place, updating all node references as needed.
162    ///
163    /// The process works by:
164    /// 1. Merging the forest with itself to deduplicate identical nodes
165    /// 2. Updating internal node references and remappings
166    /// 3. Modifying the forest in-place with the compacted result
167    ///
168    /// # Examples
169    ///
170    /// ```rust
171    /// use miden_core::mast::MastForest;
172    ///
173    /// let mut forest = MastForest::new();
174    /// // Add nodes to the forest
175    ///
176    /// // First strip decorators if needed
177    /// forest.strip_decorators();
178    ///
179    /// // Then compact the forest
180    /// let root_map = forest.compact();
181    ///
182    /// // Forest is now compacted with duplicate nodes merged
183    /// ```
184    pub fn compact(&mut self) -> MastForestRootMap {
185        // Merge with itself to deduplicate nodes
186        // Note: This cannot fail for a self-merge under normal conditions.
187        // The only possible failures (TooManyNodes, TooManyDecorators) would require the
188        // original forest to be at capacity limits, at which point compaction wouldn't help.
189        let (compacted_forest, root_map) = MastForest::merge([&*self])
190            .expect("Failed to compact MastForest: this should never happen during self-merge");
191
192        // Replace current forest with compacted version
193        *self = compacted_forest;
194
195        root_map
196    }
197
198    /// Merges all `forests` into a new [`MastForest`].
199    ///
200    /// Merging two forests means combining all their constituent parts, i.e. [`MastNode`]s,
201    /// [`Decorator`]s and roots. During this process, any duplicate or
202    /// unreachable nodes are removed. Additionally, [`MastNodeId`]s of nodes as well as
203    /// [`DecoratorId`]s of decorators may change and references to them are remapped to their new
204    /// location.
205    ///
206    /// For example, consider this representation of a forest's nodes with all of these nodes being
207    /// roots:
208    ///
209    /// ```text
210    /// [Block(foo), Block(bar)]
211    /// ```
212    ///
213    /// If we merge another forest into it:
214    ///
215    /// ```text
216    /// [Block(bar), Call(0)]
217    /// ```
218    ///
219    /// then we would expect this forest:
220    ///
221    /// ```text
222    /// [Block(foo), Block(bar), Call(1)]
223    /// ```
224    ///
225    /// - The `Call` to the `bar` block was remapped to its new index (now 1, previously 0).
226    /// - The `Block(bar)` was deduplicated any only exists once in the merged forest.
227    ///
228    /// The function also returns a vector of [`MastForestRootMap`]s, whose length equals the number
229    /// of passed `forests`. The indices in the vector correspond to the ones in `forests`. The map
230    /// of a given forest contains the new locations of its roots in the merged forest. To
231    /// illustrate, the above example would return a vector of two maps:
232    ///
233    /// ```text
234    /// vec![{0 -> 0, 1 -> 1}
235    ///      {0 -> 1, 1 -> 2}]
236    /// ```
237    ///
238    /// - The root locations of the original forest are unchanged.
239    /// - For the second forest, the `bar` block has moved from index 0 to index 1 in the merged
240    ///   forest, and the `Call` has moved from index 1 to 2.
241    ///
242    /// If any forest being merged contains an `External(qux)` node and another forest contains a
243    /// node whose digest is `qux`, then the external node will be replaced with the `qux` node,
244    /// which is effectively deduplication. Decorators are ignored when it comes to merging
245    /// External nodes. This means that an External node with decorators may be replaced by a node
246    /// without decorators or vice versa.
247    pub fn merge<'forest>(
248        forests: impl IntoIterator<Item = &'forest MastForest>,
249    ) -> Result<(MastForest, MastForestRootMap), MastForestError> {
250        MastForestMerger::merge(forests)
251    }
252}
253
254// ------------------------------------------------------------------------------------------------
255/// Helpers
256impl MastForest {
257    /// Adds all provided nodes to the internal set of nodes, remapping all [`MastNodeId`]
258    /// references in those nodes.
259    ///
260    /// # Panics
261    /// - Panics if the internal set of nodes is not empty.
262    fn remap_and_add_nodes(
263        &mut self,
264        nodes_to_add: Vec<MastNode>,
265        id_remappings: &BTreeMap<MastNodeId, MastNodeId>,
266    ) {
267        assert!(self.nodes.is_empty());
268        // extract decorator information from the nodes by converting them into builders
269        let node_builders =
270            nodes_to_add.into_iter().map(|node| node.to_builder(self)).collect::<Vec<_>>();
271
272        // Clear decorator storage after extracting builders (builders contain decorator data)
273        self.debug_info.clear_mappings();
274
275        // Add each node to the new MAST forest, making sure to rewrite any outdated internal
276        // `MastNodeId`s
277        for live_node_builder in node_builders {
278            live_node_builder.remap_children(id_remappings).add_to_forest(self).unwrap();
279        }
280    }
281
282    /// Remaps and adds all old root ids to the internal set of roots.
283    ///
284    /// # Panics
285    /// - Panics if the internal set of roots is not empty.
286    fn remap_and_add_roots(
287        &mut self,
288        old_root_ids: Vec<MastNodeId>,
289        id_remappings: &BTreeMap<MastNodeId, MastNodeId>,
290    ) {
291        assert!(self.roots.is_empty());
292
293        for old_root_id in old_root_ids {
294            let new_root_id = id_remappings.get(&old_root_id).copied().unwrap_or(old_root_id);
295            self.make_root(new_root_id);
296        }
297    }
298}
299
300/// Returns the set of nodes that are live, as well as the mapping from "old ID" to "new ID" for all
301/// live nodes.
302fn remove_nodes(
303    mast_nodes: Vec<MastNode>,
304    nodes_to_remove: &BTreeSet<MastNodeId>,
305) -> (Vec<MastNode>, BTreeMap<MastNodeId, MastNodeId>) {
306    // Note: this allows us to safely use `usize as u32`, guaranteeing that it won't wrap around.
307    assert!(mast_nodes.len() < u32::MAX as usize);
308
309    let mut retained_nodes = Vec::with_capacity(mast_nodes.len());
310    let mut id_remappings = BTreeMap::new();
311
312    for (old_node_index, old_node) in mast_nodes.into_iter().enumerate() {
313        let old_node_id: MastNodeId = MastNodeId(old_node_index as u32);
314
315        if !nodes_to_remove.contains(&old_node_id) {
316            let new_node_id: MastNodeId = MastNodeId(retained_nodes.len() as u32);
317            id_remappings.insert(old_node_id, new_node_id);
318
319            retained_nodes.push(old_node);
320        }
321    }
322
323    (retained_nodes, id_remappings)
324}
325
326// ------------------------------------------------------------------------------------------------
327/// Public accessors
328impl MastForest {
329    /// Returns the [`MastNode`] associated with the provided [`MastNodeId`] if valid, or else
330    /// `None`.
331    ///
332    /// This is the fallible version of indexing (e.g. `mast_forest[node_id]`).
333    #[inline(always)]
334    pub fn get_node_by_id(&self, node_id: MastNodeId) -> Option<&MastNode> {
335        self.nodes.get(node_id)
336    }
337
338    /// Returns the [`MastNodeId`] of the procedure associated with a given digest, if any.
339    #[inline(always)]
340    pub fn find_procedure_root(&self, digest: Word) -> Option<MastNodeId> {
341        self.roots.iter().find(|&&root_id| self[root_id].digest() == digest).copied()
342    }
343
344    /// Returns true if a node with the specified ID is a root of a procedure in this MAST forest.
345    pub fn is_procedure_root(&self, node_id: MastNodeId) -> bool {
346        self.roots.contains(&node_id)
347    }
348
349    /// Returns an iterator over the digests of all procedures in this MAST forest.
350    pub fn procedure_digests(&self) -> impl Iterator<Item = Word> + '_ {
351        self.roots.iter().map(|&root_id| self[root_id].digest())
352    }
353
354    /// Returns an iterator over the digests of local procedures in this MAST forest.
355    ///
356    /// A local procedure is defined as a procedure which is not a single external node.
357    pub fn local_procedure_digests(&self) -> impl Iterator<Item = Word> + '_ {
358        self.roots.iter().filter_map(|&root_id| {
359            let node = &self[root_id];
360            if node.is_external() { None } else { Some(node.digest()) }
361        })
362    }
363
364    /// Returns an iterator over the IDs of the procedures in this MAST forest.
365    pub fn procedure_roots(&self) -> &[MastNodeId] {
366        &self.roots
367    }
368
369    /// Returns the number of procedures in this MAST forest.
370    pub fn num_procedures(&self) -> u32 {
371        self.roots
372            .len()
373            .try_into()
374            .expect("MAST forest contains more than 2^32 procedures.")
375    }
376
377    /// Returns the [Word] representing the content hash of a subset of [`MastNodeId`]s.
378    ///
379    /// # Panics
380    /// This function panics if any `node_ids` is not a node of this forest.
381    pub fn compute_nodes_commitment<'a>(
382        &self,
383        node_ids: impl IntoIterator<Item = &'a MastNodeId>,
384    ) -> Word {
385        let mut digests: Vec<Word> = node_ids.into_iter().map(|&id| self[id].digest()).collect();
386        digests.sort_unstable_by_key(|word| LexicographicWord::from(*word));
387        miden_crypto::hash::rpo::Rpo256::merge_many(&digests)
388    }
389
390    /// Returns the number of nodes in this MAST forest.
391    pub fn num_nodes(&self) -> u32 {
392        self.nodes.len() as u32
393    }
394
395    /// Returns the underlying nodes in this MAST forest.
396    pub fn nodes(&self) -> &[MastNode] {
397        self.nodes.as_slice()
398    }
399
400    pub fn advice_map(&self) -> &AdviceMap {
401        &self.advice_map
402    }
403
404    pub fn advice_map_mut(&mut self) -> &mut AdviceMap {
405        &mut self.advice_map
406    }
407}
408
409// ------------------------------------------------------------------------------------------------
410/// Decorator methods
411impl MastForest {
412    /// Returns a list of all decorators contained in this [MastForest].
413    pub fn decorators(&self) -> &[Decorator] {
414        self.debug_info.decorators()
415    }
416
417    /// Returns the [`Decorator`] associated with the provided [`DecoratorId`] if valid, or else
418    /// `None`.
419    ///
420    /// This is the fallible version of indexing (e.g. `mast_forest[decorator_id]`).
421    #[inline]
422    pub fn decorator_by_id(&self, decorator_id: DecoratorId) -> Option<&Decorator> {
423        self.debug_info.decorator(decorator_id)
424    }
425
426    /// Returns decorator indices for a specific operation within a node.
427    ///
428    /// This is the primary accessor for reading decorators from the centralized storage.
429    /// Returns a slice of decorator IDs for the given operation.
430    #[inline]
431    pub(crate) fn decorator_indices_for_op(
432        &self,
433        node_id: MastNodeId,
434        local_op_idx: usize,
435    ) -> &[DecoratorId] {
436        self.debug_info.decorators_for_operation(node_id, local_op_idx)
437    }
438
439    /// Returns an iterator over decorator references for a specific operation within a node.
440    ///
441    /// This is the preferred method for accessing decorators, as it provides direct
442    /// references to the decorator objects.
443    #[inline]
444    pub fn decorators_for_op<'a>(
445        &'a self,
446        node_id: MastNodeId,
447        local_op_idx: usize,
448    ) -> impl Iterator<Item = &'a Decorator> + 'a {
449        self.decorator_indices_for_op(node_id, local_op_idx)
450            .iter()
451            .map(move |&decorator_id| &self[decorator_id])
452    }
453
454    /// Returns the decorators to be executed before this node is executed.
455    #[inline]
456    pub fn before_enter_decorators(&self, node_id: MastNodeId) -> &[DecoratorId] {
457        self.debug_info.before_enter_decorators(node_id)
458    }
459
460    /// Returns the decorators to be executed after this node is executed.
461    #[inline]
462    pub fn after_exit_decorators(&self, node_id: MastNodeId) -> &[DecoratorId] {
463        self.debug_info.after_exit_decorators(node_id)
464    }
465
466    /// Returns decorator links for a node, including operation indices.
467    ///
468    /// This provides a flattened view of all decorators for a node with their operation indices.
469    #[inline]
470    pub(crate) fn decorator_links_for_node<'a>(
471        &'a self,
472        node_id: MastNodeId,
473    ) -> Result<DecoratedLinks<'a>, DecoratorIndexError> {
474        self.debug_info.decorator_links_for_node(node_id)
475    }
476
477    /// Adds a decorator to the forest, and returns the associated [`DecoratorId`].
478    pub fn add_decorator(&mut self, decorator: Decorator) -> Result<DecoratorId, MastForestError> {
479        self.debug_info.add_decorator(decorator)
480    }
481
482    /// Adds decorator IDs for a node to the storage.
483    ///
484    /// Used when building nodes for efficient decorator access during execution.
485    ///
486    /// # Note
487    /// This method does not validate decorator IDs immediately. Validation occurs during
488    /// operations that need to access the actual decorator data (e.g., merging, serialization).
489    #[inline]
490    pub(crate) fn register_node_decorators(
491        &mut self,
492        node_id: MastNodeId,
493        before_enter: &[DecoratorId],
494        after_exit: &[DecoratorId],
495    ) {
496        self.debug_info.register_node_decorators(node_id, before_enter, after_exit);
497    }
498}
499
500// ------------------------------------------------------------------------------------------------
501/// Error message methods
502impl MastForest {
503    /// Given an error code as a Felt, resolves it to its corresponding error message.
504    pub fn resolve_error_message(&self, code: Felt) -> Option<Arc<str>> {
505        let key = u64::from(code);
506        self.debug_info.error_message(key)
507    }
508
509    /// Registers an error message in the MAST Forest and returns the corresponding error code as a
510    /// Felt.
511    pub fn register_error(&mut self, msg: Arc<str>) -> Felt {
512        let code: Felt = error_code_from_msg(&msg);
513        // we use u64 as keys for the map
514        self.debug_info.insert_error_code(code.as_int(), msg);
515        code
516    }
517}
518
519// MAST FOREST INDEXING
520// ------------------------------------------------------------------------------------------------
521
522impl Index<MastNodeId> for MastForest {
523    type Output = MastNode;
524
525    #[inline(always)]
526    fn index(&self, node_id: MastNodeId) -> &Self::Output {
527        &self.nodes[node_id]
528    }
529}
530
531impl IndexMut<MastNodeId> for MastForest {
532    #[inline(always)]
533    fn index_mut(&mut self, node_id: MastNodeId) -> &mut Self::Output {
534        &mut self.nodes[node_id]
535    }
536}
537
538impl Index<DecoratorId> for MastForest {
539    type Output = Decorator;
540
541    #[inline(always)]
542    fn index(&self, decorator_id: DecoratorId) -> &Self::Output {
543        self.debug_info.decorator(decorator_id).expect("DecoratorId out of bounds")
544    }
545}
546
547impl IndexMut<DecoratorId> for MastForest {
548    #[inline(always)]
549    fn index_mut(&mut self, decorator_id: DecoratorId) -> &mut Self::Output {
550        self.debug_info.decorator_mut(decorator_id).expect("DecoratorId out of bounds")
551    }
552}
553
554// MAST NODE ID
555// ================================================================================================
556
557/// An opaque handle to a [`MastNode`] in some [`MastForest`]. It is the responsibility of the user
558/// to use a given [`MastNodeId`] with the corresponding [`MastForest`].
559///
560/// Note that the [`MastForest`] does *not* ensure that equal [`MastNode`]s have equal
561/// [`MastNodeId`] handles. Hence, [`MastNodeId`] equality must not be used to test for equality of
562/// the underlying [`MastNode`].
563#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
564#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
565#[cfg_attr(feature = "serde", serde(transparent))]
566#[cfg_attr(all(feature = "arbitrary", test), miden_test_serde_macros::serde_test)]
567pub struct MastNodeId(u32);
568
569/// Operations that mutate a MAST often produce this mapping between old and new NodeIds.
570pub type Remapping = BTreeMap<MastNodeId, MastNodeId>;
571
572impl MastNodeId {
573    /// Returns a new `MastNodeId` with the provided inner value, or an error if the provided
574    /// `value` is greater than the number of nodes in the forest.
575    ///
576    /// For use in deserialization.
577    pub fn from_u32_safe(
578        value: u32,
579        mast_forest: &MastForest,
580    ) -> Result<Self, DeserializationError> {
581        Self::from_u32_with_node_count(value, mast_forest.nodes.len())
582    }
583
584    /// Returns a new [`MastNodeId`] with the provided `node_id`, or an error if `node_id` is
585    /// greater than the number of nodes in the [`MastForest`] for which this ID is being
586    /// constructed.
587    pub fn from_usize_safe(
588        node_id: usize,
589        mast_forest: &MastForest,
590    ) -> Result<Self, DeserializationError> {
591        let node_id: u32 = node_id.try_into().map_err(|_| {
592            DeserializationError::InvalidValue(format!(
593                "node id '{node_id}' does not fit into a u32"
594            ))
595        })?;
596        MastNodeId::from_u32_safe(node_id, mast_forest)
597    }
598
599    /// Returns a new [`MastNodeId`] from the given `value` without checking its validity.
600    pub(crate) fn new_unchecked(value: u32) -> Self {
601        Self(value)
602    }
603
604    /// Returns a new [`MastNodeId`] with the provided `id`, or an error if `id` is greater or equal
605    /// to `node_count`. The `node_count` is the total number of nodes in the [`MastForest`] for
606    /// which this ID is being constructed.
607    ///
608    /// This function can be used when deserializing an id whose corresponding node is not yet in
609    /// the forest and [`Self::from_u32_safe`] would fail. For instance, when deserializing the ids
610    /// referenced by the Join node in this forest:
611    ///
612    /// ```text
613    /// [Join(1, 2), Block(foo), Block(bar)]
614    /// ```
615    ///
616    /// Since it is less safe than [`Self::from_u32_safe`] and usually not needed it is not public.
617    pub(super) fn from_u32_with_node_count(
618        id: u32,
619        node_count: usize,
620    ) -> Result<Self, DeserializationError> {
621        if (id as usize) < node_count {
622            Ok(Self(id))
623        } else {
624            Err(DeserializationError::InvalidValue(format!(
625                "Invalid deserialized MAST node ID '{id}', but {node_count} is the number of nodes in the forest",
626            )))
627        }
628    }
629
630    /// Remap the NodeId to its new position using the given [`Remapping`].
631    pub fn remap(&self, remapping: &Remapping) -> Self {
632        *remapping.get(self).unwrap_or(self)
633    }
634}
635
636impl From<u32> for MastNodeId {
637    fn from(value: u32) -> Self {
638        MastNodeId::new_unchecked(value)
639    }
640}
641
642impl Idx for MastNodeId {}
643
644impl From<MastNodeId> for u32 {
645    fn from(value: MastNodeId) -> Self {
646        value.0
647    }
648}
649
650impl fmt::Display for MastNodeId {
651    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
652        write!(f, "MastNodeId({})", self.0)
653    }
654}
655
656#[cfg(any(test, feature = "arbitrary"))]
657impl proptest::prelude::Arbitrary for MastNodeId {
658    type Parameters = ();
659
660    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
661        use proptest::prelude::*;
662        any::<u32>().prop_map(MastNodeId).boxed()
663    }
664
665    type Strategy = proptest::prelude::BoxedStrategy<Self>;
666}
667
668// ITERATOR
669
670/// Iterates over all the nodes a root depends on, in pre-order. The iteration can include other
671/// roots in the same forest.
672pub struct SubtreeIterator<'a> {
673    forest: &'a MastForest,
674    discovered: Vec<MastNodeId>,
675    unvisited: Vec<MastNodeId>,
676}
677impl<'a> SubtreeIterator<'a> {
678    pub fn new(root: &MastNodeId, forest: &'a MastForest) -> Self {
679        let discovered = vec![];
680        let unvisited = vec![*root];
681        SubtreeIterator { forest, discovered, unvisited }
682    }
683}
684impl Iterator for SubtreeIterator<'_> {
685    type Item = MastNodeId;
686    fn next(&mut self) -> Option<MastNodeId> {
687        while let Some(id) = self.unvisited.pop() {
688            let node = &self.forest[id];
689            if !node.has_children() {
690                return Some(id);
691            } else {
692                self.discovered.push(id);
693                node.append_children_to(&mut self.unvisited);
694            }
695        }
696        self.discovered.pop()
697    }
698}
699
700// DECORATOR ID
701// ================================================================================================
702
703/// An opaque handle to a [`Decorator`] in some [`MastForest`]. It is the responsibility of the user
704/// to use a given [`DecoratorId`] with the corresponding [`MastForest`].
705#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
706#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
707#[cfg_attr(feature = "serde", serde(transparent))]
708pub struct DecoratorId(u32);
709
710impl DecoratorId {
711    /// Returns a new `DecoratorId` with the provided inner value, or an error if the provided
712    /// `value` is greater than the number of nodes in the forest.
713    ///
714    /// For use in deserialization.
715    pub fn from_u32_safe(
716        value: u32,
717        mast_forest: &MastForest,
718    ) -> Result<Self, DeserializationError> {
719        Self::from_u32_bounded(value, mast_forest.debug_info.num_decorators())
720    }
721
722    /// Returns a new `DecoratorId` with the provided inner value, or an error if the provided
723    /// `value` is greater than or equal to `bound`.
724    ///
725    /// For use in deserialization when the bound is known without needing the full MastForest.
726    pub fn from_u32_bounded(value: u32, bound: usize) -> Result<Self, DeserializationError> {
727        if (value as usize) < bound {
728            Ok(Self(value))
729        } else {
730            Err(DeserializationError::InvalidValue(format!(
731                "Invalid deserialized MAST decorator id '{}', but allows only {} decorators",
732                value, bound,
733            )))
734        }
735    }
736
737    /// Creates a new [`DecoratorId`] without checking its validity.
738    pub(crate) fn new_unchecked(value: u32) -> Self {
739        Self(value)
740    }
741}
742
743impl From<u32> for DecoratorId {
744    fn from(value: u32) -> Self {
745        DecoratorId::new_unchecked(value)
746    }
747}
748
749impl Idx for DecoratorId {}
750
751impl From<DecoratorId> for u32 {
752    fn from(value: DecoratorId) -> Self {
753        value.0
754    }
755}
756
757impl fmt::Display for DecoratorId {
758    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
759        write!(f, "DecoratorId({})", self.0)
760    }
761}
762
763impl Serializable for DecoratorId {
764    fn write_into<W: ByteWriter>(&self, target: &mut W) {
765        self.0.write_into(target)
766    }
767}
768
769/// Derives an error code from an error message by hashing the message and returning the 0th element
770/// of the resulting [`Word`].
771pub fn error_code_from_msg(msg: impl AsRef<str>) -> Felt {
772    // hash the message and return 0th felt of the resulting Word
773    hash_string_to_word(msg.as_ref())[0]
774}
775
776// MAST FOREST ERROR
777// ================================================================================================
778
779/// Represents the types of errors that can occur when dealing with MAST forest.
780#[derive(Debug, thiserror::Error, PartialEq)]
781pub enum MastForestError {
782    #[error("MAST forest decorator count exceeds the maximum of {} decorators", u32::MAX)]
783    TooManyDecorators,
784    #[error("MAST forest node count exceeds the maximum of {} nodes", MastForest::MAX_NODES)]
785    TooManyNodes,
786    #[error("node id {0} is greater than or equal to forest length {1}")]
787    NodeIdOverflow(MastNodeId, usize),
788    #[error("decorator id {0} is greater than or equal to decorator count {1}")]
789    DecoratorIdOverflow(DecoratorId, usize),
790    #[error("basic block cannot be created from an empty list of operations")]
791    EmptyBasicBlock,
792    #[error(
793        "decorator root of child with node id {0} is missing but is required for fingerprint computation"
794    )]
795    ChildFingerprintMissing(MastNodeId),
796    #[error("advice map key {0} already exists when merging forests")]
797    AdviceMapKeyCollisionOnMerge(Word),
798    #[error("decorator storage error: {0}")]
799    DecoratorError(DecoratorIndexError),
800    #[error("digest is required for deserialization")]
801    DigestRequiredForDeserialization,
802}
803
804// Custom serde implementations for MastForest that handle linked decorators properly
805// by delegating to the existing winter-utils serialization which already handles
806// the conversion between linked and owned decorator formats.
807#[cfg(feature = "serde")]
808impl serde::Serialize for MastForest {
809    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
810    where
811        S: serde::Serializer,
812    {
813        // Use the existing winter-utils serialization which already handles linked decorators
814        let bytes = crate::utils::Serializable::to_bytes(self);
815        serializer.serialize_bytes(&bytes)
816    }
817}
818
819#[cfg(feature = "serde")]
820impl<'de> serde::Deserialize<'de> for MastForest {
821    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
822    where
823        D: serde::Deserializer<'de>,
824    {
825        // Deserialize bytes, then use winter-utils Deserializable
826        let bytes = Vec::<u8>::deserialize(deserializer)?;
827        let mut slice_reader = winter_utils::SliceReader::new(&bytes);
828        crate::utils::Deserializable::read_from(&mut slice_reader).map_err(serde::de::Error::custom)
829    }
830}