Implement partial reduce for m31 (#333)

src/core/fields/m31.rs

@@ -22,6 +22,12 @@ impl M31 {
         (result.square() == *self).then_some(result)
     }
 
+    /// Assumes that `val` is in the range [0, 2 * `P`) and returns `val` % `P`.
+    pub fn partial_reduce(val: u32) -> Self {
+        Self(val.checked_sub(P).unwrap_or(val))
+    }
+
+    /// Assumes that `val` is in the range [0, `P`.pow(2)) and returns `val` % `P`.
     pub fn reduce(val: u64) -> Self {
         Self((((((val >> MODULUS_BITS) + val + 1) >> MODULUS_BITS) + val) & (P as u64)) as u32)
     }
@@ -41,23 +47,23 @@ impl Add for M31 {
     type Output = Self;
 
     fn add(self, rhs: Self) -> Self::Output {
-        Self::reduce((self.0 as u64) + (rhs.0 as u64))
+        Self::partial_reduce(self.0 + rhs.0)
     }
 }
 
 impl Neg for M31 {
     type Output = Self;
 
     fn neg(self) -> Self::Output {
-        Self::reduce(P as u64 - (self.0 as u64))
+        Self::partial_reduce(P - self.0)
     }
 }
 
 impl Sub for M31 {
     type Output = Self;
 
     fn sub(self, rhs: Self) -> Self::Output {
-        Self::reduce((self.0 as u64) + (P as u64) - (rhs.0 as u64))
+        Self::partial_reduce(self.0 + P - rhs.0)
     }
 }