Add DivRemLimb and RemLimb traits

Author: fjarri

src/traits.rs

@@ -8,7 +8,7 @@ pub use num_traits::{
 
 pub(crate) use sealed::PrecomputeInverterWithAdjuster;
 
-use crate::{Limb, NonZero, Odd};
+use crate::{Limb, NonZero, Odd, Reciprocal};
 use core::fmt::Debug;
 use core::ops::{
     Add, AddAssign, BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Div, DivAssign,
@@ -110,6 +110,7 @@ pub trait Integer:
     + for<'a> Div<&'a NonZero<Self>, Output = Self>
     + DivAssign<NonZero<Self>>
     + for<'a> DivAssign<&'a NonZero<Self>>
+    + DivRemLimb
     + Eq
     + From<u8>
     + From<u16>
@@ -124,6 +125,7 @@ pub trait Integer:
     + Ord
     + Rem<NonZero<Self>, Output = Self>
     + for<'a> Rem<&'a NonZero<Self>, Output = Self>
+    + RemLimb
     + Send
     + Sized
     + Shl<u32, Output = Self>
@@ -453,6 +455,29 @@ pub trait SquareAssign {
     fn square_assign(&mut self);
 }
 
+/// Support for optimized division by a single limb.
+pub trait DivRemLimb: Sized {
+    /// Computes `self / rhs` using a pre-made reciprocal,
+    /// returns the quotient (q) and remainder (r).
+    fn div_rem_limb(&self, rhs: NonZero<Limb>) -> (Self, Limb) {
+        self.div_rem_limb_with_reciprocal(&Reciprocal::new(rhs))
+    }
+
+    /// Computes `self / rhs`, returns the quotient (q) and remainder (r).
+    fn div_rem_limb_with_reciprocal(&self, reciprocal: &Reciprocal) -> (Self, Limb);
+}
+
+/// Support for optimized division by a single limb.
+pub trait RemLimb: Sized {
+    /// Computes `self % rhs` using a pre-made reciprocal.
+    fn rem_limb(&self, rhs: NonZero<Limb>) -> Limb {
+        self.rem_limb_with_reciprocal(&Reciprocal::new(rhs))
+    }
+
+    /// Computes `self % rhs`.
+    fn rem_limb_with_reciprocal(&self, reciprocal: &Reciprocal) -> Limb;
+}
+
 /// Constant-time exponentiation.
 pub trait Pow<Exponent> {
     /// Raises to the `exponent` power.

src/uint/boxed/div.rs

@@ -1,23 +1,36 @@
 //! [`BoxedUint`] division operations.
 
 use crate::{
-    uint::boxed, BoxedUint, CheckedDiv, ConstantTimeSelect, Limb, NonZero, Reciprocal, Wrapping,
+    uint::boxed, BoxedUint, CheckedDiv, ConstantTimeSelect, DivRemLimb, Limb, NonZero, Reciprocal,
+    RemLimb, Wrapping,
 };
 use core::ops::{Div, DivAssign, Rem, RemAssign};
 use subtle::{Choice, ConstantTimeEq, ConstantTimeLess, CtOption};
 
 impl BoxedUint {
-    /// Computes `self` / `rhs` using a pre-made reciprocal,
+    /// Computes `self / rhs` using a pre-made reciprocal,
     /// returns the quotient (q) and remainder (r).
     pub fn div_rem_limb_with_reciprocal(&self, reciprocal: &Reciprocal) -> (Self, Limb) {
         boxed::div_limb::div_rem_limb_with_reciprocal(self, reciprocal)
     }
 
-    /// Computes `self` / `rhs`, returns the quotient (q) and remainder (r).
+    /// Computes `self / rhs`, returns the quotient (q) and remainder (r).
     pub fn div_rem_limb(&self, rhs: NonZero<Limb>) -> (Self, Limb) {
         boxed::div_limb::div_rem_limb_with_reciprocal(self, &Reciprocal::new(rhs))
     }
 
+    /// Computes `self % rhs` using a pre-made reciprocal.
+    #[inline(always)]
+    pub fn rem_limb_with_reciprocal(&self, reciprocal: &Reciprocal) -> Limb {
+        boxed::div_limb::rem_limb_with_reciprocal(self, reciprocal)
+    }
+
+    /// Computes `self % rhs`.
+    #[inline(always)]
+    pub fn rem_limb(&self, rhs: NonZero<Limb>) -> Limb {
+        boxed::div_limb::rem_limb_with_reciprocal(self, &Reciprocal::new(rhs))
+    }
+
     /// Computes self / rhs, returns the quotient, remainder.
     pub fn div_rem(&self, rhs: &NonZero<Self>) -> (Self, Self) {
         // Since `rhs` is nonzero, this should always hold.
@@ -295,6 +308,18 @@ impl RemAssign<NonZero<BoxedUint>> for BoxedUint {
     }
 }
 
+impl DivRemLimb for BoxedUint {
+    fn div_rem_limb_with_reciprocal(&self, reciprocal: &Reciprocal) -> (Self, Limb) {
+        Self::div_rem_limb_with_reciprocal(self, reciprocal)
+    }
+}
+
+impl RemLimb for BoxedUint {
+    fn rem_limb_with_reciprocal(&self, reciprocal: &Reciprocal) -> Limb {
+        Self::rem_limb_with_reciprocal(self, reciprocal)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::{BoxedUint, NonZero};

src/uint/boxed/div_limb.rs

@@ -23,3 +23,18 @@ pub(crate) fn div_rem_limb_with_reciprocal(
     }
     (BoxedUint { limbs: q.into() }, Limb(r >> reciprocal.shift()))
 }
+
+/// Divides `u` by the divisor encoded in the `reciprocal`, and returns the remainder.
+#[inline(always)]
+pub(crate) fn rem_limb_with_reciprocal(u: &BoxedUint, reciprocal: &Reciprocal) -> Limb {
+    let (u_shifted, u_hi) = u.shl_limb(reciprocal.shift());
+    let mut r = u_hi.0;
+
+    let mut j = u.limbs.len();
+    while j > 0 {
+        j -= 1;
+        let (_, rj) = div2by1(r, u_shifted.as_limbs()[j].0, reciprocal);
+        r = rj;
+    }
+    Limb(r >> reciprocal.shift())
+}

src/uint/div.rs

@@ -1,24 +1,36 @@
 //! [`Uint`] division operations.
 
-use super::div_limb::{div_rem_limb_with_reciprocal, Reciprocal};
-use crate::{CheckedDiv, ConstChoice, Limb, NonZero, Uint, Word, Wrapping};
+use super::div_limb::{div_rem_limb_with_reciprocal, rem_limb_with_reciprocal, Reciprocal};
+use crate::{CheckedDiv, ConstChoice, DivRemLimb, Limb, NonZero, RemLimb, Uint, Word, Wrapping};
 use core::ops::{Div, DivAssign, Rem, RemAssign};
 use subtle::CtOption;
 
 impl<const LIMBS: usize> Uint<LIMBS> {
-    /// Computes `self` / `rhs` using a pre-made reciprocal,
+    /// Computes `self / rhs` using a pre-made reciprocal,
     /// returns the quotient (q) and remainder (r).
     #[inline(always)]
     pub const fn div_rem_limb_with_reciprocal(&self, reciprocal: &Reciprocal) -> (Self, Limb) {
         div_rem_limb_with_reciprocal(self, reciprocal)
     }
 
-    /// Computes `self` / `rhs`, returns the quotient (q) and remainder (r).
+    /// Computes `self / rhs`, returns the quotient (q) and remainder (r).
     #[inline(always)]
     pub const fn div_rem_limb(&self, rhs: NonZero<Limb>) -> (Self, Limb) {
         div_rem_limb_with_reciprocal(self, &Reciprocal::new(rhs))
     }
 
+    /// Computes `self % rhs` using a pre-made reciprocal.
+    #[inline(always)]
+    pub const fn rem_limb_with_reciprocal(&self, reciprocal: &Reciprocal) -> Limb {
+        rem_limb_with_reciprocal(self, reciprocal)
+    }
+
+    /// Computes `self % rhs`.
+    #[inline(always)]
+    pub const fn rem_limb(&self, rhs: NonZero<Limb>) -> Limb {
+        rem_limb_with_reciprocal(self, &Reciprocal::new(rhs))
+    }
+
     /// Computes `self` / `rhs`, returns the quotient (q) and the remainder (r)
     ///
     /// This function is constant-time with respect to both `self` and `rhs`.
@@ -584,6 +596,18 @@ impl<const LIMBS: usize> RemAssign<&NonZero<Uint<LIMBS>>> for Wrapping<Uint<LIMB
     }
 }
 
+impl<const LIMBS: usize> DivRemLimb for Uint<LIMBS> {
+    fn div_rem_limb_with_reciprocal(&self, reciprocal: &Reciprocal) -> (Self, Limb) {
+        Self::div_rem_limb_with_reciprocal(self, reciprocal)
+    }
+}
+
+impl<const LIMBS: usize> RemLimb for Uint<LIMBS> {
+    fn rem_limb_with_reciprocal(&self, reciprocal: &Reciprocal) -> Limb {
+        Self::rem_limb_with_reciprocal(self, reciprocal)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::{Limb, NonZero, Uint, Word, U256};

src/uint/div_limb.rs

@@ -236,6 +236,24 @@ pub(crate) const fn div_rem_limb_with_reciprocal<const L: usize>(
     (Uint::<L>::new(q), Limb(r >> reciprocal.shift))
 }
 
+/// Divides `u` by the divisor encoded in the `reciprocal`, and returns the remainder.
+#[inline(always)]
+pub(crate) const fn rem_limb_with_reciprocal<const L: usize>(
+    u: &Uint<L>,
+    reciprocal: &Reciprocal,
+) -> Limb {
+    let (u_shifted, u_hi) = u.shl_limb(reciprocal.shift);
+    let mut r = u_hi.0;
+
+    let mut j = L;
+    while j > 0 {
+        j -= 1;
+        let (_, rj) = div2by1(r, u_shifted.as_limbs()[j].0, reciprocal);
+        r = rj;
+    }
+    Limb(r >> reciprocal.shift)
+}
+
 #[cfg(test)]
 mod tests {
     use super::{div2by1, Reciprocal};