Add a RemMixed trait and impl for Uints (#746)

Author: dvdplm

src/traits.rs

@@ -626,6 +626,12 @@ pub trait DivRemLimb: Sized {
     fn div_rem_limb_with_reciprocal(&self, reciprocal: &Reciprocal) -> (Self, Limb);
 }
 
+/// Support for calculating the remainder of two differently sized integers.
+pub trait RemMixed<Reductor>: Sized {
+    /// Calculate the remainder of `self` by the `reductor`.
+    fn rem_mixed(&self, reductor: &NonZero<Reductor>) -> Reductor;
+}
+
 /// Support for optimized division by a single limb.
 pub trait RemLimb: Sized {
     /// Computes `self % rhs` using a pre-made reciprocal.

src/uint.rs

@@ -439,7 +439,7 @@ impl_uint_concat_split_even! {
     U16384,
 }
 
-// Implement mixed concat and split for combinations not implemented by
+// Implement mixed concat, split and reduce for combinations not implemented by
 // impl_uint_concat_split_even. The numbers represent the size of each
 // component Uint in multiple of 64 bits. For example,
 // (U256, [1, 3]) will allow splitting U256 into (U64, U192) as well as

src/uint/boxed/div.rs

@@ -6,7 +6,7 @@ use crate::{
         div_limb::{div2by1, div3by2},
     },
     BoxedUint, CheckedDiv, ConstChoice, ConstantTimeSelect, DivRemLimb, Limb, NonZero, Reciprocal,
-    RemLimb, Wrapping,
+    RemLimb, RemMixed, Wrapping,
 };
 use core::ops::{Div, DivAssign, Rem, RemAssign};
 use subtle::CtOption;
@@ -391,6 +391,12 @@ impl RemLimb for BoxedUint {
     }
 }
 
+impl RemMixed<BoxedUint> for BoxedUint {
+    fn rem_mixed(&self, reductor: &NonZero<BoxedUint>) -> BoxedUint {
+        Self::div_rem_vartime(self, reductor).1
+    }
+}
+
 /// Computes `limbs << shift` inplace, where `0 <= shift < Limb::BITS`, returning the carry.
 fn shl_limb_vartime(limbs: &mut [Limb], shift: u32) -> Limb {
     if shift == 0 {

src/uint/div.rs

@@ -953,7 +953,7 @@ impl<const LIMBS: usize> RemLimb for Uint<LIMBS> {
 
 #[cfg(test)]
 mod tests {
-    use crate::{Limb, NonZero, Uint, Word, U128, U256, U64};
+    use crate::{Limb, NonZero, RemMixed, Uint, Word, Zero, U1024, U128, U256, U64, U896};
 
     #[cfg(feature = "rand")]
     use {
@@ -1145,4 +1145,44 @@ mod tests {
         assert_eq!(&a % b, c);
         assert_eq!(&a % &b, c);
     }
+
+    #[test]
+    fn rem_mixed() {
+        let x = U1024::from_be_hex("3740C11DB8F260753BC6B97DD2B8746D3E2694412772AC6ABD975119EE0A6190F27F6F0969BCA069D8D151031AF83EE2283CC2E3E4FADBBDB9EEDBF0B8F4C1FD51912C0D329FDC37D49176DB0A1A2D17E5E6D4F9F6B217FE9412EAA2F881F7027A831C1B06D31D3618D218D6E667DBD85BFC7B6B6B93422D52516989376AA29A");
+        let y = U128::from_u64(1234567890987654321);
+        let rem = x.rem_mixed(&y.to_nz().unwrap());
+
+        let y2: U1024 = U128::concat_mixed(&y, &U896::ZERO);
+        let rem_control = x.rem(&NonZero::new(y2).unwrap());
+
+        assert_eq!(rem.bits(), rem_control.bits());
+        assert_eq!(rem.as_words(), &rem_control.as_words()[0..U128::LIMBS]);
+        assert!(rem_control.as_words()[U128::LIMBS..]
+            .iter()
+            .all(|w| *w == 0));
+    }
+
+    #[test]
+    fn rem_mixed_through_traits() {
+        struct A<T, U> {
+            t: T,
+            u: U,
+        }
+        impl<T, U> A<T, U>
+        where
+            T: RemMixed<U>,
+            U: Clone + Zero,
+        {
+            fn reduce_t_by_u(&self) -> U {
+                let rhs = &NonZero::new(self.u.clone()).unwrap();
+                self.t.rem_mixed(rhs)
+            }
+        }
+
+        let a = A {
+            t: U1024::from(1234567890u64),
+            u: U128::from(456u64),
+        };
+        assert_eq!(a.reduce_t_by_u(), U128::from(330u64));
+    }
 }

src/uint/macros.rs

@@ -97,6 +97,13 @@ macro_rules! impl_uint_concat_split_mixed {
                 self.split_mixed()
             }
         }
+
+        impl $crate::traits::RemMixed<Uint<{ U64::LIMBS * $size }>> for $name
+        {
+            fn rem_mixed(&self, reductor: &NonZero<Uint<{ U64::LIMBS * $size }>>) -> Uint<{ U64::LIMBS * $size }> {
+                self.div_rem_vartime(reductor).1
+            }
+        }
     };
     ($name:ident, [ $($size:literal),+ ]) => {
         $(

src/uint/mul.rs

@@ -94,13 +94,13 @@ pub(crate) const fn schoolbook_squaring(limbs: &[Limb], lo: &mut [Limb], hi: &mu
     let mut carry = Limb::ZERO;
     let mut i = 0;
     while i < limbs.len() {
-        (lo[i].0, carry) = (lo[i].0 << 1 | carry.0, lo[i].shr(Limb::BITS - 1));
+        (lo[i].0, carry) = ((lo[i].0 << 1) | carry.0, lo[i].shr(Limb::BITS - 1));
         i += 1;
     }
 
     let mut i = 0;
     while i < limbs.len() - 1 {
-        (hi[i].0, carry) = (hi[i].0 << 1 | carry.0, hi[i].shr(Limb::BITS - 1));
+        (hi[i].0, carry) = ((hi[i].0 << 1) | carry.0, hi[i].shr(Limb::BITS - 1));
         i += 1;
     }
     hi[limbs.len() - 1] = carry;