Add qm31 constructor from u32

starkware-libs · Oct 4, 2024 · efcba52 · efcba52
1 parent 947e147
commit efcba52
Showing 1 changed file with 26 additions and 4 deletions.
diff --git a/stwo_cairo_verifier/src/fields/qm31.cairo b/stwo_cairo_verifier/src/fields/qm31.cairo
@@ -1,4 +1,4 @@
-use super::m31::{M31, m31};
+use super::m31::{M31, m31, M31Impl};
 use super::cm31::{CM31, cm31, CM31Trait};
 use core::num::traits::zero::Zero;
 use core::num::traits::one::One;
@@ -7,8 +7,8 @@ pub const R: CM31 = CM31 { a: M31 { inner: 2 }, b: M31 { inner: 1 } };
 
 #[derive(Copy, Drop, Debug, PartialEq, Eq)]
 pub struct QM31 {
-    a: CM31,
-    b: CM31,
+    pub a: CM31,
+    pub b: CM31,
 }
 
 #[generate_trait]
@@ -18,6 +18,16 @@ pub impl QM31Impl of QM31Trait {
         let [a, b, c, d] = arr;
         QM31 { a: CM31 { a: a, b: b }, b: CM31 { a: c, b: d } }
     }
+
+    fn from_u32(arr: [u32; 4]) -> QM31 {
+        let [a, b, c, d] = arr;
+        let a_mod_p = M31Impl::reduce_u32(a);
+        let b_mod_p = M31Impl::reduce_u32(b);
+        let c_mod_p = M31Impl::reduce_u32(c);
+        let d_mod_p = M31Impl::reduce_u32(d);
+
+        QM31 { a: CM31 { a: a_mod_p, b: b_mod_p }, b: CM31 { a: c_mod_p, b: d_mod_p } }
+    }
     #[inline]
     fn to_array(self: QM31) -> [M31; 4] {
         [self.a.a, self.a.b, self.b.a, self.b.b]
@@ -103,7 +113,8 @@ pub fn qm31(a: u32, b: u32, c: u32, d: u32) -> QM31 {
 
 #[cfg(test)]
 mod tests {
-    use super::{QM31, qm31, QM31Trait};
+    use super::{QM31, qm31, QM31Trait, QM31Impl};
+    use super::{CM31, cm31};
     use super::super::m31::{m31, P, M31Trait};
 
     #[test]
@@ -125,4 +136,15 @@ mod tests {
         assert_eq!(qm1 * m.inverse().into(), qm1 * qm.inverse());
         assert_eq!(qm1.mul_m31(m), qm1 * m.into());
     }
+
+    #[test]
+    fn test_qm31_from_u32() {
+        let arr = [2147483648, 2, 3, 4];
+        let felt = QM31Impl::from_u32(arr);
+        let expected_felt = QM31 {
+            a: CM31 { a: m31(1), b: m31(2) }, b: CM31 { a: m31(3), b: m31(4) }
+        };
+
+        assert_eq!(felt, expected_felt)
+    }
 }