Sort_ecdsa_and_mod_builtins_private_inputs_by_idx (#1851)

* Sort_ecdsa_and_mod_builtins_private_inputs_by_idx * Add signature `get_air_private_input` test and msg if somehow reaching unreachable clause * sort sigs before iterating --------- Co-authored-by: Pedro Fontana <[email protected]>
lambdaclass · Nov 12, 2024 · 4bd5794 · 4bd5794
1 parent 6025257
commit 4bd5794
Show file tree

Hide file tree

Showing 3 changed files with 79 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -18,6 +18,9 @@
 
 * chore: update Rust required version to 1.81.0 [#1857](https://github.com/lambdaclass/cairo-vm/pull/1857)
 
+* fix: [#1851](https://github.com/lambdaclass/cairo-vm/pull/1851):
+  * Fix unsorted signature and mod builtin outputs in air_private_input.
+
 * feat(BREAKING): [#1824](https://github.com/lambdaclass/cairo-vm/pull/1824)[#1838](https://github.com/lambdaclass/cairo-vm/pull/1838):
     * Add support for dynamic layout
     * CLI change(BREAKING): The flag `cairo_layout_params_file` must be specified when using dynamic layout.

diff --git a/vm/src/vm/runners/builtin_runner/modulo.rs b/vm/src/vm/runners/builtin_runner/modulo.rs
@@ -263,6 +263,8 @@ impl ModBuiltinRunner {
             });
         }
 
+        instances.sort_by_key(|input| input.index);
+
         vec![PrivateInput::Mod(ModInput {
             instances,
             zero_value_address: relocation_table

diff --git a/vm/src/vm/runners/builtin_runner/signature.rs b/vm/src/vm/runners/builtin_runner/signature.rs
@@ -210,7 +210,17 @@ impl SignatureBuiltinRunner {
 
     pub fn air_private_input(&self, memory: &Memory) -> Vec<PrivateInput> {
         let mut private_inputs = vec![];
-        for (addr, signature) in self.signatures.borrow().iter() {
+
+        // Collect and sort the signatures by their index before the loop
+        let binding = self.signatures.borrow();
+        let mut sorted_signatures: Vec<_> = binding.iter().collect();
+        sorted_signatures.sort_by_key(|(addr, _)| {
+            addr.offset
+                .checked_div(CELLS_PER_SIGNATURE as usize)
+                .unwrap_or_default()
+        });
+
+        for (addr, signature) in sorted_signatures {
             if let (Ok(pubkey), Some(msg)) = (
                 memory.get_integer(*addr),
                 (*addr + 1_usize)
@@ -554,4 +564,67 @@ mod tests {
             assert_eq!(signature_a.s, signature_b.s);
         }
     }
+    #[test]
+    fn get_air_private_input() {
+        let mut builtin = SignatureBuiltinRunner::new(Some(512), true);
+
+        builtin.base = 0;
+
+        let signature1_r = Felt252::from(1234);
+        let signature1_s = Felt252::from(5678);
+        let signature2_r = Felt252::from(8765);
+        let signature2_s = Felt252::from(4321);
+
+        let sig1_addr = Relocatable::from((builtin.base as isize, 0));
+        let sig2_addr = Relocatable::from((builtin.base as isize, CELLS_PER_SIGNATURE as usize));
+
+        builtin
+            .add_signature(sig1_addr, &(signature1_r, signature1_s))
+            .unwrap();
+        builtin
+            .add_signature(sig2_addr, &(signature2_r, signature2_s))
+            .unwrap();
+
+        let pubkey1 = Felt252::from(1111);
+        let msg1 = Felt252::from(2222);
+        let pubkey2 = Felt252::from(3333);
+        let msg2 = Felt252::from(4444);
+
+        let segments = segments![
+            ((0, 0), 1111),
+            ((0, 1), 2222),
+            ((0, 2), 3333),
+            ((0, 3), 4444)
+        ];
+        let w1 =
+            Felt252::from(&div_mod(&BigInt::one(), &signature1_s.to_bigint(), &EC_ORDER).unwrap());
+
+        let w2 =
+            Felt252::from(&div_mod(&BigInt::one(), &signature2_s.to_bigint(), &EC_ORDER).unwrap());
+
+        let expected_private_inputs = vec![
+            PrivateInput::Signature(PrivateInputSignature {
+                index: 0,
+                pubkey: pubkey1,
+                msg: msg1,
+                signature_input: SignatureInput {
+                    r: signature1_r,
+                    w: w1,
+                },
+            }),
+            PrivateInput::Signature(PrivateInputSignature {
+                index: 1,
+                pubkey: pubkey2,
+                msg: msg2,
+                signature_input: SignatureInput {
+                    r: signature2_r,
+                    w: w2,
+                },
+            }),
+        ];
+
+        let private_inputs = builtin.air_private_input(&segments.memory);
+
+        assert_eq!(private_inputs, expected_private_inputs);
+    }
 }