From f552f63260f0087c5919279b471c69e4889fa3a5 Mon Sep 17 00:00:00 2001
From: Mauro Toscano <12560266+MauroToscano@users.noreply.github.com>
Date: Sun, 18 Aug 2024 17:07:15 -0300
Subject: [PATCH] Fix secp square  (#890)

* Add failing tests for secp

* Disable sos_square when no spare bit is found

* Clippy

* Clippy
---
 .../fields/montgomery_backed_prime_fields.rs  | 29 ++++++++++++++++++-
 1 file changed, 28 insertions(+), 1 deletion(-)

diff --git a/math/src/field/fields/montgomery_backed_prime_fields.rs b/math/src/field/fields/montgomery_backed_prime_fields.rs
index d87869060..bef1d312f 100644
--- a/math/src/field/fields/montgomery_backed_prime_fields.rs
+++ b/math/src/field/fields/montgomery_backed_prime_fields.rs
@@ -150,7 +150,11 @@ where
 
     #[inline(always)]
     fn square(a: &UnsignedInteger<NUM_LIMBS>) -> UnsignedInteger<NUM_LIMBS> {
-        MontgomeryAlgorithms::sos_square(a, &M::MODULUS, &Self::MU)
+        if Self::MODULUS_HAS_ONE_SPARE_BIT {
+            MontgomeryAlgorithms::sos_square(a, &M::MODULUS, &Self::MU)
+        } else {
+            MontgomeryAlgorithms::cios(a, a, &M::MODULUS, &Self::MU)
+        }
     }
 
     #[inline(always)]
@@ -1220,6 +1224,29 @@ mod tests_u256_prime_fields {
     type GoldilocksField = U64PrimeField<GoldilocksModulus>;
     type GoldilocksElement = FieldElement<GoldilocksField>;
 
+    #[derive(Clone, Debug)]
+    struct SecpModulus;
+    impl IsModulus<U256> for SecpModulus {
+        const MODULUS: U256 = UnsignedInteger::from_hex_unchecked(
+            "0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEFFFFFC2F",
+        );
+    }
+    type SecpMontField = U256PrimeField<SecpModulus>;
+    type SecpMontElement = FieldElement<SecpMontField>;
+
+    #[test]
+    fn secp256k1_minus_three_pow_2_is_9_with_all_operations() {
+        let minus_3 = -SecpMontElement::from_hex_unchecked("0x3");
+        let minus_3_mul_minus_3 = &minus_3 * &minus_3;
+        let minus_3_squared = minus_3.square();
+        let minus_3_pow_2 = minus_3.pow(2_u32);
+        let nine = SecpMontElement::from_hex_unchecked("0x9");
+
+        assert_eq!(minus_3_mul_minus_3, nine);
+        assert_eq!(minus_3_squared, nine);
+        assert_eq!(minus_3_pow_2, nine);
+    }
+
     #[test]
     fn test_cios_overflow_case() {
         let a = GoldilocksElement::from(732582227915286439);