Fix MastForest serialization

assembly/src/assembler/mod.rs

@@ -214,8 +214,8 @@ impl Assembler {
     /// calls to library procedures will be compiled down to a [`vm_core::mast::ExternalNode`] (i.e.
     /// a reference to the procedure's MAST root). This means that when executing a program compiled
     /// against a library, the processor will not be able to differentiate procedures with the same
-    /// MAST root but different decorators. 
-    /// 
+    /// MAST root but different decorators.
+    ///
     /// Hence, it is not recommended to export two procedures that have the same MAST root (i.e. are
     /// identical except for their decorators). Note however that we don't expect this scenario to
     /// be frequent in practice. For example, this could occur when APIs are being renamed and/or

core/src/mast/mod.rs

@@ -14,7 +14,7 @@ pub use node::{
     BasicBlockNode, CallNode, DynNode, ExternalNode, JoinNode, LoopNode, MastNode, OpBatch,
     OperationOrDecorator, SplitNode, OP_BATCH_SIZE, OP_GROUP_SIZE,
 };
-use winter_utils::DeserializationError;
+use winter_utils::{ByteWriter, DeserializationError, Serializable};
 
 use crate::{Decorator, DecoratorList, Operation};
 
@@ -56,6 +56,8 @@ 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> {
@@ -185,6 +187,10 @@ impl MastForest {
         self[node_id].set_before_enter(decorator_ids)
     }
 
+    pub fn set_after_exit(&mut self, node_id: MastNodeId, decorator_ids: Vec<DecoratorId>) {
+        self[node_id].set_after_exit(decorator_ids)
+    }
+
     /// Adds a basic block node to the forest, and returns the [`MastNodeId`] associated with it.
     ///
     /// It is assumed that the decorators have not already been added to the MAST forest. If they
@@ -554,6 +560,12 @@ impl fmt::Display for DecoratorId {
     }
 }
 
+impl Serializable for DecoratorId {
+    fn write_into<W: ByteWriter>(&self, target: &mut W) {
+        self.0.write_into(target)
+    }
+}
+
 // MAST FOREST ERROR
 // ================================================================================================
 

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

@@ -184,6 +184,11 @@ impl BasicBlockNode {
         DecoratorIterator::new(&self.decorators)
     }
 
+    /// Returns an iterator over the operations in the order in which they appear in the program.
+    pub fn operation_iter(&self) -> impl Iterator<Item = &Operation> {
+        self.op_batches.iter().flat_map(|batch| batch.ops())
+    }
+
     /// Returns the total number of operations and decorators in this basic block.
     pub fn num_operations_and_decorators(&self) -> u32 {
         let num_ops: usize = self.num_operations() as usize;