miden-core 0.22.1

use alloc::{boxed::Box, string::String, vec::Vec};
use core::{fmt, iter::repeat_n};

use crate::{
    Felt, Word, ZERO,
    chiplets::hasher,
    crypto::hash::Blake3_256,
    mast::{
        DecoratedLinksIter, DecoratedOpLink, DecoratorId, DecoratorStore, MastForest,
        MastForestError, MastNode, MastNodeFingerprint, MastNodeId,
    },
    operations::{DecoratorList, Operation},
    prettier::PrettyPrint,
    utils::LookupByIdx,
};

mod op_batch;
pub use op_batch::OpBatch;
use op_batch::OpBatchAccumulator;
pub(crate) use op_batch::collect_immediate_placements;

use super::{MastForestContributor, MastNodeExt};

#[cfg(any(test, feature = "arbitrary"))]
pub mod arbitrary;

#[cfg(test)]
mod tests;

// CONSTANTS
// ================================================================================================

/// Maximum number of operations per group.
pub const GROUP_SIZE: usize = 9;

/// Maximum number of groups per batch.
pub const BATCH_SIZE: usize = 8;
const _: [(); 1] = [(); ((BATCH_SIZE & (BATCH_SIZE - 1)) == 0) as usize];

// BASIC BLOCK NODE
// ================================================================================================

/// Block for a linear sequence of operations (i.e., no branching or loops).
///
/// Executes its operations in order. Fails if any of the operations fails.
///
/// A basic block is composed of operation batches, operation batches are composed of operation
/// groups, operation groups encode the VM's operations and immediate values. These values are
/// created according to these rules:
///
/// - A basic block contains one or more batches.
/// - A batch contains up to 8 groups, and the number of groups must be a power of 2.
/// - A group contains up to 9 operations or 1 immediate value.
/// - Last operation in a group cannot be an operation that requires an immediate value.
/// - NOOPs are used to fill a group or batch when necessary.
/// - An immediate value follows the operation that requires it, using the next available group in
///   the batch. If there are no groups available in the batch, then both the operation and its
///   immediate value are moved to the next batch.
///
/// Example: 8 pushes result in two operation batches:
///
/// - First batch: First group with 7 push opcodes and 2 zero-paddings packed together, followed by
///   7 groups with their respective immediate values.
/// - Second batch: First group with the last push opcode and 8 zero-paddings packed together,
///   followed by one immediate and 6 padding groups.
///
/// The hash of a basic block is:
///
/// > hash(batches, domain=BASIC_BLOCK_DOMAIN)
///
/// Where `batches` is the concatenation of each `batch` in the basic block, and each batch is 8
/// field elements (512 bits).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BasicBlockNode {
    /// The primitive operations contained in this basic block.
    ///
    /// The operations are broken up into batches of 8 groups, with each group containing up to 9
    /// operations, or a single immediates. Thus the maximum size of each batch is 72 operations.
    /// Multiple batches are used for blocks consisting of more than 72 operations.
    op_batches: Vec<OpBatch>,
    digest: Word,
    /// Stores both operation-level and node-level decorators
    /// Custom serialization is handled via Serialize/Deserialize impls
    decorators: DecoratorStore,
}

// ------------------------------------------------------------------------------------------------
// SERIALIZATION
// ================================================================================================

// ------------------------------------------------------------------------------------------------
/// Constants
impl BasicBlockNode {
    /// The domain of the basic block node (used for control block hashing).
    pub const DOMAIN: Felt = ZERO;
}

// ------------------------------------------------------------------------------------------------
/// Constructors
impl BasicBlockNode {
    /// Returns a new [`BasicBlockNode`] instantiated with the specified operations and decorators.
    ///
    /// Returns an error if:
    /// - `operations` vector is empty.
    #[cfg(any(test, feature = "arbitrary"))]
    pub(crate) fn new_owned_with_decorators(
        operations: Vec<Operation>,
        decorators: DecoratorList,
    ) -> Result<Self, MastForestError> {
        if operations.is_empty() {
            return Err(MastForestError::EmptyBasicBlock);
        }

        // Validate decorators list (only in debug mode).
        #[cfg(debug_assertions)]
        validate_decorators(operations.len(), &decorators);

        let (op_batches, digest) = batch_and_hash_ops(operations);
        // the prior line may have inserted some padding Noops in the op_batches
        // the decorator mapping should still point to the correct operation when that happens
        let reflowed_decorators = BasicBlockNode::adjust_decorators(decorators, &op_batches);

        Ok(Self {
            op_batches,
            digest,
            decorators: DecoratorStore::Owned {
                decorators: reflowed_decorators,
                before_enter: Vec::new(),
                after_exit: Vec::new(),
            },
        })
    }

    // Takes a `DecoratorList` which operation indexes are defined against un-padded operations, and
    // adjusts those indexes to point into the padded `&[OpBatches]` passed as argument.
    //
    // IOW this makes its `decorators` padding-aware, or equivalently "adds" the padding to these
    // decorators
    pub fn adjust_decorators(decorators: DecoratorList, op_batches: &[OpBatch]) -> DecoratorList {
        let raw2pad = RawToPaddedPrefix::new(op_batches);
        decorators
            .into_iter()
            .map(|(raw_idx, dec_id)| (raw_idx + raw2pad[raw_idx], dec_id))
            .collect()
    }

    /// Adjusts raw operation indices to padded indices for AssemblyOp mappings.
    ///
    /// Similar to `adjust_decorators`, but works with AssemblyOp mappings `(raw_idx, id)` pairs.
    /// The op_batches contain padding NOOPs that shift operation indices. This method adjusts
    /// the raw indices to account for this padding so lookups during execution use correct indices.
    pub fn adjust_asm_op_indices<T: Copy>(
        asm_ops: Vec<(usize, T)>,
        op_batches: &[OpBatch],
    ) -> Vec<(usize, T)> {
        let raw2pad = RawToPaddedPrefix::new(op_batches);
        asm_ops
            .into_iter()
            .map(|(raw_idx, id)| {
                let padded = raw_idx + raw2pad[raw_idx];
                (padded, id)
            })
            .collect()
    }
}

// ------------------------------------------------------------------------------------------------
/// Public accessors
impl BasicBlockNode {
    /// Returns a reference to the operation batches in this basic block.
    pub fn op_batches(&self) -> &[OpBatch] {
        &self.op_batches
    }

    /// Returns the number of operation batches in this basic block.
    pub fn num_op_batches(&self) -> usize {
        self.op_batches.len()
    }

    /// Returns the total number of operation groups in this basic block.
    ///
    /// Then number of operation groups is computed as follows:
    /// - For all batches but the last one we set the number of groups to 8, regardless of the
    ///   actual number of groups in the batch. The reason for this is that when operation batches
    ///   are concatenated together each batch contributes 8 elements to the hash.
    /// - For the last batch, we take the number of actual groups and round it up to the next power
    ///   of two. The reason for rounding is that the VM always executes a number of operation
    ///   groups which is a power of two.
    pub fn num_op_groups(&self) -> usize {
        let last_batch_num_groups = self.op_batches.last().expect("no last group").num_groups();
        (self.op_batches.len() - 1) * BATCH_SIZE + last_batch_num_groups.next_power_of_two()
    }

    /// Returns the number of operations in this basic block.
    pub fn num_operations(&self) -> u32 {
        let num_ops: usize = self.op_batches.iter().map(|batch| batch.ops().len()).sum();
        num_ops.try_into().expect("basic block contains more than 2^32 operations")
    }

    /// Returns a [`DecoratorOpLinkIterator`] which allows us to iterate through the op-indexed
    /// decorators of this basic block node.
    ///
    /// This method borrows from the forest's storage, avoiding unnecessary Arc clones and providing
    /// efficient access to decorators.
    ///
    /// This iterator is intended for e.g. processor consumption and provides access to only the
    /// operation-indexed decorators (excluding `before_enter` and `after_exit` decorators).
    pub fn indexed_decorator_iter<'a>(
        &'a self,
        forest: &'a MastForest,
    ) -> DecoratorOpLinkIterator<'a> {
        match &self.decorators {
            DecoratorStore::Owned { decorators, .. } => {
                DecoratorOpLinkIterator::from_slice(decorators)
            },
            DecoratorStore::Linked { id } => {
                // This is used in MastForestMerger::merge_nodes, which strips the `MastForest` of
                // some nodes before remapping decorators, so calling
                // verify_node_in_forest will not work here.

                // For linked nodes, borrow from forest storage
                // Check if the node has any decorators at all
                let has_decorators = forest
                    .decorator_links_for_node(*id)
                    .map(|links| links.into_iter().next().is_some())
                    .unwrap_or(false);

                if !has_decorators {
                    return DecoratorOpLinkIterator::from_slice(&[]);
                }

                let view = forest.decorator_links_for_node(*id).expect(
                    "linked node decorators should be available; forest may be inconsistent",
                );

                DecoratorOpLinkIterator::from_linked(view.into_iter())
            },
        }
    }

    /// Returns an iterator which allows us to iterate through the decorator list of
    /// this basic block node with op indexes aligned to the "raw" (un-padded)) op
    /// batches of the basic block node.
    ///
    /// Though this adjusts the indexation of op-indexed decorators, this iterator returns all
    /// decorators of the [`BasicBlockNode`] in the order in which they appear in the program.
    /// This includes `before_enter`, op-indexed decorators, and after_exit.
    ///
    /// Returns an iterator which allows us to iterate through the decorator list of
    /// this basic block node with op indexes aligned to the "raw" (un-padded)) op
    /// batches of the basic block node.
    ///
    /// This method borrows from the forest's storage, avoiding unnecessary Arc clones and
    /// providing efficient access to decorators.
    ///
    /// Though this adjusts the indexation of op-indexed decorators, this iterator returns all
    /// decorators of the [`BasicBlockNode`] in the order in which they appear in the program.
    /// This includes `before_enter`, op-indexed decorators, and after_exit`.
    pub fn raw_decorator_iter<'a>(
        &'a self,
        forest: &'a MastForest,
    ) -> RawDecoratorOpLinkIterator<'a> {
        match &self.decorators {
            DecoratorStore::Owned { decorators, before_enter, after_exit } => {
                RawDecoratorOpLinkIterator::from_slice_iters(
                    before_enter,
                    decorators,
                    after_exit,
                    &self.op_batches,
                )
            },
            DecoratorStore::Linked { id } => {
                #[cfg(debug_assertions)]
                self.verify_node_in_forest(forest);
                // For linked nodes, borrow from forest storage
                // Check if the node has any decorators at all
                let has_decorators = forest
                    .decorator_links_for_node(*id)
                    .map(|links| links.into_iter().next().is_some())
                    .unwrap_or(false);

                if !has_decorators {
                    // No operation-level decorators, but still need node-level decorators
                    let before_enter = forest.before_enter_decorators(*id);
                    let after_exit = forest.after_exit_decorators(*id);
                    return RawDecoratorOpLinkIterator::from_slice_iters(
                        before_enter,
                        &[],
                        after_exit,
                        &self.op_batches,
                    );
                }

                let view = forest.decorator_links_for_node(*id).expect(
                    "linked node decorators should be available; forest may be inconsistent",
                );

                // Get node-level decorators from NodeToDecoratorIds
                let before_enter = forest.before_enter_decorators(*id);
                let after_exit = forest.after_exit_decorators(*id);

                RawDecoratorOpLinkIterator::from_linked(
                    before_enter,
                    view.into_iter(),
                    after_exit,
                    &self.op_batches,
                )
            },
        }
    }

    /// Returns only the raw op-indexed decorators (without before_enter/after_exit)
    /// with indices based on raw operations.
    ///
    /// Stores decorators with raw operation indices for serialization.
    ///
    /// Returns only the raw op-indexed decorators (without before_enter/after_exit)
    /// with indices based on raw operations.
    ///
    /// This method borrows from the forest's storage, avoiding unnecessary Arc clones and
    /// providing efficient access to decorators.
    ///
    /// Stores decorators with raw operation indices for serialization.
    pub fn raw_op_indexed_decorators(&self, forest: &MastForest) -> Vec<(usize, DecoratorId)> {
        match &self.decorators {
            DecoratorStore::Owned { decorators, .. } => {
                RawDecoratorOpLinkIterator::from_slice_iters(&[], decorators, &[], &self.op_batches)
                    .collect()
            },
            DecoratorStore::Linked { id } => {
                // This is used in MastForest::remove_nodes, which strips the `MastForest` of its
                // nodes before remapping decorators, so calling
                // verify_node_in_forest will not work here.

                let pad2raw = PaddedToRawPrefix::new(self.op_batches());
                match forest.decorator_links_for_node(*id) {
                    Ok(links) => links
                        .into_iter()
                        .map(|(padded_idx, dec_id)| {
                            let raw_idx = padded_idx - pad2raw[padded_idx];
                            (raw_idx, dec_id)
                        })
                        .collect(),
                    Err(_) => Vec::new(), // Return empty if error
                }
            },
        }
    }

    /// Returns an iterator over the operations in the order in which they appear in the program.
    pub fn operations(&self) -> impl Iterator<Item = &Operation> {
        self.op_batches.iter().flat_map(|batch| batch.ops())
    }

    /// Returns an iterator over the un-padded operations in the order in which they
    /// appear in the program.
    pub fn raw_operations(&self) -> impl Iterator<Item = &Operation> {
        self.op_batches.iter().flat_map(|batch| batch.raw_ops())
    }

    /// Returns the total number of operations and decorators in this basic block.
    pub fn num_operations_and_decorators(&self, forest: &MastForest) -> u32 {
        let num_ops: usize = self.num_operations() as usize;
        let num_decorators = match &self.decorators {
            DecoratorStore::Owned { decorators, .. } => decorators.len(),
            DecoratorStore::Linked { id } => {
                #[cfg(debug_assertions)]
                self.verify_node_in_forest(forest);
                // For linked nodes, count from forest storage
                forest
                    .decorator_links_for_node(*id)
                    .map(|links| links.into_iter().count())
                    .unwrap_or(0)
            },
        };

        (num_ops + num_decorators)
            .try_into()
            .expect("basic block contains more than 2^32 operations and decorators")
    }

    /// Returns an iterator over all operations and decorator, in the order in which they appear in
    /// the program.
    ///
    /// This method requires access to the forest to properly handle linked nodes.
    fn iter<'a>(
        &'a self,
        forest: &'a MastForest,
    ) -> impl Iterator<Item = OperationOrDecorator<'a>> + 'a {
        OperationOrDecoratorIterator::new_with_forest(self, forest)
    }

    /// Performs semantic equality comparison with another BasicBlockNode.
    ///
    /// This method compares two blocks for logical equality by comparing:
    /// - Operations (exact equality)
    /// - Before-enter decorators (by ID)
    /// - After-exit decorators (by ID)
    /// - Operation-indexed decorators (by iterating and comparing their contents)
    ///
    /// Unlike the derived PartialEq, this method works correctly with both owned and linked
    /// decorator storage by accessing the actual decorator data from the forest when needed.
    #[cfg(test)]
    pub fn semantic_eq(&self, other: &BasicBlockNode, forest: &MastForest) -> bool {
        // Compare operations by collecting and comparing
        let self_ops: Vec<_> = self.operations().collect();
        let other_ops: Vec<_> = other.operations().collect();
        if self_ops != other_ops {
            return false;
        }

        // Compare before-enter decorators
        if self.before_enter(forest) != other.before_enter(forest) {
            return false;
        }

        // Compare after-exit decorators
        if self.after_exit(forest) != other.after_exit(forest) {
            return false;
        }

        // Compare operation-indexed decorators by collecting and comparing
        let self_decorators: Vec<_> = self.indexed_decorator_iter(forest).collect();
        let other_decorators: Vec<_> = other.indexed_decorator_iter(forest).collect();

        if self_decorators != other_decorators {
            return false;
        }

        true
    }

    /// Return the MastNodeId of this `BasicBlockNode`, if in `Linked` state
    pub fn linked_id(&self) -> Option<MastNodeId> {
        self.decorators.linked_id()
    }
}

// BATCH VALIDATION
// ================================================================================================

impl BasicBlockNode {
    /// Validates that this BasicBlockNode satisfies the core invariants:
    /// 1. Non-final batches must be full (BATCH_SIZE groups), final batch must be power-of-two
    /// 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)
    /// 5. Immediate values are committed to empty groups and match group contents
    /// 6. OpBatch padding semantics (no padding on empty groups; padded groups end with NOOP)
    ///
    /// Returns an error string describing which invariant was violated if validation fails.
    pub fn validate_batch_invariants(&self) -> Result<(), String> {
        // Check invariant 1: Power-of-two groups in each batch
        self.validate_power_of_two_groups()?;

        // Check invariant 4: OpBatch structural consistency
        // This needs to be done early on as it will validate indptr indexes used in later checks.
        self.validate_batch_structure()?;

        // Control-flow opcodes are expected to be filtered upstream and enforced centrally via
        // MastForest::validate.

        // Check invariants 2 and 3: immediate-ending constraints
        self.validate_no_immediate_endings()?;

        // Check invariant 5: Immediate values must be committed to empty groups
        self.validate_immediate_commitment()?;

        // Check invariant 6: OpBatch padding semantics
        self.validate_padding_semantics()?;

        Ok(())
    }

    /// Validates that non-final batches are full and the final batch is power-of-two.
    ///
    /// This invariant is required by trace generation (see `num_op_groups`) and is expected to
    /// hold for all serialized forests produced by the assembler; violations indicate corrupted
    /// or malformed input.
    fn validate_power_of_two_groups(&self) -> Result<(), String> {
        for (batch_idx, batch) in self.op_batches.iter().enumerate() {
            let num_groups = batch.num_groups();
            if batch_idx + 1 < self.op_batches.len() {
                if num_groups != BATCH_SIZE {
                    return Err(format!(
                        "Batch {}: {} groups is not full batch size {}",
                        batch_idx, num_groups, BATCH_SIZE
                    ));
                }
            } else if !num_groups.is_power_of_two() {
                return Err(format!(
                    "Batch {}: {} groups is not power of two",
                    batch_idx, num_groups
                ));
            }
        }
        Ok(())
    }

    /// Validates that no operation group ends with an operation that has an immediate value.
    /// Also validates that the last operation group in a batch cannot contain operations
    /// requiring immediate values.
    fn validate_no_immediate_endings(&self) -> Result<(), String> {
        for (batch_idx, batch) in self.op_batches.iter().enumerate() {
            let num_groups = batch.num_groups();
            let indptr = batch.indptr();
            let ops = batch.ops();

            // Check each group in the batch
            for group_idx in 0..num_groups {
                let group_start = indptr[group_idx];
                let group_end = indptr[group_idx + 1];

                // Skip empty groups (they contain immediate values, not operations)
                if group_start == group_end {
                    continue;
                }

                let group_ops = &ops[group_start..group_end];

                // Check if this is the last group in the batch
                let is_last_group = group_idx == num_groups - 1;

                if is_last_group {
                    // Last group in a batch cannot contain ANY operations requiring immediate
                    // values
                    for (op_idx, op) in group_ops.iter().enumerate() {
                        if op.imm_value().is_some() {
                            return Err(format!(
                                "Batch {}, group {}: operation at index {} requires immediate value, but this is the last group in batch",
                                batch_idx, group_idx, op_idx
                            ));
                        }
                    }
                } else {
                    // Non-last groups: check that the last operation doesn't require an immediate
                    if let Some(last_op) = group_ops.last()
                        && last_op.imm_value().is_some()
                    {
                        return Err(format!(
                            "Batch {}, group {}: ends with operation requiring immediate value",
                            batch_idx, group_idx
                        ));
                    }
                }
            }
        }
        Ok(())
    }

    /// Validates that OpBatch structure is consistent and won't cause panics during access.
    /// Checks:
    /// - num_groups <= BATCH_SIZE
    /// - indptr array is monotonic non-decreasing
    /// - indptr values are within ops bounds
    /// - each group has at most GROUP_SIZE operations
    fn validate_batch_structure(&self) -> Result<(), String> {
        for (batch_idx, batch) in self.op_batches.iter().enumerate() {
            // Check num_groups is within bounds
            if batch.num_groups() > BATCH_SIZE {
                return Err(format!(
                    "Batch {}: num_groups {} exceeds maximum {}",
                    batch_idx,
                    batch.num_groups(),
                    BATCH_SIZE
                ));
            }

            // Check indptr array consistency
            let indptr = batch.indptr();
            let ops = batch.ops();

            // Full array must be monotonic for serialization (delta encoding)
            for i in 0..indptr.len() - 1 {
                if indptr[i] > indptr[i + 1] {
                    return Err(format!(
                        "Batch {}: indptr[{}] {} > indptr[{}] {} - full array not monotonic (required for serialization)",
                        batch_idx,
                        i,
                        indptr[i],
                        i + 1,
                        indptr[i + 1]
                    ));
                }
            }

            let ops_len = ops.len();
            if indptr[indptr.len() - 1] != ops_len {
                return Err(format!(
                    "Batch {}: final indptr value {} doesn't match ops.len() {}",
                    batch_idx,
                    indptr[indptr.len() - 1],
                    ops_len
                ));
            }

            // Check that each group has at most GROUP_SIZE operations
            for group_idx in 0..batch.num_groups() {
                let group_start = indptr[group_idx];
                let group_end = indptr[group_idx + 1];
                let group_size = group_end - group_start;

                if group_size > GROUP_SIZE {
                    return Err(format!(
                        "Batch {}, group {}: contains {} operations, exceeds maximum {}",
                        batch_idx, group_idx, group_size, GROUP_SIZE
                    ));
                }
            }
        }
        Ok(())
    }

    /// Validates that immediate values are committed to empty groups and match group contents.
    /// Checks:
    /// - operation group encodings match committed group values
    /// - each immediate maps to an empty group slot
    /// - immediate group values equal the push immediate
    /// - immediate placement does not exceed num_groups or batch size
    fn validate_immediate_commitment(&self) -> Result<(), String> {
        for (batch_idx, batch) in self.op_batches.iter().enumerate() {
            let num_groups = batch.num_groups();
            let indptr = batch.indptr();
            let ops = batch.ops();
            let groups = batch.groups();

            let mut immediate_slots = [false; BATCH_SIZE];

            for group_idx in 0..num_groups {
                let group_start = indptr[group_idx];
                let group_end = indptr[group_idx + 1];

                if group_start == group_end {
                    continue;
                }

                let mut group_value: u64 = 0;
                for (local_op_idx, op) in ops[group_start..group_end].iter().enumerate() {
                    let opcode = op.op_code() as u64;
                    group_value |= opcode << (Operation::OP_BITS * local_op_idx);
                }
                if groups[group_idx] != Felt::new(group_value) {
                    return Err(format!(
                        "Batch {}, group {}: committed opcode group does not match operations",
                        batch_idx, group_idx
                    ));
                }

                let (placements, _next_group_idx) = collect_immediate_placements(
                    ops,
                    indptr,
                    group_idx,
                    BATCH_SIZE,
                    Some(num_groups),
                )
                .map_err(|err| format!("Batch {}: {}", batch_idx, err))?;

                for (imm_group_idx, imm_value) in placements {
                    if groups[imm_group_idx] != imm_value {
                        return Err(format!(
                            "Batch {}: push immediate value mismatch at index {}",
                            batch_idx, imm_group_idx
                        ));
                    }
                    immediate_slots[imm_group_idx] = true;
                }
            }

            for group_idx in 0..num_groups {
                if indptr[group_idx] == indptr[group_idx + 1]
                    && !immediate_slots[group_idx]
                    && groups[group_idx] != ZERO
                {
                    return Err(format!(
                        "Batch {}, group {}: empty group must be zero",
                        batch_idx, group_idx
                    ));
                }
            }
        }

        Ok(())
    }

    /// Validates that padding metadata matches batch contents.
    /// - Empty groups cannot be marked as padded.
    /// - Padded groups must end with a NOOP operation.
    fn validate_padding_semantics(&self) -> Result<(), String> {
        for (batch_idx, batch) in self.op_batches.iter().enumerate() {
            batch
                .validate_padding_semantics()
                .map_err(|err| format!("Batch {}: {}", batch_idx, err))?;
        }

        Ok(())
    }
}

// PRETTY PRINTING
// ================================================================================================

impl BasicBlockNode {
    pub(super) fn to_display<'a>(&'a self, mast_forest: &'a MastForest) -> impl fmt::Display + 'a {
        BasicBlockNodePrettyPrint { block_node: self, mast_forest }
    }

    pub(super) fn to_pretty_print<'a>(
        &'a self,
        mast_forest: &'a MastForest,
    ) -> impl PrettyPrint + 'a {
        BasicBlockNodePrettyPrint { block_node: self, mast_forest }
    }
}

// MAST NODE TRAIT IMPLEMENTATION
// ================================================================================================

impl MastNodeExt for BasicBlockNode {
    /// Returns a commitment to this basic block.
    fn digest(&self) -> Word {
        self.digest
    }

    fn before_enter<'a>(&'a self, forest: &'a MastForest) -> &'a [DecoratorId] {
        match &self.decorators {
            DecoratorStore::Owned { before_enter, .. } => before_enter,
            DecoratorStore::Linked { id } => {
                // For linked nodes, get the decorators from the forest's NodeToDecoratorIds
                #[cfg(debug_assertions)]
                self.verify_node_in_forest(forest);
                forest.before_enter_decorators(*id)
            },
        }
    }

    fn after_exit<'a>(&'a self, forest: &'a MastForest) -> &'a [DecoratorId] {
        match &self.decorators {
            DecoratorStore::Owned { after_exit, .. } => after_exit,
            DecoratorStore::Linked { id } => {
                // For linked nodes, get the decorators from the forest's NodeToDecoratorIds
                #[cfg(debug_assertions)]
                self.verify_node_in_forest(forest);
                forest.after_exit_decorators(*id)
            },
        }
    }

    fn to_display<'a>(&'a self, mast_forest: &'a MastForest) -> Box<dyn fmt::Display + 'a> {
        Box::new(BasicBlockNode::to_display(self, mast_forest))
    }

    fn to_pretty_print<'a>(&'a self, mast_forest: &'a MastForest) -> Box<dyn PrettyPrint + 'a> {
        Box::new(BasicBlockNode::to_pretty_print(self, mast_forest))
    }

    fn has_children(&self) -> bool {
        false
    }

    fn append_children_to(&self, _target: &mut Vec<MastNodeId>) {
        // No children for basic blocks
    }

    fn for_each_child<F>(&self, _f: F)
    where
        F: FnMut(MastNodeId),
    {
        // BasicBlockNode has no children
    }

    fn domain(&self) -> Felt {
        Self::DOMAIN
    }

    type Builder = BasicBlockNodeBuilder;

    fn to_builder(self, forest: &MastForest) -> Self::Builder {
        // Extract padded decorators and before_enter/after_exit based on storage type
        let (padded_decorators, before_enter, after_exit) = match self.decorators {
            DecoratorStore::Owned { decorators, before_enter, after_exit } => {
                // Decorators are already padded in Owned storage
                (decorators, before_enter, after_exit)
            },
            DecoratorStore::Linked { id } => {
                // For linked nodes, get decorators from forest's centralized storage
                // The decorators are already padded in the centralized storage
                let padded_decorators: DecoratorList = forest
                    .debug_info
                    .decorator_links_for_node(id)
                    .expect("node must exist in forest")
                    .into_iter()
                    .collect();
                let before_enter = forest.before_enter_decorators(id).to_vec();
                let after_exit = forest.after_exit_decorators(id).to_vec();
                (padded_decorators, before_enter, after_exit)
            },
        };

        // Use from_op_batches to avoid re-batching and re-adjusting decorators
        BasicBlockNodeBuilder::from_op_batches(self.op_batches, padded_decorators, self.digest)
            .with_before_enter(before_enter)
            .with_after_exit(after_exit)
    }

    #[cfg(debug_assertions)]
    fn verify_node_in_forest(&self, forest: &MastForest) {
        if let DecoratorStore::Linked { id } = &self.decorators {
            // Verify that this node is the one stored at the given ID in the forest
            let self_ptr = self as *const Self;
            let forest_node = &forest.nodes[*id];
            let forest_node_ptr = match forest_node {
                MastNode::Block(block_node) => block_node as *const BasicBlockNode as *const (),
                _ => panic!("Node type mismatch at {:?}", id),
            };
            let self_as_void = self_ptr as *const ();
            debug_assert_eq!(
                self_as_void, forest_node_ptr,
                "Node pointer mismatch: expected node at {:?} to be self",
                id
            );
        }
    }
}

struct BasicBlockNodePrettyPrint<'a> {
    block_node: &'a BasicBlockNode,
    mast_forest: &'a MastForest,
}

impl PrettyPrint for BasicBlockNodePrettyPrint<'_> {
    #[rustfmt::skip]
    fn render(&self) -> crate::prettier::Document {
        use crate::prettier::*;

        // e.g. `basic_block a b c end`
        let single_line = const_text("basic_block")
            + const_text(" ")
            + self.
                block_node
                .iter(self.mast_forest)
                .map(|op_or_dec| match op_or_dec {
                    OperationOrDecorator::Operation(op) => op.render(),
                    OperationOrDecorator::Decorator(decorator_id) => {
                        self.mast_forest.decorator_by_id(decorator_id)
                            .map(|decorator| decorator.render())
                            .unwrap_or_else(|| const_text("<invalid_decorator_id>"))
                    },
                })
                .reduce(|acc, doc| acc + const_text(" ") + doc)
                .unwrap_or_default()
            + const_text(" ")
            + const_text("end");

        // e.g. `
        // basic_block
        //     a
        //     b
        //     c
        // end
        // `

        let multi_line = indent(
            4,
            const_text("basic_block")
                + nl()
                + self
                    .block_node
                    .iter(self.mast_forest)
                    .map(|op_or_dec| match op_or_dec {
                        OperationOrDecorator::Operation(op) => op.render(),
                        OperationOrDecorator::Decorator(decorator_id) => {
                            self.mast_forest.decorator_by_id(decorator_id)
                                .map(|decorator| decorator.render())
                                .unwrap_or_else(|| const_text("<invalid_decorator_id>"))
                        },
                    })
                    .reduce(|acc, doc| acc + nl() + doc)
                    .unwrap_or_default(),
        ) + nl()
            + const_text("end");

        single_line | multi_line
    }
}

impl fmt::Display for BasicBlockNodePrettyPrint<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.pretty_print(f)
    }
}

enum OpIndexed<'a> {
    Slice(core::slice::Iter<'a, (usize, DecoratorId)>),
    Linked(DecoratedLinksIter<'a>),
}

// DECORATOR ITERATION
// ================================================================================================

/// Iterator used to iterate through the op-indexed decorators of a basic block.
///
/// This lets the caller iterate through operation-indexed decorators with indexes that match the
/// standard (padded) representation of a basic block.
pub struct DecoratorOpLinkIterator<'a>(OpIndexed<'a>);

impl<'a> DecoratorOpLinkIterator<'a> {
    /// Create a new iterator from a slice of decorator links.
    pub fn from_slice(decorators: &'a [DecoratedOpLink]) -> Self {
        Self(OpIndexed::Slice(decorators.iter()))
    }

    /// Create a new iterator from a linked decorator iterator.
    pub fn from_linked(decorators: DecoratedLinksIter<'a>) -> Self {
        Self(OpIndexed::Linked(decorators.into_iter()))
    }
}

impl<'a> Iterator for DecoratorOpLinkIterator<'a> {
    type Item = (usize, DecoratorId);

    fn next(&mut self) -> Option<Self::Item> {
        match &mut self.0 {
            OpIndexed::Slice(slice_iter) => slice_iter.next().copied(),
            OpIndexed::Linked(linked_iter) => linked_iter.next(),
        }
    }
}

impl<'a> ExactSizeIterator for DecoratorOpLinkIterator<'a> {
    #[inline]
    fn len(&self) -> usize {
        match &self.0 {
            OpIndexed::Slice(slice_iter) => slice_iter.len(),
            OpIndexed::Linked(linked_iter) => linked_iter.len(),
        }
    }
}

// Driver of the Iterators' state machine (used by other iterators)
enum Segment {
    Before,
    Middle,
    After,
    Done,
}

// RAW DECORATOR ITERATION
// ================================================================================================

/// Iterator used to iterate through the decorator list of a span block
/// while executing operation batches of a span block.
///
/// This lets the caller iterate through a Decorator list with indexes that match the
/// raw (unpadded) representation of a basic block.
///
/// IOW this makes its `BasicBlockNode::raw_decorator_iter` padding-unaware, or equivalently
/// "removes" the padding of these decorators
pub struct RawDecoratorOpLinkIterator<'a> {
    before: core::slice::Iter<'a, DecoratorId>,
    middle: RawMid<'a>,
    after: core::slice::Iter<'a, DecoratorId>,
    pad2raw: PaddedToRawPrefix, // indexed by padded indices
    total_raw_ops: usize,       // count of raw ops
    seg: Segment,
}

enum RawMid<'a> {
    Slice(core::slice::Iter<'a, (usize, DecoratorId)>),
    Linked(DecoratedLinksIter<'a>),
}

impl<'a> RawDecoratorOpLinkIterator<'a> {
    pub fn from_slice_iters(
        before_enter: &'a [DecoratorId],
        decorators: &'a [(usize, DecoratorId)], // contains adjusted indices
        after_exit: &'a [DecoratorId],
        op_batches: &'a [OpBatch],
    ) -> Self {
        let pad2raw = PaddedToRawPrefix::new(op_batches);
        let raw2pad = RawToPaddedPrefix::new(op_batches);
        let total_raw_ops = raw2pad.raw_ops();

        Self {
            before: before_enter.iter(),
            middle: RawMid::Slice(decorators.iter()),
            after: after_exit.iter(),
            pad2raw,
            total_raw_ops,
            seg: Segment::Before,
        }
    }

    pub fn from_linked(
        before_enter: &'a [DecoratorId],
        decorators: DecoratedLinksIter<'a>,
        after_exit: &'a [DecoratorId],
        op_batches: &'a [OpBatch],
    ) -> Self {
        let pad2raw = PaddedToRawPrefix::new(op_batches);
        let raw2pad = RawToPaddedPrefix::new(op_batches);
        let total_raw_ops = raw2pad.raw_ops();

        Self {
            before: before_enter.iter(),
            middle: RawMid::Linked(decorators.into_iter()),
            after: after_exit.iter(),
            pad2raw,
            total_raw_ops,
            seg: Segment::Before,
        }
    }

    fn middle_next(&mut self) -> Option<(usize, DecoratorId)> {
        match &mut self.middle {
            RawMid::Slice(slice_iter) => slice_iter.next().copied(),
            RawMid::Linked(linked_iter) => linked_iter.next(),
        }
    }
}

impl<'a> Iterator for RawDecoratorOpLinkIterator<'a> {
    type Item = (usize, DecoratorId);

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            match self.seg {
                Segment::Before => {
                    if let Some(&id) = self.before.next() {
                        return Some((0, id));
                    }
                    self.seg = Segment::Middle;
                },
                Segment::Middle => {
                    if let Some((padded_idx, id)) = self.middle_next() {
                        let raw_idx = padded_idx - self.pad2raw[padded_idx];
                        return Some((raw_idx, id));
                    }
                    self.seg = Segment::After;
                },
                Segment::After => {
                    if let Some(&id) = self.after.next() {
                        // After-exit decorators attach to the sentinel raw index
                        return Some((self.total_raw_ops, id));
                    }
                    self.seg = Segment::Done;
                },
                Segment::Done => return None,
            }
        }
    }
}

// OPERATION OR DECORATOR
// ================================================================================================

/// Encodes either an [`Operation`] or a [`crate::operations::Decorator`].
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum OperationOrDecorator<'a> {
    Operation(&'a Operation),
    Decorator(DecoratorId),
}

struct OperationOrDecoratorIterator<'a> {
    node: &'a BasicBlockNode,
    forest: Option<&'a MastForest>,

    // extra segments
    before: core::slice::Iter<'a, DecoratorId>,
    after: core::slice::Iter<'a, DecoratorId>,

    // operation traversal
    batch_index: usize,
    op_index_in_batch: usize,
    op_index: usize, // across all batches

    // decorators inside the block (sorted by op index)
    decorator_list_next_index: usize,
    seg: Segment,
}

impl<'a> OperationOrDecoratorIterator<'a> {
    fn new_with_forest(node: &'a BasicBlockNode, forest: &'a MastForest) -> Self {
        Self {
            node,
            forest: Some(forest),
            before: node.before_enter(forest).iter(),
            after: node.after_exit(forest).iter(),
            batch_index: 0,
            op_index_in_batch: 0,
            op_index: 0,
            decorator_list_next_index: 0,
            seg: Segment::Before,
        }
    }

    #[inline]
    fn next_decorator_if_due(&mut self) -> Option<OperationOrDecorator<'a>> {
        match &self.node.decorators {
            DecoratorStore::Owned { decorators, .. } => {
                // Simple case for owned decorators - use index lookup
                if let Some((op_idx, deco)) = decorators.get(self.decorator_list_next_index)
                    && *op_idx == self.op_index
                {
                    self.decorator_list_next_index += 1;
                    Some(OperationOrDecorator::Decorator(*deco))
                } else {
                    None
                }
            },
            DecoratorStore::Linked { id } => {
                // For linked nodes, use forest access if available
                if let Some(forest) = self.forest {
                    // Get decorators for the current operation from the forest
                    let decorator_ids = forest.decorator_indices_for_op(*id, self.op_index);

                    if self.decorator_list_next_index < decorator_ids.len() {
                        let decorator_id = decorator_ids[self.decorator_list_next_index];
                        self.decorator_list_next_index += 1;
                        Some(OperationOrDecorator::Decorator(decorator_id))
                    } else {
                        None
                    }
                } else {
                    // No forest access available, can't retrieve decorators
                    None
                }
            },
        }
    }
}

impl<'a> Iterator for OperationOrDecoratorIterator<'a> {
    type Item = OperationOrDecorator<'a>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            match self.seg {
                Segment::Before => {
                    if let Some(id) = self.before.next() {
                        return Some(OperationOrDecorator::Decorator(*id));
                    }
                    self.seg = Segment::Middle;
                },

                Segment::Middle => {
                    // 1) emit any decorators for the current op_index
                    if let Some(d) = self.next_decorator_if_due() {
                        return Some(d);
                    }

                    // 2) otherwise emit the operation at current indices
                    if let Some(batch) = self.node.op_batches.get(self.batch_index) {
                        if let Some(op) = batch.ops.get(self.op_index_in_batch) {
                            self.op_index_in_batch += 1;
                            self.op_index += 1;
                            // Reset decorator index when moving to a new operation
                            self.decorator_list_next_index = 0;
                            return Some(OperationOrDecorator::Operation(op));
                        }
                        // advance to next batch and retry
                        self.batch_index += 1;
                        self.op_index_in_batch = 0;
                    } else {
                        // no more ops, decorators flushed through the operation index
                        // and next_decorator_if_due
                        self.seg = Segment::After;
                    }
                },

                Segment::After => {
                    if let Some(id) = self.after.next() {
                        return Some(OperationOrDecorator::Decorator(*id));
                    }
                    self.seg = Segment::Done;
                },

                Segment::Done => return None,
            }
        }
    }
}

// HELPER FUNCTIONS
// ================================================================================================

/// Checks if a given decorators list is valid (only checked in debug mode)
/// - Assert the decorator list is in ascending order.
/// - Assert the last op index in decorator list is less than or equal to the number of operations.
#[cfg(debug_assertions)]
pub(crate) fn validate_decorators(operations_len: usize, decorators: &DecoratorList) {
    if !decorators.is_empty() {
        // check if decorator list is sorted
        for i in 0..(decorators.len() - 1) {
            debug_assert!(decorators[i + 1].0 >= decorators[i].0, "unsorted decorators list");
        }
        // assert the last index in decorator list is less than or equal to operations vector length
        debug_assert!(
            operations_len >= decorators.last().expect("empty decorators list").0,
            "last op index in decorator list should be less than or equal to the number of ops"
        );
    }
}

/// Raw-indexed prefix: how many paddings strictly before raw index r
///
/// This struct provides O(1) lookup for converting raw operation indices to padded indices.
/// For any raw index r, `raw_to_padded[r] = count of padding ops strictly before raw index r`.
///
/// Length: `raw_ops + 1` (includes sentinel entry at `r == raw_ops`)
/// Usage: `padded_idx = r + raw_to_padded[r]` (addition)
#[derive(Debug, Clone)]
pub struct RawToPaddedPrefix(Vec<usize>);

impl RawToPaddedPrefix {
    /// Build a raw-indexed prefix array from op batches.
    ///
    /// For each raw index r, records how many padding operations have been inserted before r.
    /// Includes a sentinel entry at `r == raw_ops`.
    pub fn new(op_batches: &[OpBatch]) -> Self {
        let mut v = Vec::new();
        let mut pads_so_far = 0usize;

        for b in op_batches {
            let n = b.num_groups();
            let indptr = b.indptr();
            let padding = b.padding();

            for g in 0..n {
                let group_len = indptr[g + 1] - indptr[g];
                let has_pad = padding[g] as usize;
                let raw_in_g = group_len - has_pad;

                // For each raw op, record how many paddings were before it.
                v.extend(repeat_n(pads_so_far, raw_in_g));

                // After the group's raw ops, account for the (optional) padding op.
                pads_so_far += has_pad; // adds 1 if there is a padding, else 0
            }
        }

        // Sentinel for r == raw_ops
        v.push(pads_so_far);
        RawToPaddedPrefix(v)
    }

    /// Get the total number of raw operations (excluding sentinel).
    #[inline]
    pub fn raw_ops(&self) -> usize {
        self.0.len() - 1
    }
}

/// Get the number of padding operations before raw index r.
///
/// ## Sentinel Access
///
/// Some decorators have an operation index equal to the length of the
/// operations array, to ensure they are executed at the end of the block
/// (since the semantics of the decorator index is that it must be executed
/// before the operation index it points to).
impl core::ops::Index<usize> for RawToPaddedPrefix {
    type Output = usize;
    #[inline]
    fn index(&self, idx: usize) -> &Self::Output {
        &self.0[idx]
    }
}

/// Padded-indexed prefix: how many paddings strictly before padded index p
///
/// This struct provides O(1) lookup for converting padded operation indices to raw indices.
/// For any padded index p, `padded_to_raw[p] = count of padding ops strictly before padded index
/// p`.
///
/// Length: `padded_ops + 1` (includes sentinel entry at `p == padded_ops`)
/// Usage: `raw_idx = p - padded_to_raw[p]` (subtraction)
#[derive(Debug, Clone)]
pub struct PaddedToRawPrefix(Vec<usize>);

impl PaddedToRawPrefix {
    /// Build a padded-indexed prefix array from op batches.
    ///
    /// Simulates emission of the padded sequence, recording padding count before each position.
    /// Includes a sentinel entry at `p == padded_ops`.
    pub fn new(op_batches: &[OpBatch]) -> Self {
        // Exact capacity to avoid reallocations: sum of per-group lengths across all batches.
        let padded_ops = op_batches
            .iter()
            .map(|b| {
                let n = b.num_groups();
                let indptr = b.indptr();
                indptr[1..=n]
                    .iter()
                    .zip(&indptr[..n])
                    .map(|(end, start)| end - start)
                    .sum::<usize>()
            })
            .sum::<usize>();

        let mut v = Vec::with_capacity(padded_ops + 1);
        let mut pads_so_far = 0usize;

        for b in op_batches {
            let n = b.num_groups();
            let indptr = b.indptr();
            let padding = b.padding();

            for g in 0..n {
                let group_len = indptr[g + 1] - indptr[g];
                let has_pad = padding[g] as usize;
                let raw_in_g = group_len - has_pad;

                // Emit raw ops of the group.
                v.extend(repeat_n(pads_so_far, raw_in_g));

                // Emit the optional padding op.
                if has_pad == 1 {
                    v.push(pads_so_far);
                    pads_so_far += 1; // subsequent positions see one more padding before them
                }
            }
        }

        // Sentinel at p == padded_ops
        v.push(pads_so_far);

        PaddedToRawPrefix(v)
    }
}

/// Get the number of padding operations before padded index p.
///
/// ## Sentinel Access
///
/// Some decorators have an operation index equal to the length of the
/// operations array, to ensure they are executed at the end of the block
/// (since the semantics of the decorator index is that it must be executed
/// before the operation index it points to).
impl core::ops::Index<usize> for PaddedToRawPrefix {
    type Output = usize;
    #[inline]
    fn index(&self, idx: usize) -> &Self::Output {
        &self.0[idx]
    }
}

/// Groups the provided operations into batches and computes the hash of the block.
fn batch_and_hash_ops(ops: Vec<Operation>) -> (Vec<OpBatch>, Word) {
    // Group the operations into batches.
    let batches = batch_ops(ops);

    // Compute the hash of all operation groups.
    let op_groups: Vec<Felt> = batches.iter().flat_map(|batch| batch.groups).collect();
    let hash = hasher::hash_elements(&op_groups);

    (batches, hash)
}

/// Groups the provided operations into batches as described in the docs for this module (i.e., up
/// to 9 operations per group, and 8 groups per batch).
fn batch_ops(ops: Vec<Operation>) -> Vec<OpBatch> {
    let mut batches = Vec::<OpBatch>::new();
    let mut batch_acc = OpBatchAccumulator::new();

    for op in ops {
        // If the operation cannot be accepted into the current accumulator, add the contents of
        // the accumulator to the list of batches and start a new accumulator.
        if !batch_acc.can_accept_op(op) {
            let batch = batch_acc.into_batch();
            batch_acc = OpBatchAccumulator::new();

            batches.push(batch);
        }

        // Add the operation to the accumulator.
        batch_acc.add_op(op);
    }

    // Make sure we finished processing the last batch.
    if !batch_acc.is_empty() {
        let batch = batch_acc.into_batch();
        batches.push(batch);
    }

    batches
}

// ------------------------------------------------------------------------------------------------
/// Represents the operation data for a [`BasicBlockNodeBuilder`].
///
/// The decorators are bundled with the operation data to maintain the invariant that
/// decorator indices match the format of the operations:
/// - `Raw`: decorators have raw (unpadded) indices
/// - `Batched`: decorators have padded indices
#[derive(Debug)]
enum OperationData {
    /// Raw operations with raw decorator indices
    Raw {
        operations: Vec<Operation>,
        decorators: DecoratorList,
    },
    /// Pre-batched operations with padded decorator indices
    Batched {
        op_batches: Vec<OpBatch>,
        decorators: DecoratorList,
    },
}

/// Builder for creating [`BasicBlockNode`] instances with decorators.
#[derive(Debug)]
pub struct BasicBlockNodeBuilder {
    operation_data: OperationData,
    before_enter: Vec<DecoratorId>,
    after_exit: Vec<DecoratorId>,
    digest: Option<Word>,
}

impl BasicBlockNodeBuilder {
    /// Creates a new builder for a BasicBlockNode with the specified operations and decorators.
    ///
    /// The decorators must use raw (unpadded) operation indices.
    pub fn new(operations: Vec<Operation>, decorators: DecoratorList) -> Self {
        Self {
            operation_data: OperationData::Raw { operations, decorators },
            before_enter: Vec::new(),
            after_exit: Vec::new(),
            digest: None,
        }
    }

    /// Creates a builder from pre-existing OpBatches with padded decorator indices.
    ///
    /// This constructor is used during deserialization where operations are already batched
    /// and decorators already use padded indices. The digest must also be provided.
    ///
    /// The decorators must use padded operation indices that match the batched operations.
    pub(crate) fn from_op_batches(
        op_batches: Vec<OpBatch>,
        decorators: DecoratorList,
        digest: Word,
    ) -> Self {
        Self {
            operation_data: OperationData::Batched { op_batches, decorators },
            before_enter: Vec::new(),
            after_exit: Vec::new(),
            digest: Some(digest),
        }
    }

    /// Builds the BasicBlockNode with the specified decorators.
    pub fn build(self) -> Result<BasicBlockNode, MastForestError> {
        let (op_batches, digest, padded_decorators) = match self.operation_data {
            OperationData::Raw { operations, decorators } => {
                if operations.is_empty() {
                    return Err(MastForestError::EmptyBasicBlock);
                }

                // Validate decorators list (only in debug mode).
                #[cfg(debug_assertions)]
                validate_decorators(operations.len(), &decorators);

                let (op_batches, computed_digest) = batch_and_hash_ops(operations);
                // Batch operations (adds padding NOOPs)
                // Adjust decorators from raw to padded indices
                let padded_decorators = BasicBlockNode::adjust_decorators(decorators, &op_batches);

                // Use the forced digest if provided, otherwise use the computed digest
                let digest = self.digest.unwrap_or(computed_digest);

                (op_batches, digest, padded_decorators)
            },
            OperationData::Batched { op_batches, decorators } => {
                if op_batches.is_empty() {
                    return Err(MastForestError::EmptyBasicBlock);
                }

                // Decorators are already padded - no adjustment needed!
                let digest = self.digest.expect("digest must be set for batched operations");

                (op_batches, digest, decorators)
            },
        };

        Ok(BasicBlockNode {
            op_batches,
            digest,
            decorators: DecoratorStore::Owned {
                decorators: padded_decorators,
                before_enter: self.before_enter.clone(),
                after_exit: self.after_exit.clone(),
            },
        })
    }

    /// Add this node to a forest using relaxed validation.
    ///
    /// This method is used during deserialization where nodes may reference child nodes
    /// that haven't been added to the forest yet. The child node IDs have already been
    /// validated against the expected final node count during the `try_into_mast_node_builder`
    /// step, so we can safely skip validation here.
    ///
    /// Note: This is not part of the `MastForestContributor` trait because it's only
    /// intended for internal use during deserialization.
    ///
    /// For BasicBlockNode, this is equivalent to the normal `add_to_forest` since basic blocks
    /// don't have child nodes to validate.
    pub(in crate::mast) fn add_to_forest_relaxed(
        self,
        forest: &mut MastForest,
    ) -> Result<MastNodeId, MastForestError> {
        // For deserialization: decorators are already in forest.debug_info,
        // so we don't register them again. We just create the node.

        let future_node_id = MastNodeId::new_unchecked(forest.nodes.len() as u32);

        // Process based on operation data type
        let (op_batches, digest) = match self.operation_data {
            OperationData::Raw { operations, decorators: _ } => {
                if operations.is_empty() {
                    return Err(MastForestError::EmptyBasicBlock);
                }

                // Batch operations (adds padding NOOPs)
                let (op_batches, computed_digest) = batch_and_hash_ops(operations);

                // Use the forced digest if provided, otherwise use the computed digest
                let digest = self.digest.unwrap_or(computed_digest);

                (op_batches, digest)
            },
            OperationData::Batched { op_batches, decorators: _ } => {
                if op_batches.is_empty() {
                    return Err(MastForestError::EmptyBasicBlock);
                }

                // For batched operations, digest must be set
                let digest = self.digest.expect("digest must be set for batched operations");

                (op_batches, digest)
            },
        };

        // Create the node in the forest with Linked variant
        // Note: Decorators are already in forest.debug_info from deserialization
        let node_id = forest
            .nodes
            .push(MastNode::Block(BasicBlockNode {
                op_batches,
                digest,
                decorators: DecoratorStore::Linked { id: future_node_id },
            }))
            .map_err(|_| MastForestError::TooManyNodes)?;

        Ok(node_id)
    }
}

impl MastForestContributor for BasicBlockNodeBuilder {
    fn add_to_forest(self, forest: &mut MastForest) -> Result<MastNodeId, MastForestError> {
        // Determine the node ID that will be assigned
        let future_node_id = MastNodeId::new_unchecked(forest.nodes.len() as u32);

        // Process based on operation data type
        let (op_batches, digest, padded_decorators) = match self.operation_data {
            OperationData::Raw { operations, decorators } => {
                if operations.is_empty() {
                    return Err(MastForestError::EmptyBasicBlock);
                }

                // Validate decorators list (only in debug mode).
                #[cfg(debug_assertions)]
                validate_decorators(operations.len(), &decorators);

                // Batch operations (adds padding NOOPs)
                let (op_batches, computed_digest) = batch_and_hash_ops(operations);

                // Use the forced digest if provided, otherwise use the computed digest
                let digest = self.digest.unwrap_or(computed_digest);

                // Adjust decorator indices from raw to padded
                let padded_decorators = BasicBlockNode::adjust_decorators(decorators, &op_batches);

                (op_batches, digest, padded_decorators)
            },
            OperationData::Batched { op_batches, decorators } => {
                if op_batches.is_empty() {
                    return Err(MastForestError::EmptyBasicBlock);
                }

                // Decorators are already padded - no adjustment needed!
                let digest = self.digest.expect("digest must be set for batched operations");

                (op_batches, digest, decorators)
            },
        };

        // Add decorator info to the forest storage
        forest
            .debug_info
            .register_op_indexed_decorators(future_node_id, padded_decorators)
            .map_err(MastForestError::DecoratorError)?;

        // Add node-level decorators to the centralized NodeToDecoratorIds for efficient access
        forest.register_node_decorators(future_node_id, &self.before_enter, &self.after_exit);

        // Create the node in the forest with Linked variant from the start
        let node_id = forest
            .nodes
            .push(MastNode::Block(BasicBlockNode {
                op_batches,
                digest,
                decorators: DecoratorStore::Linked { id: future_node_id },
            }))
            .map_err(|_| MastForestError::TooManyNodes)?;

        Ok(node_id)
    }

    fn fingerprint_for_node(
        &self,
        forest: &MastForest,
        _hash_by_node_id: &impl LookupByIdx<MastNodeId, MastNodeFingerprint>,
    ) -> Result<MastNodeFingerprint, MastForestError> {
        // For BasicBlockNode, we need to implement custom logic because BasicBlock has special
        // decorator handling with operation indices that other nodes don't have

        // Process based on operation data type
        let (op_batches, digest, raw_decorators) = match &self.operation_data {
            OperationData::Raw { operations, decorators } => {
                // Compute digest - use forced digest if available, otherwise compute normally
                let (op_batches, computed_digest) = batch_and_hash_ops(operations.clone());
                let digest = self.digest.unwrap_or(computed_digest);

                // Decorators are already in raw form - no conversion needed
                #[cfg(debug_assertions)]
                {
                    validate_decorators(operations.len(), decorators);
                }

                (op_batches, digest, decorators.clone())
            },
            OperationData::Batched { op_batches, decorators } => {
                let digest = self.digest.expect("digest must be set for batched operations");

                // Convert from padded to raw indices for fingerprinting
                let pad2raw = PaddedToRawPrefix::new(op_batches);
                let raw_decorators: Vec<(usize, DecoratorId)> = decorators
                    .iter()
                    .map(|(padded_idx, decorator_id)| {
                        let raw_idx = padded_idx - pad2raw[*padded_idx];
                        (raw_idx, *decorator_id)
                    })
                    .collect();

                (op_batches.clone(), digest, raw_decorators)
            },
        };

        // Collect before_enter decorator fingerprints
        let before_enter_bytes: Vec<[u8; 32]> = self
            .before_enter
            .iter()
            .map(|&id| *forest[id].fingerprint().as_bytes())
            .collect();

        // Collect op-indexed decorator data (using raw indices)
        let adjusted_decorators = raw_decorators;

        // Collect op-indexed decorator data
        let mut op_decorator_data = Vec::with_capacity(adjusted_decorators.len() * 33);
        for (raw_op_idx, decorator_id) in &adjusted_decorators {
            op_decorator_data.extend_from_slice(&raw_op_idx.to_le_bytes());
            op_decorator_data.extend_from_slice(forest[*decorator_id].fingerprint().as_bytes());
        }

        // Collect after_exit decorator fingerprints
        let after_exit_bytes: Vec<[u8; 32]> =
            self.after_exit.iter().map(|&id| *forest[id].fingerprint().as_bytes()).collect();

        // Collect assert operation data
        let mut assert_data = Vec::new();
        for (op_idx, op) in op_batches.iter().flat_map(|batch| batch.ops()).enumerate() {
            if let Operation::U32assert2(inner_value)
            | Operation::Assert(inner_value)
            | Operation::MpVerify(inner_value) = op
            {
                let op_idx: u32 = op_idx
                    .try_into()
                    .expect("there are more than 2^{32}-1 operations in basic block");

                // we include the opcode to differentiate between `Assert` and `U32assert2`
                assert_data.push(op.op_code());
                // we include the operation index to distinguish between basic blocks that
                // would have the same assert instructions, but in a different order
                assert_data.extend_from_slice(&op_idx.to_le_bytes());
                let inner_value = inner_value.as_canonical_u64();
                assert_data.extend_from_slice(&inner_value.to_le_bytes());
            }
        }

        // Create iterator of slices from all collected data
        let decorator_bytes_iter = before_enter_bytes
            .iter()
            .map(|bytes| bytes.as_slice())
            .chain(core::iter::once(op_decorator_data.as_slice()))
            .chain(after_exit_bytes.iter().map(|bytes| bytes.as_slice()))
            .chain(core::iter::once(assert_data.as_slice()));

        if self.before_enter.is_empty()
            && self.after_exit.is_empty()
            && adjusted_decorators.is_empty()
            && assert_data.is_empty()
        {
            Ok(MastNodeFingerprint::new(digest))
        } else {
            let decorator_root = Blake3_256::hash_iter(decorator_bytes_iter);
            Ok(MastNodeFingerprint::with_decorator_root(digest, decorator_root))
        }
    }

    fn remap_children(self, _remapping: &impl LookupByIdx<MastNodeId, MastNodeId>) -> Self {
        // BasicBlockNode has no children to remap
        self
    }

    fn with_before_enter(mut self, decorators: impl Into<Vec<DecoratorId>>) -> Self {
        self.before_enter = decorators.into();
        self
    }

    fn with_after_exit(mut self, decorators: impl Into<Vec<DecoratorId>>) -> Self {
        self.after_exit = decorators.into();
        self
    }

    fn append_before_enter(&mut self, decorators: impl IntoIterator<Item = DecoratorId>) {
        self.before_enter.extend(decorators);
    }

    fn append_after_exit(&mut self, decorators: impl IntoIterator<Item = DecoratorId>) {
        self.after_exit.extend(decorators);
    }

    fn with_digest(mut self, digest: Word) -> Self {
        self.digest = Some(digest);
        self
    }
}

#[cfg(any(test, feature = "arbitrary"))]
impl proptest::prelude::Arbitrary for BasicBlockNodeBuilder {
    type Parameters = super::arbitrary::BasicBlockNodeParams;
    type Strategy = proptest::strategy::BoxedStrategy<Self>;

    fn arbitrary_with(params: Self::Parameters) -> Self::Strategy {
        use proptest::prelude::*;

        use super::arbitrary::{decorator_id_strategy, op_non_control_sequence_strategy};

        (op_non_control_sequence_strategy(params.max_ops_len),)
            .prop_flat_map(move |(ops,)| {
                let ops_len = ops.len().max(1); // ensure at least 1 op
                // For builders, decorator indices must be strictly less than ops_len
                // because they reference actual operation positions
                prop::collection::vec(
                    (0..ops_len, decorator_id_strategy(params.max_decorator_id_u32)),
                    0..=params.max_pairs,
                )
                .prop_map(move |mut decorators| {
                    decorators.sort_by_key(|(i, _)| *i);
                    (ops.clone(), decorators)
                })
            })
            .prop_map(|(ops, decorators)| Self::new(ops, decorators))
            .boxed()
    }
}