miden_core/mast/
mod.rs

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