From 644770d6d96b54863e2bbc945cdb68d970ad252f Mon Sep 17 00:00:00 2001
From: Tom Grosso <tgrosso@eryx.co>
Date: Wed, 2 Oct 2024 12:29:28 -0300
Subject: [PATCH] Add qm31 constructor from u32

---
 stwo_cairo_verifier/src/fields/qm31.cairo | 31 +++++++++++++++++++----
 1 file changed, 26 insertions(+), 5 deletions(-)

diff --git a/stwo_cairo_verifier/src/fields/qm31.cairo b/stwo_cairo_verifier/src/fields/qm31.cairo
index a16ad7759..e64b15bcc 100644
--- 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]
@@ -30,7 +40,6 @@ pub impl QM31Impl of QM31Trait {
         let denom_inverse = denom.inverse();
         QM31 { a: self.a * denom_inverse, b: -self.b * denom_inverse }
     }
-
     fn complex_conjugate(self: QM31) -> QM31 {
         QM31 { a: self.a, b: -self.b }
     }
@@ -97,7 +106,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}; //::{cm31, CM31};
     use super::super::m31::{m31, P, M31Trait};
 
     #[test]
@@ -118,4 +128,15 @@ mod tests {
         assert_eq!(qm0_x_qm1 * qm1.inverse(), qm31(1, 2, 3, 4));
         assert_eq!(qm1 * m.inverse().into(), qm1 * qm.inverse());
     }
+
+    #[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)
+    }
 }