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