feat: serialize decorator data at the end (#1531)

CHANGELOG.md

@@ -4,6 +4,7 @@
 - [BREAKING] `Process` no longer takes ownership of the `Host` (#1571)
 - [BREAKING] `ProcessState` was converted from a trait to a struct (#1571)
 - [BREAKING] `Host` and `AdviceProvider` traits simplified (#1572)
+- [BREAKING] `MastForest` serialization/deserialization will store/read decorator data at the end of the binary (#1531)
 
 #### Enhancements
 - Added `miden_core::mast::MastForest::advice_map` to load it into the advice provider before the `MastForest` execution (#1574).

core/src/mast/mod.rs

@@ -69,8 +69,6 @@ impl MastForest {
 impl MastForest {
     /// The maximum number of nodes that can be stored in a single MAST forest.
     const MAX_NODES: usize = (1 << 30) - 1;
-    /// The maximum number of decorators that can be stored in a single MAST forest.
-    const MAX_DECORATORS: usize = Self::MAX_NODES;
 
     /// Adds a decorator to the forest, and returns the associated [`DecoratorId`].
     pub fn add_decorator(&mut self, decorator: Decorator) -> Result<DecoratorId, MastForestError> {
@@ -539,6 +537,21 @@ impl MastNodeId {
         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(crate) fn new_unchecked(value: u32) -> Self {
         Self(value)

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

@@ -224,6 +224,12 @@ impl BasicBlockNode {
             decorator_ids.into_iter().map(|decorator_id| (after_last_op_idx, decorator_id)),
         );
     }
+
+    /// Used to initialize decorators for the [`BasicBlockNode`]. Replaces the existing decorators
+    /// with the given ['DecoratorList'].
+    pub fn set_decorators(&mut self, decorator_list: DecoratorList) {
+        self.decorators = decorator_list;
+    }
 }
 
 // PRETTY PRINTING

core/src/mast/serialization/basic_blocks.rs

@@ -2,11 +2,8 @@ use alloc::vec::Vec;
 
 use winter_utils::{ByteReader, DeserializationError, Serializable, SliceReader};
 
-use super::{DecoratorDataOffset, NodeDataOffset};
-use crate::{
-    mast::{BasicBlockNode, DecoratorId, MastForest},
-    DecoratorList, Operation,
-};
+use super::NodeDataOffset;
+use crate::{mast::BasicBlockNode, Operation};
 
 // BASIC BLOCK DATA BUILDER
 // ================================================================================================
@@ -26,24 +23,15 @@ impl BasicBlockDataBuilder {
 
 /// Mutators
 impl BasicBlockDataBuilder {
-    /// Encodes a [`BasicBlockNode`] into the serialized [`crate::mast::MastForest`] data field.
-    pub fn encode_basic_block(
-        &mut self,
-        basic_block: &BasicBlockNode,
-    ) -> (NodeDataOffset, Option<DecoratorDataOffset>) {
+    /// Encodes a [`BasicBlockNode`]'s operations into the serialized [`crate::mast::MastForest`]
+    /// data field. The decorators are not encoded because are stored separately
+    pub fn encode_basic_block(&mut self, basic_block: &BasicBlockNode) -> NodeDataOffset {
         let ops_offset = self.node_data.len() as NodeDataOffset;
 
         let operations: Vec<Operation> = basic_block.operations().copied().collect();
         operations.write_into(&mut self.node_data);
 
-        if basic_block.decorators().is_empty() {
-            (ops_offset, None)
-        } else {
-            let decorator_data_offset = self.node_data.len() as DecoratorDataOffset;
-            basic_block.decorators().write_into(&mut self.node_data);
-
-            (ops_offset, Some(decorator_data_offset))
-        }
+        ops_offset
     }
 
     /// Returns the serialized [`crate::mast::MastForest`] node data field.
@@ -68,35 +56,14 @@ impl<'a> BasicBlockDataDecoder<'a> {
 
 /// Decoding methods
 impl BasicBlockDataDecoder<'_> {
-    pub fn decode_operations_and_decorators(
+    pub fn decode_operations(
         &self,
         ops_offset: NodeDataOffset,
-        decorator_list_offset: NodeDataOffset,
-        mast_forest: &MastForest,
-    ) -> Result<(Vec<Operation>, DecoratorList), DeserializationError> {
+    ) -> Result<Vec<Operation>, DeserializationError> {
         // Read ops
         let mut ops_data_reader = SliceReader::new(&self.node_data[ops_offset as usize..]);
         let operations: Vec<Operation> = ops_data_reader.read()?;
 
-        // read decorators only if there are some
-        let decorators = if decorator_list_offset == MastForest::MAX_DECORATORS as u32 {
-            Vec::new()
-        } else {
-            let mut decorators_data_reader =
-                SliceReader::new(&self.node_data[decorator_list_offset as usize..]);
-
-            let num_decorators: usize = decorators_data_reader.read()?;
-            (0..num_decorators)
-                .map(|_| {
-                    let decorator_loc: usize = decorators_data_reader.read()?;
-                    let decorator_id =
-                        DecoratorId::from_u32_safe(decorators_data_reader.read()?, mast_forest)?;
-
-                    Ok((decorator_loc, decorator_id))
-                })
-                .collect::<Result<DecoratorList, _>>()?
-        };
-
-        Ok((operations, decorators))
+        Ok(operations)
     }
 }

core/src/mast/serialization/decorator.rs

@@ -26,13 +26,9 @@ pub struct DecoratorInfo {
 impl DecoratorInfo {
     pub fn from_decorator(
         decorator: &Decorator,
-        data_builder: &mut DecoratorDataBuilder,
-        string_table_builder: &mut StringTableBuilder,
+        decorator_data_offset: DecoratorDataOffset,
     ) -> Self {
         let variant = EncodedDecoratorVariant::from(decorator);
-        let decorator_data_offset =
-            data_builder.encode_decorator_data(decorator, string_table_builder).unwrap_or(0);
-
         Self { variant, decorator_data_offset }
     }
 
@@ -210,6 +206,8 @@ impl Deserializable for EncodedDecoratorVariant {
 #[derive(Debug, Default)]
 pub struct DecoratorDataBuilder {
     decorator_data: Vec<u8>,
+    decorator_infos: Vec<DecoratorInfo>,
+    string_table_builder: StringTableBuilder,
 }
 
 /// Constructors
@@ -221,13 +219,15 @@ impl DecoratorDataBuilder {
 
 /// Mutators
 impl DecoratorDataBuilder {
+    pub fn add_decorator(&mut self, decorator: &Decorator) {
+        let decorator_data_offset = self.encode_decorator_data(decorator).unwrap_or(0);
+        self.decorator_infos
+            .push(DecoratorInfo::from_decorator(decorator, decorator_data_offset));
+    }
+
     /// If a decorator has extra data to store, encode it in internal data buffer, and return the
     /// offset of the newly added data. If not, return `None`.
-    pub fn encode_decorator_data(
-        &mut self,
-        decorator: &Decorator,
-        string_table_builder: &mut StringTableBuilder,
-    ) -> Option<DecoratorDataOffset> {
+    pub fn encode_decorator_data(&mut self, decorator: &Decorator) -> Option<DecoratorDataOffset> {
         let data_offset = self.decorator_data.len() as DecoratorDataOffset;
 
         match decorator {
@@ -239,21 +239,22 @@ impl DecoratorDataBuilder {
                 let loc = assembly_op.location();
                 self.decorator_data.write_bool(loc.is_some());
                 if let Some(loc) = loc {
-                    let str_offset = string_table_builder.add_string(loc.path.as_ref());
+                    let str_offset = self.string_table_builder.add_string(loc.path.as_ref());
                     self.decorator_data.write_usize(str_offset);
                     self.decorator_data.write_u32(loc.start.to_u32());
                     self.decorator_data.write_u32(loc.end.to_u32());
                 }
 
                 // context name
                 {
-                    let str_offset = string_table_builder.add_string(assembly_op.context_name());
+                    let str_offset =
+                        self.string_table_builder.add_string(assembly_op.context_name());
                     self.decorator_data.write_usize(str_offset);
                 }
 
                 // op
                 {
-                    let str_index_in_table = string_table_builder.add_string(assembly_op.op());
+                    let str_index_in_table = self.string_table_builder.add_string(assembly_op.op());
                     self.decorator_data.write_usize(str_index_in_table);
                 }
 
@@ -288,7 +289,11 @@ impl DecoratorDataBuilder {
     }
 
     /// Returns the serialized [`crate::mast::MastForest`] decorator data field.
-    pub fn finalize(self) -> Vec<u8> {
-        self.decorator_data
+    pub fn finalize(self) -> (Vec<u8>, Vec<DecoratorInfo>, StringTable) {
+        (
+            self.decorator_data,
+            self.decorator_infos,
+            self.string_table_builder.into_table(),
+        )
     }
 }

core/src/mast/serialization/info.rs

@@ -1,10 +1,10 @@
+use alloc::vec::Vec;
+
 use miden_crypto::hash::rpo::RpoDigest;
 use winter_utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
 
 use super::{basic_blocks::BasicBlockDataDecoder, NodeDataOffset};
-use crate::mast::{
-    BasicBlockNode, CallNode, JoinNode, LoopNode, MastForest, MastNode, MastNodeId, SplitNode,
-};
+use crate::mast::{BasicBlockNode, CallNode, JoinNode, LoopNode, MastNode, MastNodeId, SplitNode};
 
 // MAST NODE INFO
 // ================================================================================================
@@ -21,46 +21,32 @@ pub struct MastNodeInfo {
 }
 
 impl MastNodeInfo {
-    /// Constructs a new [`MastNodeInfo`] from a [`MastNode`], along with an `ops_offset` and
-    /// `decorator_list_offset` in the case of [`BasicBlockNode`].
+    /// Constructs a new [`MastNodeInfo`] from a [`MastNode`], along with an `ops_offset`
     ///
-    /// If the represented [`MastNode`] is a [`BasicBlockNode`] that has an empty decorator list,
-    /// use `MastForest::MAX_DECORATORS` for the value of `decorator_list_offset`. For non-basic
-    /// block nodes, `ops_offset` and `decorator_list_offset` are ignored, and should be set to 0.
-    pub fn new(
-        mast_node: &MastNode,
-        ops_offset: NodeDataOffset,
-        decorator_list_offset: NodeDataOffset,
-    ) -> Self {
+    /// For non-basic block nodes, `ops_offset` is ignored, and should be set to 0.
+    pub fn new(mast_node: &MastNode, ops_offset: NodeDataOffset) -> Self {
         if !matches!(mast_node, &MastNode::Block(_)) {
             debug_assert_eq!(ops_offset, 0);
-            debug_assert_eq!(decorator_list_offset, 0);
         }
 
-        let ty = MastNodeType::new(mast_node, ops_offset, decorator_list_offset);
+        let ty = MastNodeType::new(mast_node, ops_offset);
 
         Self { ty, digest: mast_node.digest() }
     }
 
-    /// Attempts to convert this [`MastNodeInfo`] into a [`MastNode`] for the given `mast_forest`.
+    /// Attempts to convert this [`MastNodeInfo`] into a [`MastNode`].
     ///
     /// The `node_count` is the total expected number of nodes in the [`MastForest`] **after
     /// deserialization**.
     pub fn try_into_mast_node(
         self,
-        mast_forest: &MastForest,
         node_count: usize,
         basic_block_data_decoder: &BasicBlockDataDecoder,
     ) -> Result<MastNode, DeserializationError> {
         match self.ty {
-            MastNodeType::Block { ops_offset, decorator_list_offset } => {
-                let (operations, decorators) = basic_block_data_decoder
-                    .decode_operations_and_decorators(
-                        ops_offset,
-                        decorator_list_offset,
-                        mast_forest,
-                    )?;
-                let block = BasicBlockNode::new_unsafe(operations, decorators, self.digest);
+            MastNodeType::Block { ops_offset } => {
+                let operations = basic_block_data_decoder.decode_operations(ops_offset)?;
+                let block = BasicBlockNode::new_unsafe(operations, Vec::new(), self.digest);
                 Ok(MastNode::Block(block))
             },
             MastNodeType::Join { left_child_id, right_child_id } => {
@@ -151,8 +137,6 @@ pub enum MastNodeType {
     Block {
         // offset of operations in node data
         ops_offset: u32,
-        // offset of DecoratorList in node data
-        decorator_list_offset: u32,
     } = BLOCK,
     Call {
         callee_id: u32,
@@ -168,18 +152,11 @@ pub enum MastNodeType {
 /// Constructors
 impl MastNodeType {
     /// Constructs a new [`MastNodeType`] from a [`MastNode`].
-    ///
-    /// If the represented [`MastNode`] is a [`BasicBlockNode`] that has an empty decorator list,
-    /// use `MastForest::MAX_DECORATORS` for the value of `decorator_list_offset`.
-    pub fn new(
-        mast_node: &MastNode,
-        ops_offset: NodeDataOffset,
-        decorator_list_offset: NodeDataOffset,
-    ) -> Self {
+    pub fn new(mast_node: &MastNode, ops_offset: NodeDataOffset) -> Self {
         use MastNode::*;
 
         match mast_node {
-            Block(_block_node) => Self::Block { decorator_list_offset, ops_offset },
+            Block(_block_node) => Self::Block { ops_offset },
             Join(join_node) => Self::Join {
                 left_child_id: join_node.first().0,
                 right_child_id: join_node.second().0,
@@ -223,9 +200,7 @@ impl Serializable for MastNodeType {
                 else_branch_id: else_branch,
             } => Self::encode_u32_pair(if_branch, else_branch),
             MastNodeType::Loop { body_id: body } => Self::encode_u32_payload(body),
-            MastNodeType::Block { ops_offset, decorator_list_offset } => {
-                Self::encode_u32_pair(ops_offset, decorator_list_offset)
-            },
+            MastNodeType::Block { ops_offset } => Self::encode_u32_payload(ops_offset),
             MastNodeType::Call { callee_id } => Self::encode_u32_payload(callee_id),
             MastNodeType::SysCall { callee_id } => Self::encode_u32_payload(callee_id),
             MastNodeType::Dyn => 0,
@@ -300,8 +275,8 @@ impl Deserializable for MastNodeType {
                 Ok(Self::Loop { body_id })
             },
             BLOCK => {
-                let (ops_offset, decorator_list_offset) = Self::decode_u32_pair(payload);
-                Ok(Self::Block { ops_offset, decorator_list_offset })
+                let ops_offset = Self::decode_u32_payload(payload)?;
+                Ok(Self::Block { ops_offset })
             },
             CALL => {
                 let callee_id = Self::decode_u32_payload(payload)?;