feat: introduce SpanDecorator

assembly/src/assembler/basic_block_builder.rs

@@ -1,12 +1,7 @@
 use alloc::{borrow::Borrow, string::ToString, vec::Vec};
-
-use vm_core::{
-    mast::{DecoratorId, MastNodeId},
-    sys_events::SystemEvent,
-    AssemblyOp, Decorator, Operation,
-};
-
-use super::{mast_forest_builder::MastForestBuilder, BodyWrapper, DecoratorList, ProcedureContext};
+use vm_core::{mast::{DecoratorId, MastNodeId}, sys_events::SystemEvent, AssemblyOp, Decorator, Operation};
+use vm_core::mast::DecoratorSpan;
+use super::{mast_forest_builder::MastForestBuilder, BodyWrapper, ProcedureContext};
 use crate::{ast::Instruction, AssemblyError, Span};
 
 // BASIC BLOCK BUILDER
@@ -24,7 +19,7 @@ use crate::{ast::Instruction, AssemblyError, Span};
 #[derive(Debug)]
 pub struct BasicBlockBuilder<'a> {
     ops: Vec<Operation>,
-    decorators: DecoratorList,
+    decorators: Vec<(usize, DecoratorId)>,
     epilogue: Vec<Operation>,
     last_asmop_pos: usize,
     mast_forest_builder: &'a mut MastForestBuilder,
@@ -174,11 +169,7 @@ impl BasicBlockBuilder<'_> {
     pub fn make_basic_block(&mut self) -> Result<Option<MastNodeId>, AssemblyError> {
         if !self.ops.is_empty() {
             let ops = self.ops.drain(..).collect();
-            let decorators = if !self.decorators.is_empty() {
-                Some(self.decorators.drain(..).collect())
-            } else {
-                None
-            };
+            let decorators = DecoratorSpan::from_raw_ops_decorators(self.decorators.drain(..).collect());
 
             let basic_block_node_id = self.mast_forest_builder.ensure_block(ops, decorators)?;
 

assembly/src/assembler/mast_forest_builder.rs

@@ -11,7 +11,7 @@ use vm_core::{
     },
     Decorator, DecoratorList, Operation,
 };
-
+use vm_core::mast::DecoratorSpan;
 use super::{GlobalProcedureIndex, Procedure};
 use crate::AssemblyError;
 
@@ -456,14 +456,18 @@ impl MastForestBuilder {
     }
 
     pub fn set_before_enter(&mut self, node_id: MastNodeId, decorator_ids: Vec<DecoratorId>) {
-        self.mast_forest[node_id].set_before_enter(decorator_ids);
+        // Does this need to have a Vec<DecoratorSpan>?
+        let span = DecoratorSpan::new_collection(decorator_ids).into_iter().next().unwrap();
+        self.mast_forest[node_id].set_before_enter(span);
 
         let new_node_fingerprint = self.fingerprint_for_node(&self[node_id]);
         self.hash_by_node_id.insert(node_id, new_node_fingerprint);
     }
 
     pub fn set_after_exit(&mut self, node_id: MastNodeId, decorator_ids: Vec<DecoratorId>) {
-        self.mast_forest[node_id].set_after_exit(decorator_ids);
+        // Does this need to have a Vec<DecoratorSpan>?
+        let span = DecoratorSpan::new_collection(decorator_ids).into_iter().next().unwrap();
+        self.mast_forest[node_id].set_after_exit(span);
 
         let new_node_fingerprint = self.fingerprint_for_node(&self[node_id]);
         self.hash_by_node_id.insert(node_id, new_node_fingerprint);

assembly/src/assembler/mod.rs

@@ -7,9 +7,9 @@ use vm_core::{
     crypto::hash::RpoDigest,
     debuginfo::SourceSpan,
     mast::{DecoratorId, MastNodeId},
-    DecoratorList, Felt, Kernel, Operation, Program,
+    Felt, Kernel, Operation, Program,
 };
-
+use vm_core::mast::DecoratorSpan;
 use crate::{
     ast::{self, Export, InvocationTarget, InvokeKind, ModuleKind, QualifiedProcedureName},
     diagnostics::Report,
@@ -642,10 +642,13 @@ impl Assembler {
                     )?;
 
                     if let Some(decorator_ids) = block_builder.drain_decorators() {
+                        // This just picks off the first decorator span, but should there be more than one?
+                        let span = DecoratorSpan::new_collection(decorator_ids).into_iter().next().unwrap();
+
                         // Attach the decorators before the first instance of the repeated node
                         let mut first_repeat_node =
                             block_builder.mast_forest_builder_mut()[repeat_node_id].clone();
-                        first_repeat_node.set_before_enter(decorator_ids);
+                        first_repeat_node.set_before_enter(span);
                         let first_repeat_node_id = block_builder
                             .mast_forest_builder_mut()
                             .ensure_node(first_repeat_node)?;

core/src/mast/merger/mod.rs

@@ -2,10 +2,7 @@ use alloc::{collections::BTreeMap, vec::Vec};
 
 use miden_crypto::{hash::blake::Blake3Digest, utils::collections::KvMap};
 
-use crate::mast::{
-    DecoratorId, MastForest, MastForestError, MastNode, MastNodeFingerprint, MastNodeId,
-    MultiMastForestIteratorItem, MultiMastForestNodeIter,
-};
+use crate::mast::{DecoratorId, DecoratorSpan, MastForest, MastForestError, MastNode, MastNodeFingerprint, MastNodeId, MultiMastForestIteratorItem, MultiMastForestNodeIter};
 
 #[cfg(test)]
 mod tests;
@@ -151,6 +148,8 @@ impl MastForestMerger {
             let new_decorator_id = if let Some(existing_decorator) =
                 self.decorators_by_hash.get(&merging_decorator_hash)
             {
+                // This implies that in some cases it is possible to have intersecting decorator sets
+                // meaning a DecoratorSpan will potentially need to be split up
                 *existing_decorator
             } else {
                 let new_decorator_id = self.mast_forest.add_decorator(merging_decorator.clone())?;
@@ -263,8 +262,18 @@ impl MastForestMerger {
                 )
             })
         };
-        let map_decorators = |decorators: &[DecoratorId]| -> Result<Vec<_>, MastForestError> {
-            decorators.iter().map(map_decorator_id).collect()
+
+        let remap_decorator_span = |decorator_span: &DecoratorSpan| -> Result<DecoratorSpan, MastForestError>{
+            let new_decorator_ids = decorator_span
+                .iter()
+                .map(|decorator_id| map_decorator_id(&decorator_id))
+                .collect::<Result<Vec<_>, _>>()?;
+            let new_spans = DecoratorSpan::new_collection(new_decorator_ids);
+
+            new_spans
+                .into_iter()
+                .next()
+                .ok_or_else(|| MastForestError::DecoratorSpanMissing(decorator_span.offset, decorator_span.num_decorators))
         };
 
         let map_node_id = |node_id: MastNodeId| {
@@ -311,9 +320,11 @@ impl MastForestMerger {
                         basic_block_node
                             .decorators()
                             .iter()
-                            .map(|(idx, decorator_id)| match map_decorator_id(decorator_id) {
-                                Ok(mapped_decorator) => Ok((*idx, mapped_decorator)),
-                                Err(err) => Err(err),
+                            .map(|(idx, decorator_span)| {
+                                match remap_decorator_span(decorator_span) {
+                                    Ok(mapped_decorator) => Ok((*idx, mapped_decorator)),
+                                    Err(err) => Err(err),
+                                }
                             })
                             .collect::<Result<Vec<_>, _>>()?,
                     ),
@@ -327,8 +338,8 @@ impl MastForestMerger {
         // Decorators must be handled specially for basic block nodes.
         // For other node types we can handle it centrally.
         if !mapped_node.is_basic_block() {
-            mapped_node.set_before_enter(map_decorators(node.before_enter())?);
-            mapped_node.set_after_exit(map_decorators(node.after_exit())?);
+            mapped_node.set_before_enter(remap_decorator_span(node.before_enter())?);
+            mapped_node.set_after_exit(remap_decorator_span(node.after_exit())?);
         }
 
         Ok(mapped_node)

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::{ByteWriter, DeserializationError, Serializable};
+use winter_utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
 
 use crate::{AdviceMap, Decorator, DecoratorList, Operation};
 
@@ -199,11 +199,11 @@ impl MastForest {
         Some(id_remappings)
     }
 
-    pub fn set_before_enter(&mut self, node_id: MastNodeId, decorator_ids: Vec<DecoratorId>) {
+    pub fn set_before_enter(&mut self, node_id: MastNodeId, decorator_ids: DecoratorSpan) {
         self[node_id].set_before_enter(decorator_ids)
     }
 
-    pub fn set_after_exit(&mut self, node_id: MastNodeId, decorator_ids: Vec<DecoratorId>) {
+    pub fn set_after_exit(&mut self, node_id: MastNodeId, decorator_ids: DecoratorSpan) {
         self[node_id].set_after_exit(decorator_ids)
     }
 
@@ -621,6 +621,132 @@ impl fmt::Display for MastNodeId {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Default, Copy)]
+pub struct DecoratorSpan {
+    /// Offset into the decorators vector of the mast forest.
+    offset: u32,
+
+    /// The number of decorators in this span.
+    num_decorators: u32,
+}
+
+impl DecoratorSpan {
+    pub fn new_collection(decorator_ids: Vec<DecoratorId>) -> Vec<DecoratorSpan> {
+        // Test that this returns an empty decoratorspan if the DecoratorId list is empty
+        let mut sorted_ids = decorator_ids;
+        sorted_ids.sort_unstable();
+
+        sorted_ids
+            .into_iter()
+            .fold(Vec::new(), |mut spans: Vec<DecoratorSpan>, id| {
+                if let Some(last) = spans.last_mut() {
+                    if last.offset + last.num_decorators == id.as_u32() {
+                        last.num_decorators += 1;
+                        return spans;
+                    }
+                }
+                spans.push(DecoratorSpan {
+                    offset: id.as_u32(),
+                    num_decorators: 1,
+                });
+                spans
+            })
+    }
+
+    pub fn from_raw_ops_decorators(decorators: Vec<(usize, DecoratorId)>) -> Option<DecoratorList> {
+        if !decorators.is_empty() {
+            let mut decorators_list: DecoratorList = Vec::new();
+            let mut current_index: Option<usize> = None;
+            let mut current_group: Vec<DecoratorId> = Vec::new();
+
+            for (index, decorator_id) in decorators{
+                if Some(index) != current_index {
+                    //process the previous group
+                    if !current_group.is_empty() {
+                        let spans = DecoratorSpan::new_collection(current_group);
+                        if let Some(span) = spans.into_iter().next() {
+                            decorators_list.push((index, span));
+                        }
+                    }
+                    // Start a new group
+                    current_index = Some(index);
+                    current_group = Vec::new();
+                }
+                current_group.push(decorator_id);
+            }
+
+            // Process the last group
+            if let Some(idx) = current_index {
+                if !current_group.is_empty() {
+                    let spans = DecoratorSpan::new_collection(current_group);
+                    if let Some(span) = spans.into_iter().next() {
+                        decorators_list.push((idx, span));
+                    }
+                }
+            }
+            Some(decorators_list)
+        } else {
+            None
+        }
+    }
+
+    pub fn iter(&self) -> impl Iterator<Item = DecoratorId> { //Should this return a reference?
+        (self.offset..self.offset + self.num_decorators).map(DecoratorId)
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.num_decorators == 0
+    }
+
+    pub fn read_safe<R: ByteReader>(
+        source: &mut R,
+        mast_forest: &MastForest,
+    ) -> Result<Self, DeserializationError> {
+        let value = Self::read_from(source)?;
+        let last_decorator_id = (value.offset + value.num_decorators) as usize;
+        if last_decorator_id < mast_forest.decorators.len() {
+            Ok(Self { offset: value.offset, num_decorators: value.num_decorators })
+        } else {
+            Err(DeserializationError::InvalidValue(format!(
+                "Invalid deserialized MAST decorator span with offset'{}' and num_decorators'{}', but only {} decorators in the forest",
+                value.offset, value.num_decorators,
+                mast_forest.nodes.len(),
+            )))
+        }
+    }
+}
+
+impl Deserializable for DecoratorSpan {
+    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
+        let offset = source.read_u32()?;
+        let num_decorators = source.read_u32()?;
+        Ok(DecoratorSpan { offset, num_decorators })
+    }
+}
+impl From<Vec<DecoratorId>> for DecoratorSpan {
+    fn from(decorator_ids: Vec<DecoratorId>) -> Self {
+        // Create a DecoratorSpan based on the size of the list and the first offset
+        let num_decorators = decorator_ids.len() as u32;
+        let offset = if num_decorators > 0 {
+            decorator_ids[0].into()
+        } else {
+            0
+        };
+
+        DecoratorSpan {
+            offset,
+            num_decorators,
+        }
+    }
+}
+
+impl Serializable for DecoratorSpan {
+    fn write_into<W: ByteWriter>(&self, target: &mut W) {
+        self.offset.write_into(target);
+        self.num_decorators.write_into(target);
+    }
+}
+
 // DECORATOR ID
 // ================================================================================================
 
@@ -707,6 +833,8 @@ pub enum MastForestError {
     NodeIdOverflow(MastNodeId, usize),
     #[error("decorator id {0} is greater than or equal to decorator count {1}")]
     DecoratorIdOverflow(DecoratorId, usize),
+    #[error("decorator span with offset '{0}' and num_decorators '{1}' should remap when merging, but was not")]
+    DecoratorSpanMissing(u32, u32),
     #[error("basic block cannot be created from an empty list of operations")]
     EmptyBasicBlock,
     #[error("decorator root of child with node id {0} is missing but is required for fingerprint computation")]