Merge pull request #202 from 0xPolygonMiden/greenhat/i200-duplicate-s…

…tack-operand-schedule

[3/x] Always duplicate stack operands in CopyAll stack operand scheduling tactic

codegen/masm/src/codegen/emitter.rs

@@ -205,10 +205,13 @@ impl<'b, 'f: 'b> BlockEmitter<'b, 'f> {
         let args = self.function.f.dfg.inst_args(inst_info.inst);
         self.schedule_operands(args, inst_info.plain_arguments()).unwrap_or_else(|err| {
             panic!(
-                "failed to schedule operands: {:?} \n for inst:\n {}\n {:?}\n with error: {err:?}",
+                "failed to schedule operands: {:?} \n for inst: {} {:?}\n with error: {err:?}\n \
+                 constraints: {:?}\n stack: {:?}",
                 args,
                 inst_info.inst,
                 self.function.f.dfg.inst(inst_info.inst),
+                inst_info.plain_arguments(),
+                self.stack,
             )
         });
 

codegen/masm/src/codegen/opt/operands/context.rs

@@ -21,8 +21,6 @@ impl SolverContext {
         constraints: &[Constraint],
         stack: &crate::codegen::OperandStack,
     ) -> Result<Self, SolverError> {
-        use std::collections::btree_map::Entry;
-
         // Compute the expected output on the stack, as well as alias/copy information
         let mut stack = Stack::from(stack);
         let mut expected_output = Stack::default();
@@ -35,36 +33,23 @@ impl SolverContext {
                 Constraint::Move => {
                     expected_output.push(value);
                 }
-                // If we observe a value with copy semantics, then it is
-                // always referencing an alias, because the original would
-                // need to be preserved
+                // If we observe a value with copy semantics, then the expected
+                // output is always an alias, because the original would need to
+                // be preserved
                 Constraint::Copy => {
                     expected_output.push(copies.push(value));
                 }
             }
         }
 
         // Rename multiple occurrences of the same value on the operand stack, if present
-        let mut renamed = BTreeMap::<ValueOrAlias, u8>::default();
+        let mut dupes = CopyInfo::default();
         for operand in stack.iter_mut().rev() {
-            match renamed.entry(operand.value) {
-                Entry::Vacant(entry) => {
-                    entry.insert(0);
-                }
-                Entry::Occupied(mut entry) => {
-                    let next_id = entry.get_mut();
-                    *next_id += 1;
-                    operand.value.set_alias(NonZeroU8::new(*next_id).unwrap());
-                }
-            }
+            operand.value = dupes.push_if_duplicate(operand.value);
         }
 
         // Determine if the stack is already in the desired order
-        //
-        // If copies are required we can't consider the stack in order even if
-        // the operands we want are in the desired order, because we must make
-        // copies of them anyway.
-        let requires_copies = !copies.is_empty();
+        let requires_copies = copies.copies_required();
         let is_solved = !requires_copies
             && expected_output.iter().rev().all(|op| &stack[op.pos as usize] == op);
         if is_solved {
@@ -131,24 +116,59 @@ impl CopyInfo {
     }
 
     /// Push a new copy of `value`, returning an alias of that value
+    ///
+    /// NOTE: It is expected that `value` is not an alias.
     pub fn push(&mut self, value: ValueOrAlias) -> ValueOrAlias {
         use std::collections::btree_map::Entry;
 
+        assert!(!value.is_alias());
+
         self.num_copies += 1;
         match self.copies.entry(value) {
             Entry::Vacant(entry) => {
-                entry.insert(0);
+                entry.insert(1);
                 value.copy(unsafe { NonZeroU8::new_unchecked(1) })
             }
             Entry::Occupied(mut entry) => {
                 let next_id = entry.get_mut();
                 *next_id += 1;
-                value.copy(NonZeroU8::new(*next_id).unwrap())
+                value.copy(unsafe { NonZeroU8::new_unchecked(*next_id) })
             }
         }
     }
 
+    /// Push a copy of `value`, but only if `value` has already been seen
+    /// at least once, i.e. `value` is a duplicate.
+    ///
+    /// NOTE: It is expected that `value` is not an alias.
+    pub fn push_if_duplicate(&mut self, value: ValueOrAlias) -> ValueOrAlias {
+        use std::collections::btree_map::Entry;
+
+        assert!(!value.is_alias());
+
+        match self.copies.entry(value) {
+            // `value` is not a duplicate
+            Entry::Vacant(entry) => {
+                entry.insert(0);
+                value
+            }
+            // `value` is a duplicate, record it as such
+            Entry::Occupied(mut entry) => {
+                self.num_copies += 1;
+                let next_id = entry.get_mut();
+                *next_id += 1;
+                value.copy(unsafe { NonZeroU8::new_unchecked(*next_id) })
+            }
+        }
+    }
+
+    /// Returns true if `value` has at least one copy
     pub fn has_copies(&self, value: &ValueOrAlias) -> bool {
-        self.copies.contains_key(value)
+        self.copies.get(value).map(|count| *count > 0).unwrap_or(false)
+    }
+
+    /// Returns true if any of the values seen so far have copies
+    pub fn copies_required(&self) -> bool {
+        self.copies.values().any(|count| *count > 0)
     }
 }

codegen/masm/src/codegen/opt/operands/solver.rs

@@ -427,6 +427,40 @@ mod tests {
         }
     }
 
+    /// This test reproduces https://github.com/0xPolygonMiden/compiler/issues/200
+    /// where v7 value should be duplicated on the stack
+    #[test]
+    fn operand_movement_constraint_solver_duplicate() {
+        setup();
+        let v7 = hir::Value::from_u32(7);
+        let v16 = hir::Value::from_u32(16);
+        let v32 = hir::Value::from_u32(32);
+        let v0 = hir::Value::from_u32(0);
+
+        let tests = [[v32, v7, v16, v0]];
+
+        for test in tests.into_iter() {
+            let mut stack = crate::codegen::OperandStack::default();
+            for value in test.into_iter().rev() {
+                stack.push(crate::codegen::TypedValue {
+                    ty: Type::I32,
+                    value,
+                });
+            }
+            // The v7 is expected to be twice on stack
+            let expected = [v7, v7, v32, v16];
+            let constraints = [Constraint::Copy; 4];
+
+            match OperandMovementConstraintSolver::new(&expected, &constraints, &stack) {
+                Ok(solver) => {
+                    let _result = solver.solve().expect("no solution found");
+                }
+                Err(SolverError::AlreadySolved) => panic!("already solved"),
+                Err(err) => panic!("invalid solver context: {err:?}"),
+            }
+        }
+    }
+
     // Strategy:
     //
     // 1. Generate a set of 1..16 operands to form a stack (called `stack`), with no more than 2

tests/integration/Cargo.toml

@@ -40,3 +40,5 @@ proptest.workspace = true
 miden-core.workspace = true
 concat-idents = "1.1"
 blake3.workspace = true
+log.workspace = true
+env_logger.workspace = true

tests/integration/src/rust_masm_tests/abi_transform/tx_kernel.rs

@@ -2,9 +2,20 @@ use expect_test::expect_file;
 
 use crate::CompilerTest;
 
+#[allow(unused)]
+fn setup_log() {
+    use log::LevelFilter;
+    let _ = env_logger::builder()
+        .filter_level(LevelFilter::Trace)
+        .format_timestamp(None)
+        .is_test(true)
+        .try_init();
+}
+
 #[ignore = "until https://github.com/0xPolygonMiden/compiler/issues/200 is fixed"]
 #[test]
 fn test_get_inputs() {
+    // setup_log();
     let main_fn = "() -> Vec<Felt> {{ get_inputs() }}";
     let artifact_name = "abi_transform_tx_kernel_get_inputs";
     let mut test = CompilerTest::rust_fn_body_with_sdk(artifact_name, main_fn, true);