Fix random_mod performance for small moduli; NonZero moduli

Changes the modulus for `random_mod` to be `NonZero`.

Changes the algorithm used by `random_mod` to generate a number that can
be represented by the same number of bytes as the modulus.

Such a number can still be larger than the modulus, but is much more
likely not to overflow than a "full-width" number provided the modulus
is small relative to the width.

Closes #3

src/limb.rs

@@ -8,6 +8,7 @@ mod bit_and;
 mod bit_not;
 mod bit_or;
 mod bit_xor;
+mod bits;
 mod cmp;
 mod encoding;
 mod from;
@@ -99,6 +100,12 @@ impl Limb {
     /// Maximum value this [`Limb`] can express.
     pub const MAX: Self = Limb(Inner::MAX);
 
+    /// Size of the inner integer in bits.
+    pub const BIT_SIZE: usize = BIT_SIZE;
+
+    /// Size of the inner integer in bytes.
+    pub const BYTE_SIZE: usize = BYTE_SIZE;
+
     /// Return `a` if `c`!=0 or `b` if `c`==0.
     ///
     /// Const-friendly: we can't yet use `subtle` in `const fn` contexts.

src/limb/bits.rs

@@ -0,0 +1,8 @@
+use super::{Limb, BIT_SIZE};
+
+impl Limb {
+    /// Calculate the number of bits needed to represent this number.
+    pub const fn bits(self) -> usize {
+        BIT_SIZE - (self.0.leading_zeros() as usize)
+    }
+}

src/limb/rand.rs

@@ -1,20 +1,45 @@
 //! Random number generator support
 
-use super::Limb;
+use super::{Limb, BYTE_SIZE};
+use crate::{Encoding, NonZero, Random, RandomMod};
 use rand_core::{CryptoRng, RngCore};
+use subtle::ConstantTimeLess;
 
-impl Limb {
-    /// Generate a random limb
+impl Random for Limb {
     #[cfg(target_pointer_width = "32")]
     #[cfg_attr(docsrs, doc(cfg(feature = "rand")))]
-    pub fn random(mut rng: impl CryptoRng + RngCore) -> Self {
+    fn random(mut rng: impl CryptoRng + RngCore) -> Self {
         Self(rng.next_u32())
     }
 
-    /// Generate a random limb
     #[cfg(target_pointer_width = "64")]
     #[cfg_attr(docsrs, doc(cfg(feature = "rand")))]
-    pub fn random(mut rng: impl CryptoRng + RngCore) -> Self {
+    fn random(mut rng: impl CryptoRng + RngCore) -> Self {
         Self(rng.next_u64())
     }
 }
+
+impl RandomMod for Limb {
+    fn random_mod(mut rng: impl CryptoRng + RngCore, modulus: &NonZero<Self>) -> Self {
+        let mut bytes = <Self as Encoding>::Repr::default();
+
+        // TODO(tarcieri): use `div_ceil` when available
+        // See: https://github.com/rust-lang/rust/issues/88581
+        let mut n_bytes = modulus.bits() / 8;
+
+        // Ensure the randomly generated value can always be larger than
+        // the modulus in order to ensure a uniform distribution
+        if n_bytes < BYTE_SIZE {
+            n_bytes += 1;
+        }
+
+        loop {
+            rng.fill_bytes(&mut bytes[..n_bytes]);
+            let n = Limb::from_le_bytes(bytes);
+
+            if n.ct_lt(modulus).into() {
+                return n;
+            }
+        }
+    }
+}

src/traits.rs

@@ -68,7 +68,7 @@ pub trait Random: Sized {
 /// Modular random number generation support.
 #[cfg(feature = "rand")]
 #[cfg_attr(docsrs, doc(cfg(feature = "rand")))]
-pub trait RandomMod: Sized {
+pub trait RandomMod: Sized + Zero {
     /// Generate a cryptographically secure random number which is less than
     /// a given `modulus`.
     ///
@@ -80,7 +80,7 @@ pub trait RandomMod: Sized {
     /// issue so long as the underlying random number generator is truly a
     /// [`CryptoRng`], where previous outputs are unrelated to subsequent
     /// outputs and do not reveal information about the RNG's internal state.
-    fn random_mod(rng: impl CryptoRng + RngCore, modulus: &Self) -> Self;
+    fn random_mod(rng: impl CryptoRng + RngCore, modulus: &NonZero<Self>) -> Self;
 }
 
 /// Compute `self + rhs mod p`.

src/uint/bits.rs

@@ -2,8 +2,8 @@ use crate::limb::{Inner, BIT_SIZE};
 use crate::{Limb, UInt};
 
 impl<const LIMBS: usize> UInt<LIMBS> {
-    /// Calculate the number of bits needed to represent this number
-    pub(crate) const fn bits(self) -> Inner {
+    /// Calculate the number of bits needed to represent this number.
+    pub const fn bits(self) -> usize {
         let mut i = LIMBS - 1;
         while i > 0 && self.limbs[i].0 == 0 {
             i -= 1;
@@ -17,6 +17,6 @@ impl<const LIMBS: usize> UInt<LIMBS> {
             Limb::ZERO,
             !self.limbs[0].is_nonzero() & !Limb(i as Inner).is_nonzero(),
         )
-        .0
+        .0 as usize
     }
 }

src/uint/rand.rs

@@ -2,7 +2,7 @@
 // TODO(tarcieri): use `Random` and `RandomMod` impls exclusively in next breaking release
 
 use super::UInt;
-use crate::{Limb, Random, RandomMod};
+use crate::{Limb, NonZero, Random, RandomMod};
 use rand_core::{CryptoRng, RngCore};
 use subtle::ConstantTimeLess;
 
@@ -11,7 +11,7 @@ use subtle::ConstantTimeLess;
 impl<const LIMBS: usize> UInt<LIMBS> {
     /// Generate a cryptographically secure random [`UInt`].
     pub fn random(mut rng: impl CryptoRng + RngCore) -> Self {
-        let mut limbs = [Limb::default(); LIMBS];
+        let mut limbs = [Limb::ZERO; LIMBS];
 
         for limb in &mut limbs {
             *limb = Limb::random(&mut rng)
@@ -31,9 +31,30 @@ impl<const LIMBS: usize> UInt<LIMBS> {
     /// issue so long as the underlying random number generator is truly a
     /// [`CryptoRng`], where previous outputs are unrelated to subsequent
     /// outputs and do not reveal information about the RNG's internal state.
-    pub fn random_mod(mut rng: impl CryptoRng + RngCore, modulus: &Self) -> Self {
+    pub fn random_mod(mut rng: impl CryptoRng + RngCore, modulus: &NonZero<Self>) -> Self {
+        let mut n = Self::ZERO;
+
+        // TODO(tarcieri): use `div_ceil` when available
+        // See: https://github.com/rust-lang/rust/issues/88581
+        let mut n_limbs = modulus.bits() / Limb::BIT_SIZE;
+        if n_limbs < LIMBS {
+            n_limbs += 1;
+        }
+
+        // Compute the highest limb of `modulus` as a `NonZero`
+        // Note: this will always be `Some`, but this impl avoids panics
+        let modulus_hi = Option::from(NonZero::new(modulus.limbs[n_limbs.saturating_sub(1)]));
+
         loop {
-            let n = Self::random(&mut rng);
+            for i in 0..n_limbs {
+                n.limbs[i] = if i + 1 == n_limbs {
+                    modulus_hi
+                        .map(|hi| Limb::random_mod(&mut rng, &hi))
+                        .unwrap_or_default()
+                } else {
+                    Limb::random(&mut rng)
+                }
+            }
 
             if n.ct_lt(modulus).into() {
                 return n;
@@ -51,7 +72,33 @@ impl<const LIMBS: usize> Random for UInt<LIMBS> {
 
 #[cfg_attr(docsrs, doc(cfg(feature = "rand")))]
 impl<const LIMBS: usize> RandomMod for UInt<LIMBS> {
-    fn random_mod(rng: impl CryptoRng + RngCore, modulus: &Self) -> Self {
+    fn random_mod(rng: impl CryptoRng + RngCore, modulus: &NonZero<Self>) -> Self {
         Self::random_mod(rng, modulus)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::{NonZero, U256};
+    use rand_core::SeedableRng;
+
+    #[test]
+    fn random_mod() {
+        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(1);
+
+        // Ensure `random_mod` runs in a reasonable amount of time
+        let modulus = NonZero::new(U256::from(42u8)).unwrap();
+        let res = U256::random_mod(&mut rng, &modulus);
+
+        // Sanity check that the return value isn't zero
+        assert_ne!(res, U256::ZERO);
+
+        // Ensure `random_mod` runs in a reasonable amount of time
+        // when the modulus is larger than 1 limb
+        let modulus = NonZero::new(U256::from(0x10000000000000001u128)).unwrap();
+        let res = U256::random_mod(&mut rng, &modulus);
+
+        // Sanity check that the return value isn't zero
+        assert_ne!(res, U256::ZERO);
+    }
+}

src/uint/sub_mod.rs

@@ -45,18 +45,18 @@ impl_sub_mod!(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12);
 
 #[cfg(all(test, feature = "rand"))]
 mod tests {
-    use crate::UInt;
+    use crate::{Limb, NonZero, Random, UInt};
+    use rand_core::SeedableRng;
 
     macro_rules! test_sub_mod {
         ($size:expr, $test_name:ident) => {
             #[test]
             fn $test_name() {
-                use crate::Limb;
-                use rand_core::SeedableRng;
-
                 let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(1);
-
-                let moduli = [UInt::<$size>::random(&mut rng), UInt::random(&mut rng)];
+                let moduli = [
+                    NonZero::<UInt<$size>>::random(&mut rng),
+                    NonZero::<UInt<$size>>::random(&mut rng),
+                ];
 
                 for p in &moduli {
                     let base_cases = [
@@ -79,7 +79,7 @@ mod tests {
                             let (a, b) = if a < b { (b, a) } else { (a, b) };
 
                             let c = a.sub_mod(&b, p);
-                            assert!(c < *p, "not reduced");
+                            assert!(c < **p, "not reduced");
                             assert_eq!(c, a.wrapping_sub(&b), "result incorrect");
                         }
                     }
@@ -89,10 +89,10 @@ mod tests {
                         let b = UInt::<$size>::random_mod(&mut rng, p);
 
                         let c = a.sub_mod(&b, p);
-                        assert!(c < *p, "not reduced: {} >= {} ", c, p);
+                        assert!(c < **p, "not reduced: {} >= {} ", c, p);
 
                         let x = a.wrapping_sub(&b);
-                        if a >= b && x < *p {
+                        if a >= b && x < **p {
                             assert_eq!(c, x, "incorrect result");
                         }
                     }