Add support for exact division (#1256)

This adds `div_exact` and `div_exact_vartime` methods for `Uint`,
`BoxedUint`, and `Limb`. There is currently no trait support.

For cases where it applies, exact division can be much faster than
div/rem. In my tests this is about 40% faster for a 4096-bit
`BoxedUint`.

The algorithm is similar to `DivideByWord` in Modern Computer
Arithmetic (https://arxiv.org/pdf/1004.4710) but I don't have a
reference for the adjustments to support a multi-word divisor.

---------

Signed-off-by: Andrew Whitehead <cywolf@gmail.com>

Author: andrewwhitehead

benches/boxed_uint.rs

@@ -232,6 +232,42 @@ fn bench_division(c: &mut Criterion) {
         );
     });
 
+    group.bench_function("boxed_div_exact", |b| {
+        b.iter_batched(
+            || {
+                (
+                    BoxedUint::max(UINT_BITS),
+                    NonZero::new(BoxedUint::random_bits_with_precision(
+                        &mut rng,
+                        UINT_BITS / 2,
+                        UINT_BITS,
+                    ))
+                    .unwrap(),
+                )
+            },
+            |(x, y)| black_box(x.div_exact(&y)),
+            BatchSize::SmallInput,
+        );
+    });
+
+    group.bench_function("boxed_div_exact_vartime", |b| {
+        b.iter_batched(
+            || {
+                (
+                    BoxedUint::max(UINT_BITS),
+                    NonZero::new(BoxedUint::random_bits_with_precision(
+                        &mut rng,
+                        UINT_BITS / 2,
+                        UINT_BITS,
+                    ))
+                    .unwrap(),
+                )
+            },
+            |(x, y)| black_box(x.div_exact_vartime(&y)),
+            BatchSize::SmallInput,
+        );
+    });
+
     group.finish();
 }
 

benches/uint.rs

@@ -393,6 +393,44 @@ fn bench_division(c: &mut Criterion) {
         );
     });
 
+    group.bench_function("div exact, U256/U128", |b| {
+        b.iter_batched(
+            || {
+                (
+                    U256::random_from_rng(&mut rng),
+                    NonZero::<U128>::random_from_rng(&mut rng),
+                )
+            },
+            |(x, y)| x.div_exact(&y),
+            BatchSize::SmallInput,
+        );
+    });
+
+    group.bench_function("div exact, U256/U128 (in U256)", |b| {
+        b.iter_batched(
+            || {
+                let x = U256::random_from_rng(&mut rng);
+                let y_half = U128::random_from_rng(&mut rng);
+                let y: U256 = (y_half, U128::ZERO).into();
+                (x, NonZero::new(y).unwrap())
+            },
+            |(x, y)| x.div_exact(&y),
+            BatchSize::SmallInput,
+        );
+    });
+
+    group.bench_function("div exact, U256/U128 (in U512)", |b| {
+        b.iter_batched(
+            || {
+                let x = U256::random_from_rng(&mut rng);
+                let y: U512 = U128::random_from_rng(&mut rng).resize();
+                (x, NonZero::new(y).unwrap())
+            },
+            |(x, y)| x.div_exact(&y),
+            BatchSize::SmallInput,
+        );
+    });
+
     group.bench_function("div/rem_vartime, U256/U128, full size", |b| {
         b.iter_batched(
             || {
@@ -405,6 +443,18 @@ fn bench_division(c: &mut Criterion) {
         );
     });
 
+    group.bench_function("div exact vartime, U256/U128, full size", |b| {
+        b.iter_batched(
+            || {
+                let x = U256::random_from_rng(&mut rng);
+                let y = U256::from((NonZero::<U128>::random_from_rng(&mut rng).get(), U128::ZERO));
+                (x, NonZero::new(y).unwrap())
+            },
+            |(x, y)| x.div_exact_vartime(&y),
+            BatchSize::SmallInput,
+        );
+    });
+
     group.bench_function("rem, U256/U128", |b| {
         b.iter_batched(
             || {
@@ -510,6 +560,19 @@ fn bench_division(c: &mut Criterion) {
         );
     });
 
+    group.bench_function("div exact vartime, U256/Limb, full size", |b| {
+        b.iter_batched(
+            || {
+                let x = U256::random_from_rng(&mut rng);
+                let y_small = Limb::random_from_rng(&mut rng);
+                let y = U256::from_word(y_small.0);
+                (x, NonZero::new(y).unwrap())
+            },
+            |(x, y)| x.div_exact_vartime(&y),
+            BatchSize::SmallInput,
+        );
+    });
+
     group.bench_function("div/rem, U256/Limb, single limb", |b| {
         b.iter_batched(
             || {

src/limb.rs

@@ -13,6 +13,7 @@ mod div;
 mod encoding;
 mod from;
 mod gcd;
+mod invert_mod;
 mod mul;
 mod neg;
 mod shl;

src/limb/div.rs

@@ -42,6 +42,30 @@ impl Limb {
         let rem = self.div_rem(rhs_nz).1;
         CtOption::new(rem, is_nz)
     }
+
+    /// Exactly divides `self` by `rhs`, returning `CtOption::none()` if `self` is not divisible by `rhs`.
+    #[must_use]
+    pub const fn div_exact(&self, rhs: NonZero<Limb>) -> CtOption<Self> {
+        let mut quo = *self;
+        let mut div = rhs.get_copy();
+        let exact = UintRef::new_mut(slice::from_mut(&mut quo))
+            .div_exact(UintRef::new_mut(slice::from_mut(&mut div)));
+        CtOption::new(quo, exact)
+    }
+
+    /// Exactly divides `self` by `rhs`, returning `CtOption::none()` if `self` is not divisible by `rhs`.
+    ///
+    /// This is variable-time only with respect to `rhs`.
+    ///
+    /// When used with a fixed `rhs`, this function is constant-time with respect to `self`.
+    #[must_use]
+    pub const fn div_exact_vartime(&self, rhs: NonZero<Limb>) -> CtOption<Self> {
+        let mut quo = *self;
+        let mut div = rhs.get_copy();
+        let exact = UintRef::new_mut(slice::from_mut(&mut quo))
+            .div_exact_vartime(UintRef::new_mut(slice::from_mut(&mut div)));
+        CtOption::new(quo, exact)
+    }
 }
 
 impl CheckedDiv for Limb {
@@ -285,6 +309,17 @@ mod tests {
         let n = Limb::from_u32(0xffff_ffff);
         let d = NonZero::new(Limb::from_u32(0xfffe)).expect("ensured non-zero");
         assert_eq!(n.div_rem(d), (Limb::from_u32(0x10002), Limb::from_u32(0x3)));
+
+        assert_eq!(n.div_exact(d).into_option(), None);
+        assert_eq!(n.div_exact_vartime(d).into_option(), None);
+
+        let d = NonZero::new(Limb::from_u32(0xffff)).expect("ensured non-zero");
+        assert_eq!(n.div_rem(d), (Limb::from_u32(0x10001), Limb::from_u32(0)));
+        assert_eq!(n.div_exact(d).into_option(), Some(Limb::from_u32(0x10001)));
+        assert_eq!(
+            n.div_exact_vartime(d).into_option(),
+            Some(Limb::from_u32(0x10001))
+        );
     }
 
     #[test]

src/limb/invert_mod.rs

@@ -0,0 +1,17 @@
+use super::Limb;
+use crate::{Odd, primitives};
+
+impl Odd<Limb> {
+    /// Returns the multiplicative inverse of the argument modulo 2^N, where
+    /// 2^N is the capacity of a [`Limb`].
+    pub(crate) const fn multiplicative_inverse(self) -> Limb {
+        cpubits::cpubits! {
+            32 => {
+                Limb(primitives::u32_invert_odd(self.as_ref().0))
+            }
+            64 => {
+                Limb(primitives::u64_invert_odd(self.as_ref().0))
+            }
+        }
+    }
+}

src/primitives.rs

@@ -94,9 +94,62 @@ pub(crate) const fn u32_bits(n: u32) -> u32 {
     u32::BITS - n.leading_zeros()
 }
 
+/// Return a `Choice` representing whether `a < b`.
+#[allow(clippy::cast_possible_truncation)]
+#[cfg(target_pointer_width = "32")]
+#[inline]
+pub(crate) const fn usize_lt(a: usize, b: usize) -> Choice {
+    Choice::from_u32_lt(a as u32, b as u32)
+}
+
+/// Return a `Choice` representing whether `a < b`.
+#[allow(clippy::cast_possible_truncation)]
+#[cfg(target_pointer_width = "64")]
+#[inline]
+pub(crate) const fn usize_lt(a: usize, b: usize) -> Choice {
+    Choice::from_u64_lt(a as u64, b as u64)
+}
+
+cpubits::cpubits! {
+    32 => {
+        /// Returns the multiplicative inverse of the argument modulo 2^32.
+        ///
+        /// For correct results, the input `value` must be odd.
+        #[must_use]
+        pub(crate) const fn u32_invert_odd(value: u32) -> u32 {
+            debug_assert!(value & 1 == 1, "value must be odd");
+            let x = value.wrapping_mul(3) ^ 2;
+            let y = 1u32.wrapping_sub(x.wrapping_mul(value));
+            let (x, y) = (x.wrapping_mul(y.wrapping_add(1)), y.wrapping_mul(y));
+            let (x, y) = (x.wrapping_mul(y.wrapping_add(1)), y.wrapping_mul(y));
+            x.wrapping_mul(y.wrapping_add(1))
+        }
+    }
+}
+
+/// Returns the multiplicative inverse of the argument modulo 2^64. The implementation is based
+/// on Hurchalla's method for computing the multiplicative inverse modulo a power of two, and
+/// is essentially an optimized Newton iteration.
+///
+/// For correct results, the input `value` must be odd.
+///
+/// For better understanding the implementation, the following paper is recommended:
+/// J. Hurchalla, "An Improved Integer Multiplicative Inverse (modulo 2^w)",
+/// <https://arxiv.org/abs/2204.04342>
+#[must_use]
+pub(crate) const fn u64_invert_odd(value: u64) -> u64 {
+    debug_assert!(value & 1 == 1, "value must be odd");
+    let x = value.wrapping_mul(3) ^ 2;
+    let y = 1u64.wrapping_sub(x.wrapping_mul(value));
+    let (x, y) = (x.wrapping_mul(y.wrapping_add(1)), y.wrapping_mul(y));
+    let (x, y) = (x.wrapping_mul(y.wrapping_add(1)), y.wrapping_mul(y));
+    let (x, y) = (x.wrapping_mul(y.wrapping_add(1)), y.wrapping_mul(y));
+    x.wrapping_mul(y.wrapping_add(1))
+}
+
 #[cfg(test)]
 mod tests {
-    use super::{u32_max, u32_min, u32_rem};
+    use super::{u32_max, u32_min, u32_rem, usize_lt};
     use crate::Word;
 
     #[test]
@@ -133,4 +186,25 @@ mod tests {
         assert_eq!(u32_rem(7, 5), 2);
         assert_eq!(u32_rem(101, 5), 1);
     }
+
+    #[test]
+    fn test_usize_const_lt() {
+        assert!(usize_lt(0, 5).to_bool_vartime());
+        assert!(!usize_lt(7, 0).to_bool_vartime());
+        assert!(!usize_lt(7, 5).to_bool_vartime());
+        assert!(!usize_lt(7, 7).to_bool_vartime());
+    }
+
+    cpubits::cpubits! {
+        32 => {
+            #[test]
+            fn test_u32_invert_odd() {
+                use super::u32_invert_odd;
+
+                assert_eq!(u32_invert_odd(1), 1);
+                assert_eq!(u32_invert_odd(5).wrapping_mul(5), 1);
+                assert_eq!(u32_invert_odd(u32::MAX).wrapping_mul(u32::MAX), 1);
+            }
+        }
+    }
 }

src/uint/boxed/div.rs

@@ -63,6 +63,15 @@ impl BoxedUint {
         rem
     }
 
+    /// Exactly divides `self` by `rhs`, returning `CtOption::none()` if `self` is not divisible by `rhs`.
+    #[must_use]
+    pub fn div_exact<Rhs: ToUnsigned + ?Sized>(&self, rhs: &NonZero<Rhs>) -> CtOption<Self> {
+        let mut quo = self.clone();
+        let mut div = rhs.to_unsigned().get();
+        let exact = quo.as_mut_uint_ref().div_exact(div.as_mut_uint_ref());
+        CtOption::new(quo, exact)
+    }
+
     /// Computes self / rhs, returns the quotient and remainder.
     ///
     /// Variable-time with respect to `rhs`
@@ -121,6 +130,20 @@ impl BoxedUint {
         rem
     }
 
+    /// Exactly divides `self` by `rhs`, returning `CtOption::none()` if `self` is not divisible by `rhs`.
+    #[must_use]
+    pub fn div_exact_vartime<Rhs: ToUnsigned + ?Sized>(
+        &self,
+        rhs: &NonZero<Rhs>,
+    ) -> CtOption<Self> {
+        let mut quo = self.clone();
+        let mut div = rhs.to_unsigned().get();
+        let exact = quo
+            .as_mut_uint_ref()
+            .div_exact_vartime(div.as_mut_uint_ref());
+        CtOption::new(quo, exact)
+    }
+
     /// Wrapped division is just normal division i.e. `self` / `rhs`
     /// There’s no way wrapping could ever happen.
     ///

src/uint/boxed/lcm.rs

@@ -8,15 +8,19 @@ impl Lcm for BoxedUint {
     fn lcm(&self, rhs: &Self) -> Self {
         let (lhs_nz, _) = self.to_nz_or_one();
         let gcd_nz = lhs_nz.gcd(rhs);
-        self.wrapping_div(&gcd_nz).concatenating_mul(rhs)
+        self.div_exact(&gcd_nz)
+            .expect("invalid gcd")
+            .concatenating_mul(rhs)
     }
 
     fn lcm_vartime(&self, rhs: &Self) -> Self {
         let (Some(lhs_nz), false) = (self.as_nz_vartime(), rhs.is_zero_vartime()) else {
             return BoxedUint::zero_with_precision(self.bits_precision() + rhs.bits_precision());
         };
         let gcd_nz = lhs_nz.gcd_vartime(rhs);
-        self.wrapping_div_vartime(&gcd_nz).concatenating_mul(rhs)
+        self.div_exact_vartime(&gcd_nz)
+            .expect("invalid gcd")
+            .concatenating_mul(rhs)
     }
 }