miden_core/mast/

mod.rs

//! MAST forest: a collection of procedures represented as Merkle trees.
//!
//! # Deserializing from untrusted sources
//!
//! When loading a `MastForest` from bytes you don't fully trust (network, user upload, etc.),
//! use [`UntrustedMastForest`] instead of calling `MastForest::read_from_bytes` directly:
//!
//! ```ignore
//! use miden_core::mast::UntrustedMastForest;
//!
//! let forest = UntrustedMastForest::read_from_bytes(&bytes)?
//!     .validate()?;
//! ```
//!
//! For maximum protection against denial-of-service attacks from malicious input, use
//! [`UntrustedMastForest::read_from_bytes_with_budget`] which limits memory consumption:
//!
//! ```ignore
//! use miden_core::mast::UntrustedMastForest;
//!
//! // Budget limits pre-allocation sizes and total bytes consumed
//! let forest = UntrustedMastForest::read_from_bytes_with_budget(&bytes, bytes.len())?
//!     .validate()?;
//! ```
//!
//! This recomputes all node hashes and checks structural invariants before returning a usable
//! `MastForest`. Direct deserialization via `MastForest::read_from_bytes` trusts the serialized
//! hashes and should only be used for data from trusted sources (e.g. compiled locally).

use alloc::{
    collections::{BTreeMap, BTreeSet},
    string::String,
    sync::Arc,
    vec::Vec,
};
use core::{
    fmt,
    ops::{Index, IndexMut},
};

use miden_utils_sync::OnceLockCompat;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};

mod node;
#[cfg(any(test, feature = "arbitrary"))]
pub use node::arbitrary;
pub use node::{
    BasicBlockNode, BasicBlockNodeBuilder, CallNode, CallNodeBuilder, DecoratedOpLink,
    DecoratorOpLinkIterator, DecoratorStore, DynNode, DynNodeBuilder, ExternalNode,
    ExternalNodeBuilder, JoinNode, JoinNodeBuilder, LoopNode, LoopNodeBuilder,
    MastForestContributor, MastNode, MastNodeBuilder, MastNodeExt, OP_BATCH_SIZE, OP_GROUP_SIZE,
    OpBatch, OperationOrDecorator, SplitNode, SplitNodeBuilder,
};

use crate::{
    Felt, LexicographicWord, Word,
    advice::AdviceMap,
    field::PrimeField64,
    operations::{AssemblyOp, Decorator},
    serde::{
        BudgetedReader, ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable,
        SliceReader,
    },
    utils::{Idx, IndexVec, hash_string_to_word},
};

mod debuginfo;
pub use debuginfo::{
    AsmOpIndexError, DebugInfo, DecoratedLinks, DecoratedLinksIter, DecoratorIndexError,
    NodeToDecoratorIds, OpToAsmOpId, OpToDecoratorIds,
};

mod serialization;

mod merger;
pub(crate) use merger::MastForestMerger;
pub use merger::MastForestRootMap;

mod multi_forest_node_iterator;
pub(crate) use multi_forest_node_iterator::*;

mod node_fingerprint;
pub use node_fingerprint::{DecoratorFingerprint, MastNodeFingerprint};

mod node_builder_utils;
pub use node_builder_utils::build_node_with_remapped_ids;

#[cfg(test)]
mod tests;

// MAST FOREST
// ================================================================================================

/// Represents one or more procedures, represented as a collection of [`MastNode`]s.
///
/// A [`MastForest`] does not have an entrypoint, and hence is not executable. A
/// [`crate::program::Program`] can be built from a [`MastForest`] to specify an entrypoint.
#[derive(Clone, Debug, Default)]
pub struct MastForest {
    /// All of the nodes local to the trees comprising the MAST forest.
    nodes: IndexVec<MastNodeId, MastNode>,

    /// Roots of procedures defined within this MAST forest.
    roots: Vec<MastNodeId>,

    /// Advice map to be loaded into the VM prior to executing procedures from this MAST forest.
    advice_map: AdviceMap,

    /// Debug information including decorators and error codes.
    /// Always present (as per issue #1821), but can be empty for stripped builds.
    debug_info: DebugInfo,

    /// Cached commitment to this MAST forest (commitment to all roots).
    /// This is computed lazily on first access and invalidated on any mutation.
    commitment_cache: OnceLockCompat<Word>,
}

// ------------------------------------------------------------------------------------------------
/// Constructors
impl MastForest {
    /// Creates a new empty [`MastForest`].
    pub fn new() -> Self {
        Self {
            nodes: IndexVec::new(),
            roots: Vec::new(),
            advice_map: AdviceMap::default(),
            debug_info: DebugInfo::new(),
            commitment_cache: OnceLockCompat::new(),
        }
    }
}

// ------------------------------------------------------------------------------------------------
/// Equality implementations
impl PartialEq for MastForest {
    fn eq(&self, other: &Self) -> bool {
        // Compare all fields except commitment_cache, which is derived data
        self.nodes == other.nodes
            && self.roots == other.roots
            && self.advice_map == other.advice_map
            && self.debug_info == other.debug_info
    }
}

impl Eq for MastForest {}

// ------------------------------------------------------------------------------------------------
/// State mutators
impl MastForest {
    /// The maximum number of nodes that can be stored in a single MAST forest.
    const MAX_NODES: usize = (1 << 30) - 1;

    /// Marks the given [`MastNodeId`] as being the root of a procedure.
    ///
    /// If the specified node is already marked as a root, this will have no effect.
    ///
    /// # Panics
    /// - if `new_root_id`'s internal index is larger than the number of nodes in this forest (i.e.
    ///   clearly doesn't belong to this MAST forest).
    pub fn make_root(&mut self, new_root_id: MastNodeId) {
        assert!(new_root_id.to_usize() < self.nodes.len());

        if !self.roots.contains(&new_root_id) {
            self.roots.push(new_root_id);
            // Invalidate the cached commitment since we modified the roots
            self.commitment_cache.take();
        }
    }

    /// Removes all nodes in the provided set from the MAST forest. The nodes MUST be orphaned (i.e.
    /// have no parent). Otherwise, this parent's reference is considered "dangling" after the
    /// removal (i.e. will point to an incorrect node after the removal), and this removal operation
    /// would result in an invalid [`MastForest`].
    ///
    /// It also returns the map from old node IDs to new node IDs. Any [`MastNodeId`] used in
    /// reference to the old [`MastForest`] should be remapped using this map.
    pub fn remove_nodes(
        &mut self,
        nodes_to_remove: &BTreeSet<MastNodeId>,
    ) -> BTreeMap<MastNodeId, MastNodeId> {
        if nodes_to_remove.is_empty() {
            return BTreeMap::new();
        }

        let old_nodes = core::mem::replace(&mut self.nodes, IndexVec::new());
        let old_root_ids = core::mem::take(&mut self.roots);
        let (retained_nodes, id_remappings) = remove_nodes(old_nodes.into_inner(), nodes_to_remove);

        self.remap_and_add_nodes(retained_nodes, &id_remappings);
        self.remap_and_add_roots(old_root_ids, &id_remappings);

        // Remap the asm_op_storage to use the new node IDs
        self.debug_info.remap_asm_op_storage(&id_remappings);

        // Invalidate the cached commitment since we modified the forest structure
        self.commitment_cache.take();

        id_remappings
    }

    /// Clears all [`DebugInfo`] from this forest: decorators, error codes, and procedure names.
    ///
    /// ```
    /// # use miden_core::mast::MastForest;
    /// let mut forest = MastForest::new();
    /// forest.clear_debug_info();
    /// assert!(forest.decorators().is_empty());
    /// ```
    pub fn clear_debug_info(&mut self) {
        self.debug_info = DebugInfo::empty_for_nodes(self.nodes.len());
    }

    /// Compacts the forest by merging duplicate nodes.
    ///
    /// This operation performs node deduplication by merging the forest with itself.
    /// The method assumes that debug info has already been cleared if that is desired.
    /// This method consumes the forest and returns a new compacted forest.
    ///
    /// The process works by:
    /// 1. Merging the forest with itself to deduplicate identical nodes
    /// 2. Updating internal node references and remappings
    /// 3. Returning the compacted forest and root map
    ///
    /// # Examples
    ///
    /// ```rust
    /// use miden_core::mast::MastForest;
    ///
    /// let mut forest = MastForest::new();
    /// // Add nodes to the forest
    ///
    /// // First clear debug info if needed
    /// forest.clear_debug_info();
    ///
    /// // Then compact the forest (consumes the original)
    /// let (compacted_forest, root_map) = forest.compact();
    ///
    /// // compacted_forest is now compacted with duplicate nodes merged
    /// ```
    pub fn compact(self) -> (MastForest, MastForestRootMap) {
        // Merge with itself to deduplicate nodes
        // Note: This cannot fail for a self-merge under normal conditions.
        // The only possible failures (TooManyNodes, TooManyDecorators) would require the
        // original forest to be at capacity limits, at which point compaction wouldn't help.
        MastForest::merge([&self])
            .expect("Failed to compact MastForest: this should never happen during self-merge")
    }

    /// Merges all `forests` into a new [`MastForest`].
    ///
    /// Merging two forests means combining all their constituent parts, i.e. [`MastNode`]s,
    /// [`Decorator`]s and roots. During this process, any duplicate or
    /// unreachable nodes are removed. Additionally, [`MastNodeId`]s of nodes as well as
    /// [`DecoratorId`]s of decorators may change and references to them are remapped to their new
    /// location.
    ///
    /// For example, consider this representation of a forest's nodes with all of these nodes being
    /// roots:
    ///
    /// ```text
    /// [Block(foo), Block(bar)]
    /// ```
    ///
    /// If we merge another forest into it:
    ///
    /// ```text
    /// [Block(bar), Call(0)]
    /// ```
    ///
    /// then we would expect this forest:
    ///
    /// ```text
    /// [Block(foo), Block(bar), Call(1)]
    /// ```
    ///
    /// - The `Call` to the `bar` block was remapped to its new index (now 1, previously 0).
    /// - The `Block(bar)` was deduplicated any only exists once in the merged forest.
    ///
    /// The function also returns a vector of [`MastForestRootMap`]s, whose length equals the number
    /// of passed `forests`. The indices in the vector correspond to the ones in `forests`. The map
    /// of a given forest contains the new locations of its roots in the merged forest. To
    /// illustrate, the above example would return a vector of two maps:
    ///
    /// ```text
    /// vec![{0 -> 0, 1 -> 1}
    ///      {0 -> 1, 1 -> 2}]
    /// ```
    ///
    /// - The root locations of the original forest are unchanged.
    /// - For the second forest, the `bar` block has moved from index 0 to index 1 in the merged
    ///   forest, and the `Call` has moved from index 1 to 2.
    ///
    /// If any forest being merged contains an `External(qux)` node and another forest contains a
    /// node whose digest is `qux`, then the external node will be replaced with the `qux` node,
    /// which is effectively deduplication. Decorators are ignored when it comes to merging
    /// External nodes. This means that an External node with decorators may be replaced by a node
    /// without decorators or vice versa.
    pub fn merge<'forest>(
        forests: impl IntoIterator<Item = &'forest MastForest>,
    ) -> Result<(MastForest, MastForestRootMap), MastForestError> {
        MastForestMerger::merge(forests)
    }
}

// ------------------------------------------------------------------------------------------------
/// Helpers
impl MastForest {
    /// Adds all provided nodes to the internal set of nodes, remapping all [`MastNodeId`]
    /// references in those nodes.
    ///
    /// # Panics
    /// - Panics if the internal set of nodes is not empty.
    fn remap_and_add_nodes(
        &mut self,
        nodes_to_add: Vec<MastNode>,
        id_remappings: &BTreeMap<MastNodeId, MastNodeId>,
    ) {
        assert!(self.nodes.is_empty());
        // extract decorator information from the nodes by converting them into builders
        let node_builders =
            nodes_to_add.into_iter().map(|node| node.to_builder(self)).collect::<Vec<_>>();

        // Clear decorator storage after extracting builders (builders contain decorator data)
        self.debug_info.clear_mappings();

        // Add each node to the new MAST forest, making sure to rewrite any outdated internal
        // `MastNodeId`s
        for live_node_builder in node_builders {
            live_node_builder.remap_children(id_remappings).add_to_forest(self).unwrap();
        }
    }

    /// Remaps and adds all old root ids to the internal set of roots.
    ///
    /// # Panics
    /// - Panics if the internal set of roots is not empty.
    fn remap_and_add_roots(
        &mut self,
        old_root_ids: Vec<MastNodeId>,
        id_remappings: &BTreeMap<MastNodeId, MastNodeId>,
    ) {
        assert!(self.roots.is_empty());

        for old_root_id in old_root_ids {
            let new_root_id = id_remappings.get(&old_root_id).copied().unwrap_or(old_root_id);
            self.make_root(new_root_id);
        }
    }
}

/// Returns the set of nodes that are live, as well as the mapping from "old ID" to "new ID" for all
/// live nodes.
fn remove_nodes(
    mast_nodes: Vec<MastNode>,
    nodes_to_remove: &BTreeSet<MastNodeId>,
) -> (Vec<MastNode>, BTreeMap<MastNodeId, MastNodeId>) {
    // Note: this allows us to safely use `usize as u32`, guaranteeing that it won't wrap around.
    assert!(mast_nodes.len() < u32::MAX as usize);

    let mut retained_nodes = Vec::with_capacity(mast_nodes.len());
    let mut id_remappings = BTreeMap::new();

    for (old_node_index, old_node) in mast_nodes.into_iter().enumerate() {
        let old_node_id: MastNodeId = MastNodeId(old_node_index as u32);

        if !nodes_to_remove.contains(&old_node_id) {
            let new_node_id: MastNodeId = MastNodeId(retained_nodes.len() as u32);
            id_remappings.insert(old_node_id, new_node_id);

            retained_nodes.push(old_node);
        }
    }

    (retained_nodes, id_remappings)
}

// ------------------------------------------------------------------------------------------------
/// Public accessors
impl MastForest {
    /// Returns the [`MastNode`] associated with the provided [`MastNodeId`] if valid, or else
    /// `None`.
    ///
    /// This is the fallible version of indexing (e.g. `mast_forest[node_id]`).
    #[inline(always)]
    pub fn get_node_by_id(&self, node_id: MastNodeId) -> Option<&MastNode> {
        self.nodes.get(node_id)
    }

    /// Returns the [`MastNodeId`] of the procedure associated with a given digest, if any.
    #[inline(always)]
    pub fn find_procedure_root(&self, digest: Word) -> Option<MastNodeId> {
        self.roots.iter().find(|&&root_id| self[root_id].digest() == digest).copied()
    }

    /// Returns true if a node with the specified ID is a root of a procedure in this MAST forest.
    pub fn is_procedure_root(&self, node_id: MastNodeId) -> bool {
        self.roots.contains(&node_id)
    }

    /// Returns an iterator over the digests of all procedures in this MAST forest.
    pub fn procedure_digests(&self) -> impl Iterator<Item = Word> + '_ {
        self.roots.iter().map(|&root_id| self[root_id].digest())
    }

    /// Returns an iterator over the digests of local procedures in this MAST forest.
    ///
    /// A local procedure is defined as a procedure which is not a single external node.
    pub fn local_procedure_digests(&self) -> impl Iterator<Item = Word> + '_ {
        self.roots.iter().filter_map(|&root_id| {
            let node = &self[root_id];
            if node.is_external() { None } else { Some(node.digest()) }
        })
    }

    /// Returns an iterator over the IDs of the procedures in this MAST forest.
    pub fn procedure_roots(&self) -> &[MastNodeId] {
        &self.roots
    }

    /// Returns the number of procedures in this MAST forest.
    pub fn num_procedures(&self) -> u32 {
        self.roots
            .len()
            .try_into()
            .expect("MAST forest contains more than 2^32 procedures.")
    }

    /// Returns the [Word] representing the content hash of a subset of [`MastNodeId`]s.
    ///
    /// # Panics
    /// This function panics if any `node_ids` is not a node of this forest.
    pub fn compute_nodes_commitment<'a>(
        &self,
        node_ids: impl IntoIterator<Item = &'a MastNodeId>,
    ) -> Word {
        let mut digests: Vec<Word> = node_ids.into_iter().map(|&id| self[id].digest()).collect();
        digests.sort_unstable_by_key(|word| LexicographicWord::from(*word));
        miden_crypto::hash::poseidon2::Poseidon2::merge_many(&digests)
    }

    /// Returns the commitment to this MAST forest.
    ///
    /// The commitment is computed as the sequential hash of all procedure roots in the forest.
    /// This value is cached after the first computation and reused for subsequent calls,
    /// unless the forest is mutated (in which case the cache is invalidated).
    ///
    /// The commitment uniquely identifies the forest's structure, as each root's digest
    /// transitively includes all of its descendants. Therefore, a commitment to all roots
    /// is a commitment to the entire forest.
    pub fn commitment(&self) -> Word {
        *self.commitment_cache.get_or_init(|| self.compute_nodes_commitment(&self.roots))
    }

    /// Returns the number of nodes in this MAST forest.
    pub fn num_nodes(&self) -> u32 {
        self.nodes.len() as u32
    }

    /// Returns the underlying nodes in this MAST forest.
    pub fn nodes(&self) -> &[MastNode] {
        self.nodes.as_slice()
    }

    pub fn advice_map(&self) -> &AdviceMap {
        &self.advice_map
    }

    pub fn advice_map_mut(&mut self) -> &mut AdviceMap {
        &mut self.advice_map
    }

    // SERIALIZATION
    // --------------------------------------------------------------------------------------------

    /// Serializes this MastForest without debug information.
    ///
    /// This produces a smaller output by omitting decorators, error codes, and procedure names.
    /// The resulting bytes can be deserialized with the standard [`Deserializable`] impl,
    /// which auto-detects the format and creates an empty [`DebugInfo`].
    ///
    /// Use this for production builds where debug info is not needed.
    ///
    /// # Example
    ///
    /// ```
    /// use miden_core::{mast::MastForest, serde::Serializable};
    ///
    /// let forest = MastForest::new();
    ///
    /// // Full serialization (with debug info)
    /// let full_bytes = forest.to_bytes();
    ///
    /// // Stripped serialization (without debug info)
    /// let mut stripped_bytes = Vec::new();
    /// forest.write_stripped(&mut stripped_bytes);
    ///
    /// // Both can be deserialized the same way
    /// // let restored = MastForest::read_from_bytes(&stripped_bytes).unwrap();
    /// ```
    pub fn write_stripped<W: ByteWriter>(&self, target: &mut W) {
        use serialization::StrippedMastForest;
        StrippedMastForest(self).write_into(target);
    }
}

// ------------------------------------------------------------------------------------------------
/// Decorator methods
impl MastForest {
    /// Returns a list of all decorators contained in this [MastForest].
    pub fn decorators(&self) -> &[Decorator] {
        self.debug_info.decorators()
    }

    /// Returns the [`Decorator`] associated with the provided [`DecoratorId`] if valid, or else
    /// `None`.
    ///
    /// This is the fallible version of indexing (e.g. `mast_forest[decorator_id]`).
    #[inline]
    pub fn decorator_by_id(&self, decorator_id: DecoratorId) -> Option<&Decorator> {
        self.debug_info.decorator(decorator_id)
    }

    /// Returns decorator indices for a specific operation within a node.
    ///
    /// This is the primary accessor for reading decorators from the centralized storage.
    /// Returns a slice of decorator IDs for the given operation.
    #[inline]
    pub(crate) fn decorator_indices_for_op(
        &self,
        node_id: MastNodeId,
        local_op_idx: usize,
    ) -> &[DecoratorId] {
        self.debug_info.decorators_for_operation(node_id, local_op_idx)
    }

    /// Returns an iterator over decorator references for a specific operation within a node.
    ///
    /// This is the preferred method for accessing decorators, as it provides direct
    /// references to the decorator objects.
    #[inline]
    pub fn decorators_for_op<'a>(
        &'a self,
        node_id: MastNodeId,
        local_op_idx: usize,
    ) -> impl Iterator<Item = &'a Decorator> + 'a {
        self.decorator_indices_for_op(node_id, local_op_idx)
            .iter()
            .map(move |&decorator_id| &self[decorator_id])
    }

    /// Returns the decorators to be executed before this node is executed.
    #[inline]
    pub fn before_enter_decorators(&self, node_id: MastNodeId) -> &[DecoratorId] {
        self.debug_info.before_enter_decorators(node_id)
    }

    /// Returns the decorators to be executed after this node is executed.
    #[inline]
    pub fn after_exit_decorators(&self, node_id: MastNodeId) -> &[DecoratorId] {
        self.debug_info.after_exit_decorators(node_id)
    }

    /// Returns decorator links for a node, including operation indices.
    ///
    /// This provides a flattened view of all decorators for a node with their operation indices.
    #[inline]
    pub(crate) fn decorator_links_for_node<'a>(
        &'a self,
        node_id: MastNodeId,
    ) -> Result<DecoratedLinks<'a>, DecoratorIndexError> {
        self.debug_info.decorator_links_for_node(node_id)
    }

    /// Adds a decorator to the forest, and returns the associated [`DecoratorId`].
    pub fn add_decorator(&mut self, decorator: Decorator) -> Result<DecoratorId, MastForestError> {
        self.debug_info.add_decorator(decorator)
    }

    /// Adds decorator IDs for a node to the storage.
    ///
    /// Used when building nodes for efficient decorator access during execution.
    ///
    /// # Note
    /// This method does not validate decorator IDs immediately. Validation occurs during
    /// operations that need to access the actual decorator data (e.g., merging, serialization).
    #[inline]
    pub(crate) fn register_node_decorators(
        &mut self,
        node_id: MastNodeId,
        before_enter: &[DecoratorId],
        after_exit: &[DecoratorId],
    ) {
        self.debug_info.register_node_decorators(node_id, before_enter, after_exit);
    }

    /// Returns the [`AssemblyOp`] associated with a node.
    ///
    /// For basic block nodes with a `target_op_idx`, returns the AssemblyOp for that operation.
    /// For other nodes or when no `target_op_idx` is provided, returns the first AssemblyOp.
    pub fn get_assembly_op(
        &self,
        node_id: MastNodeId,
        target_op_idx: Option<usize>,
    ) -> Option<&AssemblyOp> {
        match target_op_idx {
            Some(op_idx) => self.debug_info.asm_op_for_operation(node_id, op_idx),
            None => self.debug_info.first_asm_op_for_node(node_id),
        }
    }
}

// ------------------------------------------------------------------------------------------------
/// Validation methods
impl MastForest {
    /// Validates that all BasicBlockNodes in this forest satisfy the core invariants:
    /// 1. Power-of-two number of groups in each batch
    /// 2. No operation group ends with an operation requiring an immediate value
    /// 3. The last operation group in a batch cannot contain operations requiring immediate values
    /// 4. OpBatch structural consistency (num_groups <= BATCH_SIZE, group size <= GROUP_SIZE,
    ///    indptr integrity, bounds checking)
    ///
    /// This addresses the gap created by PR 2094, where padding NOOPs are now inserted
    /// at assembly time rather than dynamically during execution, and adds comprehensive
    /// structural validation to prevent deserialization-time panics.
    pub fn validate(&self) -> Result<(), MastForestError> {
        // Validate basic block batch invariants
        for (node_id_idx, node) in self.nodes.iter().enumerate() {
            let node_id =
                MastNodeId::new_unchecked(node_id_idx.try_into().expect("too many nodes"));
            if let MastNode::Block(basic_block) = node {
                basic_block.validate_batch_invariants().map_err(|error_msg| {
                    MastForestError::InvalidBatchPadding(node_id, error_msg)
                })?;
            }
        }

        // Validate that all procedure name digests correspond to procedure roots in the forest
        for (digest, _) in self.debug_info.procedure_names() {
            if self.find_procedure_root(digest).is_none() {
                return Err(MastForestError::InvalidProcedureNameDigest(digest));
            }
        }

        Ok(())
    }

    /// Validates topological ordering of nodes and recomputes all node hashes.
    ///
    /// This method iterates through all nodes in index order, verifying:
    /// 1. All child references point to nodes with smaller indices (topological order)
    /// 2. Each node's recomputed digest matches its stored digest
    ///
    /// # Errors
    ///
    /// Returns `MastForestError::ForwardReference` if any node references a child that
    /// appears later in the forest.
    ///
    /// Returns `MastForestError::HashMismatch` if any node's recomputed digest doesn't
    /// match its stored digest.
    fn validate_node_hashes(&self) -> Result<(), MastForestError> {
        use crate::chiplets::hasher;

        /// Checks that child_id references a node that appears before node_id in topological order.
        fn check_no_forward_ref(
            node_id: MastNodeId,
            child_id: MastNodeId,
        ) -> Result<(), MastForestError> {
            if child_id.0 >= node_id.0 {
                return Err(MastForestError::ForwardReference(node_id, child_id));
            }
            Ok(())
        }

        for (node_idx, node) in self.nodes.iter().enumerate() {
            let node_id = MastNodeId::new_unchecked(node_idx as u32);

            // Check topological ordering and compute expected digest
            let computed_digest = match node {
                MastNode::Block(block) => {
                    let op_groups: Vec<Felt> =
                        block.op_batches().iter().flat_map(|batch| *batch.groups()).collect();
                    hasher::hash_elements(&op_groups)
                },
                MastNode::Join(join) => {
                    let left_id = join.first();
                    let right_id = join.second();
                    check_no_forward_ref(node_id, left_id)?;
                    check_no_forward_ref(node_id, right_id)?;

                    let left_digest = self.nodes[left_id].digest();
                    let right_digest = self.nodes[right_id].digest();
                    hasher::merge_in_domain(&[left_digest, right_digest], JoinNode::DOMAIN)
                },
                MastNode::Split(split) => {
                    let true_id = split.on_true();
                    let false_id = split.on_false();
                    check_no_forward_ref(node_id, true_id)?;
                    check_no_forward_ref(node_id, false_id)?;

                    let true_digest = self.nodes[true_id].digest();
                    let false_digest = self.nodes[false_id].digest();
                    hasher::merge_in_domain(&[true_digest, false_digest], SplitNode::DOMAIN)
                },
                MastNode::Loop(loop_node) => {
                    let body_id = loop_node.body();
                    check_no_forward_ref(node_id, body_id)?;

                    let body_digest = self.nodes[body_id].digest();
                    hasher::merge_in_domain(&[body_digest, Word::default()], LoopNode::DOMAIN)
                },
                MastNode::Call(call) => {
                    let callee_id = call.callee();
                    check_no_forward_ref(node_id, callee_id)?;

                    let callee_digest = self.nodes[callee_id].digest();
                    let domain = if call.is_syscall() {
                        CallNode::SYSCALL_DOMAIN
                    } else {
                        CallNode::CALL_DOMAIN
                    };
                    hasher::merge_in_domain(&[callee_digest, Word::default()], domain)
                },
                MastNode::Dyn(dyn_node) => {
                    if dyn_node.is_dyncall() {
                        DynNode::DYNCALL_DEFAULT_DIGEST
                    } else {
                        DynNode::DYN_DEFAULT_DIGEST
                    }
                },
                MastNode::External(_) => {
                    // External nodes have externally-provided digests that cannot be recomputed
                    continue;
                },
            };

            let stored_digest = node.digest();
            if computed_digest != stored_digest {
                return Err(MastForestError::HashMismatch {
                    node_id,
                    expected: stored_digest,
                    computed: computed_digest,
                });
            }
        }

        Ok(())
    }
}

// ------------------------------------------------------------------------------------------------
/// Error message methods
impl MastForest {
    /// Given an error code as a Felt, resolves it to its corresponding error message.
    pub fn resolve_error_message(&self, code: Felt) -> Option<Arc<str>> {
        let key = code.as_canonical_u64();
        self.debug_info.error_message(key)
    }

    /// Registers an error message in the MAST Forest and returns the corresponding error code as a
    /// Felt.
    pub fn register_error(&mut self, msg: Arc<str>) -> Felt {
        let code: Felt = error_code_from_msg(&msg);
        // we use u64 as keys for the map
        self.debug_info.insert_error_code(code.as_canonical_u64(), msg);
        code
    }
}

// ------------------------------------------------------------------------------------------------
/// Procedure name methods
impl MastForest {
    /// Returns the procedure name for the given MAST root digest, if present.
    pub fn procedure_name(&self, digest: &Word) -> Option<&str> {
        self.debug_info.procedure_name(digest)
    }

    /// Returns an iterator over all (digest, name) pairs of procedure names.
    pub fn procedure_names(&self) -> impl Iterator<Item = (Word, &Arc<str>)> {
        self.debug_info.procedure_names()
    }

    /// Inserts a procedure name for the given MAST root digest.
    pub fn insert_procedure_name(&mut self, digest: Word, name: Arc<str>) {
        assert!(
            self.find_procedure_root(digest).is_some(),
            "attempted to insert procedure name for digest that is not a procedure root"
        );
        self.debug_info.insert_procedure_name(digest, name);
    }

    /// Returns a reference to the debug info for this forest.
    pub fn debug_info(&self) -> &DebugInfo {
        &self.debug_info
    }

    /// Returns a mutable reference to the debug info.
    ///
    /// This is intended for use by the assembler to register AssemblyOps and other debug
    /// information during compilation.
    pub fn debug_info_mut(&mut self) -> &mut DebugInfo {
        &mut self.debug_info
    }
}

// TEST HELPERS
// ================================================================================================

#[cfg(test)]
impl MastForest {
    /// Returns all decorators for a given node as a vector of (position, DecoratorId) tuples.
    ///
    /// This helper method combines before_enter, operation-indexed, and after_exit decorators
    /// into a single collection, which is useful for testing decorator positions and ordering.
    ///
    /// **Performance Warning**: This method performs multiple allocations through collect() calls
    /// and should not be relied upon for performance-critical code. It is intended for testing
    /// only.
    pub fn all_decorators(&self, node_id: MastNodeId) -> Vec<(usize, DecoratorId)> {
        let node = &self[node_id];

        // For non-basic blocks, just get before_enter and after_exit decorators at position 0
        if !node.is_basic_block() {
            let before_enter_decorators: Vec<_> = self
                .before_enter_decorators(node_id)
                .iter()
                .map(|&deco_id| (0, deco_id))
                .collect();

            let after_exit_decorators: Vec<_> = self
                .after_exit_decorators(node_id)
                .iter()
                .map(|&deco_id| (1, deco_id))
                .collect();

            return [before_enter_decorators, after_exit_decorators].concat();
        }

        // For basic blocks, we need to handle operation-indexed decorators with proper positioning
        let block = node.unwrap_basic_block();

        // Before-enter decorators are at position 0
        let before_enter_decorators: Vec<_> = self
            .before_enter_decorators(node_id)
            .iter()
            .map(|&deco_id| (0, deco_id))
            .collect();

        // Operation-indexed decorators with their actual positions
        let op_indexed_decorators: Vec<_> =
            self.decorator_links_for_node(node_id).unwrap().into_iter().collect();

        // After-exit decorators are positioned after all operations
        let after_exit_decorators: Vec<_> = self
            .after_exit_decorators(node_id)
            .iter()
            .map(|&deco_id| (block.num_operations() as usize, deco_id))
            .collect();

        [before_enter_decorators, op_indexed_decorators, after_exit_decorators].concat()
    }
}

// MAST FOREST INDEXING
// ------------------------------------------------------------------------------------------------

impl Index<MastNodeId> for MastForest {
    type Output = MastNode;

    #[inline(always)]
    fn index(&self, node_id: MastNodeId) -> &Self::Output {
        &self.nodes[node_id]
    }
}

impl IndexMut<MastNodeId> for MastForest {
    #[inline(always)]
    fn index_mut(&mut self, node_id: MastNodeId) -> &mut Self::Output {
        &mut self.nodes[node_id]
    }
}

impl Index<DecoratorId> for MastForest {
    type Output = Decorator;

    #[inline(always)]
    fn index(&self, decorator_id: DecoratorId) -> &Self::Output {
        self.debug_info.decorator(decorator_id).expect("DecoratorId out of bounds")
    }
}

impl IndexMut<DecoratorId> for MastForest {
    #[inline(always)]
    fn index_mut(&mut self, decorator_id: DecoratorId) -> &mut Self::Output {
        self.debug_info.decorator_mut(decorator_id).expect("DecoratorId out of bounds")
    }
}

// MAST NODE ID
// ================================================================================================

/// An opaque handle to a [`MastNode`] in some [`MastForest`]. It is the responsibility of the user
/// to use a given [`MastNodeId`] with the corresponding [`MastForest`].
///
/// Note that the [`MastForest`] does *not* ensure that equal [`MastNode`]s have equal
/// [`MastNodeId`] handles. Hence, [`MastNodeId`] equality must not be used to test for equality of
/// the underlying [`MastNode`].
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "serde", serde(transparent))]
#[cfg_attr(all(feature = "arbitrary", test), miden_test_serde_macros::serde_test)]
pub struct MastNodeId(u32);

/// Operations that mutate a MAST often produce this mapping between old and new NodeIds.
pub type Remapping = BTreeMap<MastNodeId, MastNodeId>;

impl MastNodeId {
    /// Returns a new `MastNodeId` with the provided inner value, or an error if the provided
    /// `value` is greater than the number of nodes in the forest.
    ///
    /// For use in deserialization.
    pub fn from_u32_safe(
        value: u32,
        mast_forest: &MastForest,
    ) -> Result<Self, DeserializationError> {
        Self::from_u32_with_node_count(value, mast_forest.nodes.len())
    }

    /// Returns a new [`MastNodeId`] with the provided `node_id`, or an error if `node_id` is
    /// greater than the number of nodes in the [`MastForest`] for which this ID is being
    /// constructed.
    pub fn from_usize_safe(
        node_id: usize,
        mast_forest: &MastForest,
    ) -> Result<Self, DeserializationError> {
        let node_id: u32 = node_id.try_into().map_err(|_| {
            DeserializationError::InvalidValue(format!(
                "node id '{node_id}' does not fit into a u32"
            ))
        })?;
        MastNodeId::from_u32_safe(node_id, mast_forest)
    }

    /// Returns a new [`MastNodeId`] from the given `value` without checking its validity.
    pub fn new_unchecked(value: u32) -> Self {
        Self(value)
    }

    /// Returns a new [`MastNodeId`] with the provided `id`, or an error if `id` is greater or equal
    /// to `node_count`. The `node_count` is the total number of nodes in the [`MastForest`] for
    /// which this ID is being constructed.
    ///
    /// This function can be used when deserializing an id whose corresponding node is not yet in
    /// the forest and [`Self::from_u32_safe`] would fail. For instance, when deserializing the ids
    /// referenced by the Join node in this forest:
    ///
    /// ```text
    /// [Join(1, 2), Block(foo), Block(bar)]
    /// ```
    ///
    /// Since it is less safe than [`Self::from_u32_safe`] and usually not needed it is not public.
    pub(super) fn from_u32_with_node_count(
        id: u32,
        node_count: usize,
    ) -> Result<Self, DeserializationError> {
        if (id as usize) < node_count {
            Ok(Self(id))
        } else {
            Err(DeserializationError::InvalidValue(format!(
                "Invalid deserialized MAST node ID '{id}', but {node_count} is the number of nodes in the forest",
            )))
        }
    }

    /// Remap the NodeId to its new position using the given [`Remapping`].
    pub fn remap(&self, remapping: &Remapping) -> Self {
        *remapping.get(self).unwrap_or(self)
    }
}

impl From<u32> for MastNodeId {
    fn from(value: u32) -> Self {
        MastNodeId::new_unchecked(value)
    }
}

impl Idx for MastNodeId {}

impl From<MastNodeId> for u32 {
    fn from(value: MastNodeId) -> Self {
        value.0
    }
}

impl fmt::Display for MastNodeId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "MastNodeId({})", self.0)
    }
}

#[cfg(any(test, feature = "arbitrary"))]
impl proptest::prelude::Arbitrary for MastNodeId {
    type Parameters = ();

    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
        use proptest::prelude::*;
        any::<u32>().prop_map(MastNodeId).boxed()
    }

    type Strategy = proptest::prelude::BoxedStrategy<Self>;
}

// ITERATOR

/// Iterates over all the nodes a root depends on, in pre-order. The iteration can include other
/// roots in the same forest.
pub struct SubtreeIterator<'a> {
    forest: &'a MastForest,
    discovered: Vec<MastNodeId>,
    unvisited: Vec<MastNodeId>,
}
impl<'a> SubtreeIterator<'a> {
    pub fn new(root: &MastNodeId, forest: &'a MastForest) -> Self {
        let discovered = vec![];
        let unvisited = vec![*root];
        SubtreeIterator { forest, discovered, unvisited }
    }
}
impl Iterator for SubtreeIterator<'_> {
    type Item = MastNodeId;
    fn next(&mut self) -> Option<MastNodeId> {
        while let Some(id) = self.unvisited.pop() {
            let node = &self.forest[id];
            if !node.has_children() {
                return Some(id);
            } else {
                self.discovered.push(id);
                node.append_children_to(&mut self.unvisited);
            }
        }
        self.discovered.pop()
    }
}

// DECORATOR ID
// ================================================================================================

/// An opaque handle to a [`Decorator`] in some [`MastForest`]. It is the responsibility of the user
/// to use a given [`DecoratorId`] with the corresponding [`MastForest`].
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "serde", serde(transparent))]
pub struct DecoratorId(u32);

impl DecoratorId {
    /// Returns a new `DecoratorId` with the provided inner value, or an error if the provided
    /// `value` is greater than the number of nodes in the forest.
    ///
    /// For use in deserialization.
    pub fn from_u32_safe(
        value: u32,
        mast_forest: &MastForest,
    ) -> Result<Self, DeserializationError> {
        Self::from_u32_bounded(value, mast_forest.debug_info.num_decorators())
    }

    /// Returns a new `DecoratorId` with the provided inner value, or an error if the provided
    /// `value` is greater than or equal to `bound`.
    ///
    /// For use in deserialization when the bound is known without needing the full MastForest.
    pub fn from_u32_bounded(value: u32, bound: usize) -> Result<Self, DeserializationError> {
        if (value as usize) < bound {
            Ok(Self(value))
        } else {
            Err(DeserializationError::InvalidValue(format!(
                "Invalid deserialized MAST decorator id '{}', but allows only {} decorators",
                value, bound,
            )))
        }
    }

    /// Creates a new [`DecoratorId`] without checking its validity.
    pub(crate) fn new_unchecked(value: u32) -> Self {
        Self(value)
    }
}

impl From<u32> for DecoratorId {
    fn from(value: u32) -> Self {
        DecoratorId::new_unchecked(value)
    }
}

impl Idx for DecoratorId {}

impl From<DecoratorId> for u32 {
    fn from(value: DecoratorId) -> Self {
        value.0
    }
}

impl fmt::Display for DecoratorId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "DecoratorId({})", self.0)
    }
}

impl Serializable for DecoratorId {
    fn write_into<W: ByteWriter>(&self, target: &mut W) {
        self.0.write_into(target)
    }
}

impl Deserializable for DecoratorId {
    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
        let value = u32::read_from(source)?;
        Ok(Self(value))
    }
}

// ASM OP ID
// ================================================================================================

/// Unique identifier for an [`AssemblyOp`] within a [`MastForest`].
///
/// Unlike decorators (which are executed at runtime), AssemblyOps are metadata
/// used only for error context and debugging tools.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "serde", serde(transparent))]
pub struct AsmOpId(u32);

impl AsmOpId {
    /// Creates a new [`AsmOpId`] with the provided inner value.
    pub const fn new(value: u32) -> Self {
        Self(value)
    }
}

impl From<u32> for AsmOpId {
    fn from(value: u32) -> Self {
        AsmOpId::new(value)
    }
}

impl Idx for AsmOpId {}

impl From<AsmOpId> for u32 {
    fn from(id: AsmOpId) -> Self {
        id.0
    }
}

impl fmt::Display for AsmOpId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "AsmOpId({})", self.0)
    }
}

impl Serializable for AsmOpId {
    fn write_into<W: ByteWriter>(&self, target: &mut W) {
        self.0.write_into(target)
    }
}

impl Deserializable for AsmOpId {
    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
        let value = u32::read_from(source)?;
        Ok(Self(value))
    }
}

/// Derives an error code from an error message by hashing the message and returning the 0th element
/// of the resulting [`Word`].
pub fn error_code_from_msg(msg: impl AsRef<str>) -> Felt {
    // hash the message and return 0th felt of the resulting Word
    hash_string_to_word(msg.as_ref())[0]
}

// MAST FOREST ERROR
// ================================================================================================

/// Represents the types of errors that can occur when dealing with MAST forest.
#[derive(Debug, thiserror::Error, PartialEq)]
pub enum MastForestError {
    #[error("MAST forest decorator count exceeds the maximum of {} decorators", u32::MAX)]
    TooManyDecorators,
    #[error("MAST forest node count exceeds the maximum of {} nodes", MastForest::MAX_NODES)]
    TooManyNodes,
    #[error("node id {0} is greater than or equal to forest length {1}")]
    NodeIdOverflow(MastNodeId, usize),
    #[error("decorator id {0} is greater than or equal to decorator count {1}")]
    DecoratorIdOverflow(DecoratorId, usize),
    #[error("basic block cannot be created from an empty list of operations")]
    EmptyBasicBlock,
    #[error(
        "decorator root of child with node id {0} is missing but is required for fingerprint computation"
    )]
    ChildFingerprintMissing(MastNodeId),
    #[error("advice map key {0} already exists when merging forests")]
    AdviceMapKeyCollisionOnMerge(Word),
    #[error("decorator storage error: {0}")]
    DecoratorError(DecoratorIndexError),
    #[error("digest is required for deserialization")]
    DigestRequiredForDeserialization,
    #[error("invalid batch in basic block node {0:?}: {1}")]
    InvalidBatchPadding(MastNodeId, String),
    #[error("procedure name references digest that is not a procedure root: {0:?}")]
    InvalidProcedureNameDigest(Word),
    #[error(
        "node {0:?} references child {1:?} which comes after it in the forest (forward reference)"
    )]
    ForwardReference(MastNodeId, MastNodeId),
    #[error("hash mismatch for node {node_id:?}: expected {expected:?}, computed {computed:?}")]
    HashMismatch {
        node_id: MastNodeId,
        expected: Word,
        computed: Word,
    },
}

// Custom serde implementations for MastForest that handle linked decorators properly
// by delegating to the existing miden-crypto serialization which already handles
// the conversion between linked and owned decorator formats.
#[cfg(feature = "serde")]
impl serde::Serialize for MastForest {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        // Use the existing miden-crypto serialization which already handles linked decorators
        let bytes = Serializable::to_bytes(self);
        serializer.serialize_bytes(&bytes)
    }
}

#[cfg(feature = "serde")]
impl<'de> serde::Deserialize<'de> for MastForest {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        // Deserialize bytes, then use miden-crypto Deserializable
        let bytes = Vec::<u8>::deserialize(deserializer)?;
        let mut slice_reader = SliceReader::new(&bytes);
        Deserializable::read_from(&mut slice_reader).map_err(serde::de::Error::custom)
    }
}

// UNTRUSTED MAST FOREST
// ================================================================================================

/// A [`MastForest`] deserialized from untrusted input that has not yet been validated.
///
/// This type wraps a `MastForest` that was deserialized from bytes but has not had its
/// node hashes verified. Before using the forest, callers must call [`validate()`](Self::validate)
/// to verify structural integrity and recompute all node hashes.
///
/// # Usage
///
/// ```ignore
/// // Deserialize from untrusted bytes
/// let untrusted = UntrustedMastForest::read_from_bytes(&bytes)?;
///
/// // Validate structure and hashes
/// let forest = untrusted.validate()?;
///
/// // Now safe to use
/// let root = forest.procedure_roots()[0];
/// ```
///
/// # Security
///
/// This type exists to provide type-level safety for untrusted deserialization. The validation
/// performed by [`validate()`](Self::validate) includes:
///
/// 1. **Structural validation**: Checks that basic block batch invariants are satisfied and
///    procedure names reference valid roots.
/// 2. **Topological ordering**: Verifies that all node references point to nodes that appear
///    earlier in the forest (no forward references).
/// 3. **Hash recomputation**: Recomputes the digest for every node and verifies it matches the
///    stored digest.
#[derive(Debug, Clone)]
pub struct UntrustedMastForest(MastForest);

impl UntrustedMastForest {
    /// Validates the forest by checking structural invariants and recomputing all node hashes.
    ///
    /// This method performs a complete validation of the deserialized forest:
    ///
    /// 1. Validates structural invariants (batch padding, procedure names)
    /// 2. Validates topological ordering (no forward references)
    /// 3. Recomputes all node hashes and compares against stored digests
    ///
    /// # Returns
    ///
    /// - `Ok(MastForest)` if validation succeeds
    /// - `Err(MastForestError)` with details about the first validation failure
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - Any basic block has invalid batch structure ([`MastForestError::InvalidBatchPadding`])
    /// - Any procedure name references a non-root digest
    ///   ([`MastForestError::InvalidProcedureNameDigest`])
    /// - Any node references a child that appears later in the forest
    ///   ([`MastForestError::ForwardReference`])
    /// - Any node's recomputed hash doesn't match its stored digest
    ///   ([`MastForestError::HashMismatch`])
    pub fn validate(self) -> Result<MastForest, MastForestError> {
        let forest = self.0;

        // Step 1: Validate structural invariants (existing validate() checks)
        forest.validate()?;

        // Step 2: Validate topological ordering and recompute hashes
        forest.validate_node_hashes()?;

        Ok(forest)
    }

    /// Deserializes an [`UntrustedMastForest`] from bytes.
    ///
    /// This method uses a [`BudgetedReader`] with a budget equal to the input size to protect
    /// against denial-of-service attacks from malicious input.
    ///
    /// For stricter limits, use
    /// [`read_from_bytes_with_budget`](Self::read_from_bytes_with_budget) with a custom budget.
    ///
    /// # Example
    ///
    /// ```ignore
    /// // Read from untrusted source
    /// let untrusted = UntrustedMastForest::read_from_bytes(&bytes)?;
    ///
    /// // Validate before use
    /// let forest = untrusted.validate()?;
    /// ```
    pub fn read_from_bytes(bytes: &[u8]) -> Result<Self, DeserializationError> {
        Self::read_from_bytes_with_budget(bytes, bytes.len())
    }

    /// Deserializes an [`UntrustedMastForest`] from bytes with a byte budget.
    ///
    /// This method uses a [`BudgetedReader`] to limit memory consumption during deserialization,
    /// protecting against denial-of-service attacks from malicious input that claims to contain
    /// an excessive number of elements.
    ///
    /// # Arguments
    ///
    /// * `bytes` - The serialized forest bytes
    /// * `budget` - Maximum bytes to consume during deserialization. Set this to `bytes.len()` for
    ///   typical use cases, or lower to enforce stricter limits.
    ///
    /// # Example
    ///
    /// ```ignore
    /// // Read from untrusted source with budget equal to input size
    /// let untrusted = UntrustedMastForest::read_from_bytes_with_budget(&bytes, bytes.len())?;
    ///
    /// // Validate before use
    /// let forest = untrusted.validate()?;
    /// ```
    ///
    /// # Security
    ///
    /// The budget limits:
    /// - Pre-allocation sizes when deserializing collections (via `max_alloc`)
    /// - Total bytes consumed during deserialization
    ///
    /// This prevents attacks where malicious input claims an unrealistic number of elements
    /// (e.g., `len = 2^60`), causing excessive memory allocation before any data is read.
    pub fn read_from_bytes_with_budget(
        bytes: &[u8],
        budget: usize,
    ) -> Result<Self, DeserializationError> {
        let mut reader = BudgetedReader::new(SliceReader::new(bytes), budget);
        Self::read_from(&mut reader)
    }
}