feat: serialize BasicBlocks in padded representation (1/3)

CHANGELOG.md

@@ -26,6 +26,7 @@
 - [BREAKING] Removed `miden debug`, `miden analyze` and `miden repl` ([#2483](https://github.com/0xMiden/miden-vm/pull/2483)).
 - [BREAKING] Change backend from winterfell to Plonky3 ([#2472](https://github.com/0xMiden/miden-vm/pull/2472)).
 - Added validation of `core_trace_fragment_size` in `ExecutionOptions` ([#2528](https://github.com/0xMiden/miden-vm/pull/2528)).
+- [BREAKING] Change serialization of `BasicBlockNode`s to use padded indices ([#2466](https://github.com/0xMiden/miden-vm/pull/2466/)).
 
 ## 0.20.2 (TBD)
 - Fix issue where decorator access was not bypassed properly in release mode ([#2529](https://github.com/0xMiden/miden-vm/pull/2529)).

core/src/mast/node/basic_block_node/mod.rs

@@ -648,20 +648,29 @@ impl MastNodeExt for BasicBlockNode {
     type Builder = BasicBlockNodeBuilder;
 
     fn to_builder(self, forest: &MastForest) -> Self::Builder {
-        let operations: Vec<Operation> = self.raw_operations().cloned().collect();
-        let un_adjusted_decorators = self.raw_op_indexed_decorators(forest);
-
-        let (before_enter, after_exit) = match self.decorators {
-            DecoratorStore::Owned { before_enter, after_exit, .. } => (before_enter, after_exit),
+        // 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 the decorators from the forest's NodeToDecoratorIds
+                // 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();
-                (before_enter, after_exit)
+                (padded_decorators, before_enter, after_exit)
             },
         };
 
-        BasicBlockNodeBuilder::new(operations, un_adjusted_decorators)
+        // 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)
     }
@@ -1253,57 +1262,115 @@ fn batch_ops(ops: Vec<Operation>) -> Vec<OpBatch> {
 }
 
 // ------------------------------------------------------------------------------------------------
+/// 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 {
-    operations: Vec<Operation>,
-    decorators: DecoratorList,
+    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 {
-            operations,
-            decorators,
+            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),
+        }
+    }
+
     /// Used to initialize decorators for the [`BasicBlockNodeBuilder`]. Replaces the existing
     /// decorators with the given ['DecoratorList'].
     pub(crate) fn set_decorators(&mut self, decorators: DecoratorList) {
-        self.decorators = decorators;
+        match &mut self.operation_data {
+            OperationData::Raw { decorators: dec, .. } => *dec = decorators,
+            OperationData::Batched { decorators: dec, .. } => *dec = decorators,
+        }
     }
 
     /// Builds the BasicBlockNode with the specified decorators.
     pub fn build(self) -> Result<BasicBlockNode, MastForestError> {
-        if self.operations.is_empty() {
-            return Err(MastForestError::EmptyBasicBlock);
-        }
+        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(self.operations.len(), &self.decorators);
+                // Validate decorators list (only in debug mode).
+                #[cfg(debug_assertions)]
+                validate_decorators(operations.len(), &decorators);
 
-        let (op_batches, computed_digest) = batch_and_hash_ops(self.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(self.decorators, &op_batches);
+                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");
 
-        // Use the forced digest if provided, otherwise use the computed digest
-        let digest = self.digest.unwrap_or(computed_digest);
+                (op_batches, digest, decorators)
+            },
+        };
 
         Ok(BasicBlockNode {
             op_batches,
             digest,
             decorators: DecoratorStore::Owned {
-                decorators: reflowed_decorators,
+                decorators: padded_decorators,
                 before_enter: self.before_enter.clone(),
                 after_exit: self.after_exit.clone(),
             },
@@ -1334,31 +1401,47 @@ impl BasicBlockNodeBuilder {
 
 impl MastForestContributor for BasicBlockNodeBuilder {
     fn add_to_forest(self, forest: &mut MastForest) -> Result<MastNodeId, MastForestError> {
-        if self.operations.is_empty() {
-            return Err(MastForestError::EmptyBasicBlock);
-        }
-
-        // Validate decorators list (only in debug mode).
-        #[cfg(debug_assertions)]
-        validate_decorators(self.operations.len(), &self.decorators);
-
         // Determine the node ID that will be assigned
         let future_node_id = MastNodeId::new_unchecked(forest.nodes.len() as u32);
 
-        // Process operations and decorators directly without creating an intermediate Owned node
-        let (op_batches, computed_digest) = batch_and_hash_ops(self.operations);
+        // 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);
+                }
 
-        // Use the forced digest if provided, otherwise use the computed digest
-        let digest = self.digest.unwrap_or(computed_digest);
+                // 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 using the same logic as
-        // BasicBlockNode::adjust_decorators
-        let reflowed_decorators = BasicBlockNode::adjust_decorators(self.decorators, &op_batches);
+                // 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, reflowed_decorators)
+            .register_op_indexed_decorators(future_node_id, padded_decorators)
             .map_err(MastForestError::DecoratorError)?;
 
         // Add node-level decorators to the centralized NodeToDecoratorIds for efficient access
@@ -1385,9 +1468,37 @@ impl MastForestContributor for BasicBlockNodeBuilder {
         // For BasicBlockNode, we need to implement custom logic because BasicBlock has special
         // decorator handling with operation indices that other nodes don't have
 
-        // Compute digest - use forced digest if available, otherwise compute normally
-        let (op_batches, computed_digest) = batch_and_hash_ops(self.operations.clone());
-        let digest = self.digest.unwrap_or(computed_digest);
+        // 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
@@ -1396,21 +1507,8 @@ impl MastForestContributor for BasicBlockNodeBuilder {
             .map(|&id| forest[id].fingerprint().as_bytes())
             .collect();
 
-        // Hash op-indexed decorators using the same logic as node.indexed_decorator_iter()
-        #[cfg(debug_assertions)]
-        {
-            validate_decorators(self.operations.len(), &self.decorators);
-        }
-        // For BasicBlockNodeBuilder, convert from padded to raw indices
-        let pad2raw = PaddedToRawPrefix::new(&op_batches);
-        let adjusted_decorators: Vec<(usize, DecoratorId)> = self
-            .decorators
-            .iter()
-            .map(|(padded_idx, decorator_id)| {
-                let raw_idx = padded_idx - pad2raw[*padded_idx];
-                (raw_idx, *decorator_id)
-            })
-            .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);

core/src/mast/node/basic_block_node/op_batch.rs

@@ -116,6 +116,27 @@ impl OpBatch {
         self.num_groups
     }
 
+    /// Creates a new OpBatch from its constituent parts.
+    ///
+    /// This constructor is used during deserialization to reconstruct OpBatches with the exact
+    /// structure they had when serialized.
+    ///
+    /// # Arguments
+    /// * `ops` - The operations in this batch (including padding NOOPs)
+    /// * `indptr` - Array of group boundary indices
+    /// * `padding` - Array of padding flags for each group
+    /// * `groups` - Array of group hashes and immediate values
+    /// * `num_groups` - Number of groups in this batch
+    pub(crate) fn new_from_parts(
+        ops: Vec<Operation>,
+        indptr: [usize; Self::BATCH_SIZE_PLUS_ONE],
+        padding: [bool; BATCH_SIZE],
+        groups: [Felt; BATCH_SIZE],
+        num_groups: usize,
+    ) -> Self {
+        Self { ops, indptr, padding, groups, num_groups }
+    }
+
     /// Returns the (op_group_idx, op_idx_in_group) given an operation index in the batch. Returns
     /// `None` if the index is out of bounds.
     ///