diff --git a/core/src/mast/mod.rs b/core/src/mast/mod.rs index c5186052cd..37f1c69efe 100644 --- a/core/src/mast/mod.rs +++ b/core/src/mast/mod.rs @@ -167,7 +167,7 @@ impl MastForest { /// /// This method modifies the forest in-place, removing all decorator information /// including operation-indexed decorators, before-enter decorators, after-exit - /// decorators, and error codes while keeping the CSR structure valid. + /// decorators, and error codes. /// /// # Examples /// @@ -179,7 +179,8 @@ impl MastForest { /// forest.strip_decorators(); // forest is now stripped /// ``` pub fn strip_decorators(&mut self) { - self.debug_info = DebugInfo::empty_for_nodes(self.nodes.len()); + // Clear all debug info (decorators and error codes) + self.debug_info.clear(); } /// Compacts the forest by merging duplicate nodes. @@ -573,18 +574,58 @@ impl MastForest { // If a target operation index is provided, return the assembly op associated with that // operation. Some(target_op_idx) => { - for (op_idx, decorator_id) in decorator_links { - if let Some(Decorator::AsmOp(assembly_op)) = self.decorator_by_id(decorator_id) - { - // when an instruction compiles down to multiple operations, only the first - // operation is associated with the assembly op. We need to check if the - // target operation index falls within the range of operations associated - // with the assembly op. - if target_op_idx >= op_idx - && target_op_idx < op_idx + assembly_op.num_cycles() as usize - { - return Some(assembly_op); + // With the 1-1 mapping, every operation covered by an AsmOp is explicitly linked + // to that AsmOp. We need to find the minimum op_idx for each unique AsmOp decorator + // to determine the start of its range, then check if target_op_idx falls within + // that range. + let mut asmop_ranges: std::collections::HashMap = + std::collections::HashMap::new(); + + // For Block nodes with target_op_idx, we should only consider operation-indexed + // decorators (not before_enter/after_exit) for operation-specific lookups + match node { + MastNode::Block(_) => { + // For Block nodes, use only operation-indexed decorators + let op_iter = self + .decorator_links_for_node(node_id) + .expect("Block node must have some valid set of decorator links"); + + for (op_idx, decorator_id) in op_iter { + if let Some(Decorator::AsmOp(assembly_op)) = self.decorator_by_id(decorator_id) + { + let num_cycles = assembly_op.num_cycles() as usize; + asmop_ranges + .entry(decorator_id) + .and_modify(|(min_idx, cycles)| { + *min_idx = (*min_idx).min(op_idx); + *cycles = num_cycles; // num_cycles should be the same for same decorator_id + }) + .or_insert((op_idx, num_cycles)); + } } + }, + _ => { + // For non-Block nodes, before_enter is at index 0, after_exit is at index 1 + // These are not operation ranges, so we check exact matches + for (op_idx, decorator_id) in decorator_links { + if op_idx == target_op_idx { + if let Some(Decorator::AsmOp(assembly_op)) = self.decorator_by_id(decorator_id) + { + return Some(assembly_op); + } + } + } + }, + } + + // Second pass: check if target_op_idx falls within any AsmOp's range + for (decorator_id, (min_op_idx, num_cycles)) in asmop_ranges { + if target_op_idx >= min_op_idx + && target_op_idx < min_op_idx + num_cycles + && let Some(Decorator::AsmOp(assembly_op)) = + self.decorator_by_id(decorator_id) + { + return Some(assembly_op); } } }, @@ -652,7 +693,7 @@ impl MastForest { // TEST HELPERS // ================================================================================================ -#[cfg(test)] +#[cfg(any(test, feature = "testing"))] impl MastForest { /// Returns all decorators for a given node as a vector of (position, DecoratorId) tuples. /// diff --git a/core/src/mast/tests.rs b/core/src/mast/tests.rs index 72aa0021d6..bc27ac0f3e 100644 --- a/core/src/mast/tests.rs +++ b/core/src/mast/tests.rs @@ -869,13 +869,14 @@ fn test_mast_forest_get_assembly_op_with_target_index() { let decorator_id = forest.add_decorator(Decorator::AsmOp(assembly_op.clone())).unwrap(); // Add a basic block node with the decorator at index 2 + // With 1-to-1 mapping, all operations covered by the AsmOp (indices 2, 3, 4) should be linked let operations = vec![ Operation::Push(Felt::new(1)), Operation::Push(Felt::new(2)), Operation::Mul, Operation::Add, ]; - let decorators = vec![(2, decorator_id)]; // Decorator at operation index 2 + let decorators = vec![(2, decorator_id), (3, decorator_id), (4, decorator_id)]; // Decorator linked to all operations it covers let node_id = BasicBlockNodeBuilder::new(operations, decorators) .add_to_forest(&mut forest) .unwrap(); diff --git a/core/src/operations/decorators/mod.rs b/core/src/operations/decorators/mod.rs index 6b56120cb9..e1a7f65321 100644 --- a/core/src/operations/decorators/mod.rs +++ b/core/src/operations/decorators/mod.rs @@ -87,5 +87,6 @@ impl fmt::Display for Decorator { /// index. /// /// Note: for `AssemblyOp` decorators, when an instruction compiles down to multiple operations, -/// only the first operation is associated with the assembly op. +/// all operations are associated with the same assembly op, creating a 1-1 mapping between +/// operations and AsmOp decorators. pub type DecoratorList = Vec; diff --git a/crates/assembly/Cargo.toml b/crates/assembly/Cargo.toml index 47f228a26b..c108f0e037 100644 --- a/crates/assembly/Cargo.toml +++ b/crates/assembly/Cargo.toml @@ -43,6 +43,7 @@ thiserror.workspace = true miden-assembly = { path = ".", default-features = false, features = ["testing"] } miden-mast-package = { workspace = true, features = ["arbitrary"] } # core lib and processor are imported using paths to avoid specifying dependency versions +miden-core = { path = "../../core", default-features = false, features = ["testing"] } miden-core-lib = { path = "../lib/core", default-features = false } miden-processor = { path = "../../processor", features = ["testing"] } insta = { workspace = true } diff --git a/crates/assembly/src/basic_block_builder.rs b/crates/assembly/src/basic_block_builder.rs index 4d8cf89fb5..2f27e6d9e2 100644 --- a/crates/assembly/src/basic_block_builder.rs +++ b/crates/assembly/src/basic_block_builder.rs @@ -138,21 +138,24 @@ impl BasicBlockBuilder<'_> { /// Computes the number of cycles elapsed since the last invocation of track_instruction() and /// updates the related AsmOp decorator to include this cycle count. /// + /// Creates a 1-1 mapping between operations and AsmOp: every operation covered by the AsmOp + /// is linked to that AsmOp. This means if an instruction compiles to multiple operations, all + /// of those operations will be associated with the same AsmOp decorator. + /// /// If the cycle count is 0, the original decorator is removed from the list and returned. This /// can happen for instructions which do not contribute any operations to the span block - e.g., /// exec, call, and syscall. pub fn set_instruction_cycle_count(&mut self) -> Option { // get the last asmop decorator and the cycle at which it was added - let (op_start, assembly_op_id) = - self.decorators.get_mut(self.last_asmop_pos).expect("no asmop decorator"); + let (op_start, assembly_op_id) = self.decorators[self.last_asmop_pos]; - let assembly_op = match &mut self.mast_forest_builder[*assembly_op_id] { + let assembly_op = match &mut self.mast_forest_builder[assembly_op_id] { Decorator::AsmOp(assembly_op) => assembly_op, _ => panic!("internal error: last asmop decorator is not an AsmOp"), }; // compute the cycle count for the instruction - let cycle_count = self.ops.len() - *op_start; + let cycle_count = self.ops.len() - op_start; // if the cycle count is 0, remove the decorator; otherwise update its cycle count if cycle_count == 0 { @@ -161,6 +164,14 @@ impl BasicBlockBuilder<'_> { } else { assembly_op.set_num_cycles(cycle_count as u8); + // Create a 1-1 mapping: link the AsmOp to all operations it covers. + // The first operation is already linked at op_start, so we add links for + // operations from op_start + 1 to op_start + cycle_count - 1. + // We insert them in reverse order to maintain the sorted order of the decorators list. + for op_idx in ((op_start + 1)..(op_start + cycle_count)).rev() { + self.decorators.insert(self.last_asmop_pos + 1, (op_idx, assembly_op_id)); + } + None } } diff --git a/crates/assembly/src/tests.rs b/crates/assembly/src/tests.rs index 5834287da7..01b56f42d2 100644 --- a/crates/assembly/src/tests.rs +++ b/crates/assembly/src/tests.rs @@ -8,12 +8,10 @@ use std::{ use miden_assembly_syntax::{ast::Path, diagnostics::WrapErr, library::LibraryExport}; use miden_core::{ EventId, Operation, Program, StackInputs, Word, assert_matches, - mast::{MastNodeExt, MastNodeId}, + mast::{MastForest, MastNodeExt, MastNodeId}, utils::{Deserializable, Serializable}, }; -use miden_mast_package::{ - MastArtifact, MastForest, Package, PackageExport, PackageKind, PackageManifest, -}; +use miden_mast_package::{MastArtifact, Package, PackageExport, PackageKind, PackageManifest}; #[cfg(test)] use miden_processor::{AdviceInputs, DefaultHost, ExecutionOptions}; use proptest::{ @@ -4573,3 +4571,106 @@ fn test_cross_module_constant_resolution() -> TestResult { Ok(()) } + +#[test] +fn test_multi_cycle_instruction_decorator_mapping() -> TestResult { + let context = TestContext::default(); + // Assemble MASM code with multi-cycle instructions. + // `sub` compiles to [Neg, Add] - 2 operations + // `div` compiles to [Inv, Mul] - 2 operations + let source = source_file!( + &context, + r#" + begin + push.10 + push.5 + sub # This compiles to [Neg, Add] - 2 operations + push.20 + push.4 + div # This compiles to [Inv, Mul] - 2 operations + end + "# + ); + + let program = context.assemble(source)?; + let mast_forest: &MastForest = program.mast_forest().as_ref(); + let entrypoint = program.entrypoint(); + + // Get all decorators for the entrypoint node + let all_decorators = mast_forest.all_decorators(entrypoint); + + // Find AsmOp decorators and group them by DecoratorId + // to verify that all operations in a multi-cycle range share the same DecoratorId + let mut decorator_id_to_positions: std::collections::HashMap<_, Vec> = + std::collections::HashMap::new(); + + for (position, decorator_id) in &all_decorators { + // Check if this is an AsmOp decorator + if let miden_core::Decorator::AsmOp(_) = &mast_forest[*decorator_id] { + decorator_id_to_positions + .entry(*decorator_id) + .or_insert_with(Vec::new) + .push(*position); + } + } + + // Verify that we have decorators for multi-cycle instructions + assert!( + !decorator_id_to_positions.is_empty(), + "Expected to find AsmOp decorators for instructions" + ); + + // For each decorator ID, verify that all operations it covers share the same ID + // and that they form a contiguous range + for (decorator_id, positions) in &decorator_id_to_positions { + // Sort positions to check contiguity + let mut sorted_positions = positions.clone(); + sorted_positions.sort(); + + // Verify that all positions share the same decorator ID + for position in positions { + let decorators_at_pos: Vec<_> = all_decorators + .iter() + .filter(|(pos, _)| pos == position) + .map(|(_, id)| *id) + .collect(); + + // Check that at least one decorator at this position matches our decorator ID + assert!( + decorators_at_pos.contains(decorator_id), + "Position {} should have decorator ID {:?}", + position, + decorator_id + ); + } + + // For multi-cycle instructions (more than 1 position), verify they form a contiguous range + if sorted_positions.len() > 1 { + // Check that positions form a contiguous range + for i in 1..sorted_positions.len() { + assert_eq!( + sorted_positions[i], + sorted_positions[i - 1] + 1, + "Multi-cycle instruction decorator positions should be contiguous. \ + Found positions: {:?} for decorator ID {:?}", + sorted_positions, + decorator_id + ); + } + + // Verify the decorator has the correct cycle count + if let miden_core::Decorator::AsmOp(asm_op) = &mast_forest[*decorator_id] { + assert_eq!( + asm_op.num_cycles() as usize, + sorted_positions.len(), + "Decorator cycle count should match number of operations. \ + Expected {} cycles for positions {:?}", + sorted_positions.len(), + sorted_positions + ); + } + } + } + + Ok(()) +}