Extend to primes less than 2^31

ext/crates/fp/src/limb.rs

@@ -13,16 +13,17 @@ pub(crate) struct LimbBitIndexPair {
 
 /// Return the number of bits an element of $\mathbb{F}_P$ occupies in a limb.
 pub(crate) fn bit_length<P: Prime>(p: P) -> usize {
-    match p.as_u32() {
+    let p = p.as_u32() as u64;
+    match p {
         2 => 1,
-        _ => (32u32 - (p * (p - 1)).leading_zeros()) as usize,
+        _ => (BITS_PER_LIMB as u32 - (p * (p - 1)).leading_zeros()) as usize,
     }
 }
 
 /// If `l` is a limb of elements of $\\mathbb{F}_p$, then `l & bitmask::<P>()` is the value of the
 /// first entry of `l`.
 pub(crate) fn bitmask<P: Prime>(p: P) -> Limb {
-    (1 << bit_length(p)) - 1
+    ((1u128 << bit_length(p)) - 1) as Limb
 }
 
 // this function is never called if `odd-primes` is disabled
@@ -77,7 +78,7 @@ pub(crate) fn reduce<P: Prime>(p: P, limb: Limb) -> Limb {
             let d = m & BOTTOM_THREE_BITS;
             d + c - BOTTOM_TWO_BITS
         }
-        _ => pack(p, unpack(p, limb).map(|x| x % p.as_u32())),
+        _ => pack(p, unpack(p, limb).map(|x| (x % (p.as_u32() as u64)) as u32)),
     }
 }
 
@@ -104,14 +105,16 @@ pub(crate) fn pack<T: Iterator<Item = u32>, P: Prime>(p: P, entries: T) -> Limb
     result
 }
 
-/// Give an iterator over the entries of `limb`.
-pub(crate) fn unpack<P: Prime>(p: P, mut limb: Limb) -> impl Iterator<Item = u32> {
+/// Give an iterator over the entries of `limb`. We return `u64` instead of `u32` because the
+/// entries might not be reduced. In general we can only guarantee that the entries are less than p
+/// * (p - 1), so they might not fit into a single `u32`.
+pub(crate) fn unpack<P: Prime>(p: P, mut limb: Limb) -> impl Iterator<Item = u64> {
     let entries = entries_per_limb(p);
     let bit_length = bit_length(p);
     let bit_mask = bitmask(p);
 
     (0..entries).map(move |_| {
-        let result = (limb & bit_mask) as u32;
+        let result = limb & bit_mask;
         limb >>= bit_length;
         result
     })

ext/crates/fp/src/prime/mod.rs

@@ -255,8 +255,14 @@ mod validprime {
 
     impl ValidPrime {
         pub const fn new(p: u32) -> ValidPrime {
-            // We need the size restriction because we need p * (p - 1) to fit in a u32. See `limb::bit_length`.
-            assert!(p < (1 << 16), "Tried to construct a prime larger than 2^16");
+            // We need the size restriction because we need `bit_length(p)` to be smaller than 64.
+            // Otherwise, shifting a u64 by 64 bits is considered an overflow. We could special case
+            // these shifts to be setting to 0, but that doesn't seem worth it.
+            //
+            // Also, we could in theory support primes up to sqrt(2^63 - 1), but that makes the
+            // assert below more complicated, and primes up to 2^31 are good enough for most
+            // purposes anyway.
+            assert!(p < (1 << 31), "Tried to construct a prime larger than 2^31");
             assert!(is_prime(p), "Tried to construct a composite dynamic prime");
             ValidPrime { p }
         }

ext/crates/fp/src/vector/impl_fpvectorp.rs

@@ -120,7 +120,8 @@ impl<P: Prime> FpVectorP<P> {
                 for limb in &mut self.limbs {
                     *limb = limb::pack(
                         self.p,
-                        limb::unpack(self.p, *limb).map(|x| (x * c) % self.p.as_u32()),
+                        limb::unpack(self.p, *limb)
+                            .map(|x| ((x * c as u64) % (self.p.as_u32() as u64)) as u32),
                     );
                 }
             }

ext/crates/fp/src/vector/mod.rs

@@ -102,12 +102,12 @@ mod test {
         (0..dimension).map(|_| rng.gen_range(0..p)).collect()
     }
 
-    /// An arbitrary `ValidPrime` in the range `2..65536`.
+    /// An arbitrary `ValidPrime` in the range `2..(1 << 24)`, plus the largest prime that we support.
     fn arb_prime() -> impl Strategy<Value = ValidPrime> {
         static TEST_PRIMES: OnceLock<Vec<ValidPrime>> = OnceLock::new();
         let test_primes = TEST_PRIMES.get_or_init(|| {
             // Sieve of erathosthenes
-            const MAX: usize = 65536;
+            const MAX: usize = 1 << 24;
             let mut is_prime = Vec::new();
             is_prime.resize_with(MAX, || true);
             is_prime[0] = false;
@@ -122,6 +122,7 @@ mod test {
             (0..MAX)
                 .filter(|&i| is_prime[i])
                 .map(|p| ValidPrime::new_unchecked(p as u32))
+                .chain(std::iter::once(ValidPrime::new_unchecked(2147483647)))
                 .collect()
         });
         (0..test_primes.len()).prop_map(|i| test_primes[i])
@@ -212,7 +213,7 @@ mod test {
 
         #[test]
         fn test_bit_length(p in arb_prime()) {
-            bit_length(p);
+            prop_assert!(bit_length(p) <= 63);
         }
 
         #[test]
@@ -249,7 +250,7 @@ mod test {
             let mut v = FpVector::from_slice(p, &v_arr);
             v.scale(c);
             for entry in &mut v_arr {
-                *entry = (*entry * c) % p;
+                *entry = (((*entry as u64) * (c as u64)) % (p.as_u32() as u64)) as u32;
             }
             v.assert_list_eq(&v_arr);
         }
@@ -264,7 +265,7 @@ mod test {
             v.slice_mut(slice_start, slice_end).scale(c);
 
             for entry in &mut v_arr[slice_start..slice_end] {
-                *entry = (*entry * c) % p;
+                *entry = (((*entry as u64) * (c as u64)) % (p.as_u32() as u64)) as u32
             }
             v.assert_list_eq(&v_arr);
         }