Fix CIOS overflow check (#602)

* fix cios overflow check * rename test
lambdaclass · Oct 11, 2023 · 5a938e7 · 5a938e7
1 parent c0314d8
commit 5a938e7
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 2 deletions.
diff --git a/math/src/field/fields/montgomery_backed_prime_fields.rs b/math/src/field/fields/montgomery_backed_prime_fields.rs
@@ -12,6 +12,7 @@ use core::marker::PhantomData;
 
 pub type U384PrimeField<M> = MontgomeryBackendPrimeField<M, 6>;
 pub type U256PrimeField<M> = MontgomeryBackendPrimeField<M, 4>;
+pub type U64PrimeField<M> = MontgomeryBackendPrimeField<M, 1>;
 
 /// This trait is necessary for us to be able to use unsigned integer types bigger than
 /// `u128` (the biggest native `unit`) as constant generics.
@@ -759,8 +760,10 @@ mod tests_u256_prime_fields {
     use crate::field::traits::IsPrimeField;
     #[cfg(feature = "std")]
     use crate::traits::ByteConversion;
-    use crate::unsigned_integer::element::UnsignedInteger;
     use crate::unsigned_integer::element::U256;
+    use crate::unsigned_integer::element::{UnsignedInteger, U64};
+
+    use super::U64PrimeField;
 
     #[derive(Clone, Debug)]
     struct U256Modulus29;
@@ -1122,4 +1125,24 @@ mod tests_u256_prime_fields {
         let b = U256F29Element::zero();
         assert_eq!(a, b);
     }
+
+    // Goldilocks
+    #[derive(Clone, Debug)]
+    struct GoldilocksModulus;
+    impl IsModulus<U64> for GoldilocksModulus {
+        const MODULUS: U64 = UnsignedInteger {
+            limbs: [18446744069414584321],
+        };
+    }
+
+    type GoldilocksField = U64PrimeField<GoldilocksModulus>;
+    type GoldilocksElement = FieldElement<GoldilocksField>;
+
+    #[test]
+    fn test_cios_overflow_case() {
+        let a = GoldilocksElement::from(732582227915286439);
+        let b = GoldilocksElement::from(3906369333256140342);
+        let expected_sum = GoldilocksElement::from(4638951561171426781);
+        assert_eq!(a + b, expected_sum);
+    }
 }
diff --git a/math/src/unsigned_integer/montgomery.rs b/math/src/unsigned_integer/montgomery.rs
@@ -67,7 +67,7 @@ impl MontgomeryAlgorithms {
         }
         let mut result = UnsignedInteger { limbs: t };
 
-        let overflow = t_extra[0] > 0;
+        let overflow = t_extra[1] > 0;
 
         if overflow || UnsignedInteger::const_le(q, &result) {
             (result, _) = UnsignedInteger::sub(&result, q);