Merge #142

142: Implement shift with broad RHS types r=cuviper a=cuviper

The primitive integers all support shift operators with the right-hand side as any other primitive integer, by value or by reference. Now `BigInt` and `BigUint` support this too. This is a minor breaking change, because it can cause type inference failures if a user has an ambiguous integer for a shift count, whereas before it could only be `usize`.

Co-authored-by: Josh Stone <[email protected]>

Author: bors[bot]

benches/bigint.rs

@@ -293,9 +293,10 @@ fn rand_131072(b: &mut Bencher) {
 
 #[bench]
 fn shl(b: &mut Bencher) {
-    let n = BigUint::one() << 1000;
+    let n = BigUint::one() << 1000u32;
+    let mut m = n.clone();
     b.iter(|| {
-        let mut m = n.clone();
+        m.clone_from(&n);
         for i in 0..50 {
             m <<= i;
         }
@@ -304,9 +305,10 @@ fn shl(b: &mut Bencher) {
 
 #[bench]
 fn shr(b: &mut Bencher) {
-    let n = BigUint::one() << 2000;
+    let n = BigUint::one() << 2000u32;
+    let mut m = n.clone();
     b.iter(|| {
-        let mut m = n.clone();
+        m.clone_from(&n);
         for i in 0..50 {
             m >>= i;
         }

src/algorithms.rs

@@ -3,7 +3,7 @@ use core::cmp;
 use core::cmp::Ordering::{self, Equal, Greater, Less};
 use core::iter::repeat;
 use core::mem;
-use num_traits::{One, Zero};
+use num_traits::{One, PrimInt, Zero};
 
 use crate::biguint::biguint_from_vec;
 use crate::biguint::BigUint;
@@ -720,34 +720,46 @@ fn div_rem_core(mut a: BigUint, b: &BigUint) -> (BigUint, BigUint) {
 
 /// Find last set bit
 /// fls(0) == 0, fls(u32::MAX) == 32
-pub(crate) fn fls<T: num_traits::PrimInt>(v: T) -> usize {
+pub(crate) fn fls<T: PrimInt>(v: T) -> usize {
     mem::size_of::<T>() * 8 - v.leading_zeros() as usize
 }
 
-pub(crate) fn ilog2<T: num_traits::PrimInt>(v: T) -> usize {
+pub(crate) fn ilog2<T: PrimInt>(v: T) -> usize {
     fls(v) - 1
 }
 
 #[inline]
-pub(crate) fn biguint_shl(n: Cow<'_, BigUint>, bits: usize) -> BigUint {
-    let n_unit = bits / big_digit::BITS;
-    let mut data = match n_unit {
+pub(crate) fn biguint_shl<T: PrimInt>(n: Cow<'_, BigUint>, shift: T) -> BigUint {
+    if shift < T::zero() {
+        panic!("attempt to shift left with negative");
+    }
+    if n.is_zero() {
+        return n.into_owned();
+    }
+    let bits = T::from(big_digit::BITS).unwrap();
+    let digits = (shift / bits).to_usize().expect("capacity overflow");
+    let shift = (shift % bits).to_u8().unwrap();
+    biguint_shl2(n, digits, shift)
+}
+
+fn biguint_shl2(n: Cow<'_, BigUint>, digits: usize, shift: u8) -> BigUint {
+    let mut data = match digits {
         0 => n.into_owned().data,
         _ => {
-            let len = n_unit + n.data.len() + 1;
+            let len = digits.saturating_add(n.data.len() + 1);
             let mut data = Vec::with_capacity(len);
-            data.extend(repeat(0).take(n_unit));
-            data.extend(n.data.iter().cloned());
+            data.extend(repeat(0).take(digits));
+            data.extend(n.data.iter());
             data
         }
     };
 
-    let n_bits = bits % big_digit::BITS;
-    if n_bits > 0 {
+    if shift > 0 {
         let mut carry = 0;
-        for elem in data[n_unit..].iter_mut() {
-            let new_carry = *elem >> (big_digit::BITS - n_bits);
-            *elem = (*elem << n_bits) | carry;
+        let carry_shift = big_digit::BITS as u8 - shift;
+        for elem in data[digits..].iter_mut() {
+            let new_carry = *elem >> carry_shift;
+            *elem = (*elem << shift) | carry;
             carry = new_carry;
         }
         if carry != 0 {
@@ -759,25 +771,39 @@ pub(crate) fn biguint_shl(n: Cow<'_, BigUint>, bits: usize) -> BigUint {
 }
 
 #[inline]
-pub(crate) fn biguint_shr(n: Cow<'_, BigUint>, bits: usize) -> BigUint {
-    let n_unit = bits / big_digit::BITS;
-    if n_unit >= n.data.len() {
-        return Zero::zero();
+pub(crate) fn biguint_shr<T: PrimInt>(n: Cow<'_, BigUint>, shift: T) -> BigUint {
+    if shift < T::zero() {
+        panic!("attempt to shift right with negative");
+    }
+    if n.is_zero() {
+        return n.into_owned();
+    }
+    let bits = T::from(big_digit::BITS).unwrap();
+    let digits = (shift / bits).to_usize().unwrap_or(core::usize::MAX);
+    let shift = (shift % bits).to_u8().unwrap();
+    biguint_shr2(n, digits, shift)
+}
+
+fn biguint_shr2(n: Cow<'_, BigUint>, digits: usize, shift: u8) -> BigUint {
+    if digits >= n.data.len() {
+        let mut n = n.into_owned();
+        n.set_zero();
+        return n;
     }
     let mut data = match n {
-        Cow::Borrowed(n) => n.data[n_unit..].to_vec(),
+        Cow::Borrowed(n) => n.data[digits..].to_vec(),
         Cow::Owned(mut n) => {
-            n.data.drain(..n_unit);
+            n.data.drain(..digits);
             n.data
         }
     };
 
-    let n_bits = bits % big_digit::BITS;
-    if n_bits > 0 {
+    if shift > 0 {
         let mut borrow = 0;
+        let borrow_shift = big_digit::BITS as u8 - shift;
         for elem in data.iter_mut().rev() {
-            let new_borrow = *elem << (big_digit::BITS - n_bits);
-            *elem = (*elem >> n_bits) | borrow;
+            let new_borrow = *elem << borrow_shift;
+            *elem = (*elem >> shift) | borrow;
             borrow = new_borrow;
         }
     }

src/bigint.rs

@@ -25,7 +25,7 @@ use serde;
 
 use num_integer::{Integer, Roots};
 use num_traits::{
-    CheckedAdd, CheckedDiv, CheckedMul, CheckedSub, FromPrimitive, Num, One, Pow, Signed,
+    CheckedAdd, CheckedDiv, CheckedMul, CheckedSub, FromPrimitive, Num, One, Pow, PrimInt, Signed,
     ToPrimitive, Zero,
 };
 
@@ -779,69 +779,105 @@ impl Num for BigInt {
     }
 }
 
-impl Shl<usize> for BigInt {
-    type Output = BigInt;
+macro_rules! impl_shift {
+    (@ref $Shx:ident :: $shx:ident, $ShxAssign:ident :: $shx_assign:ident, $rhs:ty) => {
+        impl<'b> $Shx<&'b $rhs> for BigInt {
+            type Output = BigInt;
 
-    #[inline]
-    fn shl(mut self, rhs: usize) -> BigInt {
-        self <<= rhs;
-        self
-    }
-}
+            #[inline]
+            fn $shx(self, rhs: &'b $rhs) -> BigInt {
+                $Shx::$shx(self, *rhs)
+            }
+        }
+        impl<'a, 'b> $Shx<&'b $rhs> for &'a BigInt {
+            type Output = BigInt;
 
-impl<'a> Shl<usize> for &'a BigInt {
-    type Output = BigInt;
+            #[inline]
+            fn $shx(self, rhs: &'b $rhs) -> BigInt {
+                $Shx::$shx(self, *rhs)
+            }
+        }
+        impl<'b> $ShxAssign<&'b $rhs> for BigInt {
+            #[inline]
+            fn $shx_assign(&mut self, rhs: &'b $rhs) {
+                $ShxAssign::$shx_assign(self, *rhs);
+            }
+        }
+    };
+    ($($rhs:ty),+) => {$(
+        impl Shl<$rhs> for BigInt {
+            type Output = BigInt;
 
-    #[inline]
-    fn shl(self, rhs: usize) -> BigInt {
-        BigInt::from_biguint(self.sign, &self.data << rhs)
-    }
-}
+            #[inline]
+            fn shl(self, rhs: $rhs) -> BigInt {
+                BigInt::from_biguint(self.sign, self.data << rhs)
+            }
+        }
+        impl<'a> Shl<$rhs> for &'a BigInt {
+            type Output = BigInt;
 
-impl ShlAssign<usize> for BigInt {
-    #[inline]
-    fn shl_assign(&mut self, rhs: usize) {
-        self.data <<= rhs;
-    }
-}
+            #[inline]
+            fn shl(self, rhs: $rhs) -> BigInt {
+                BigInt::from_biguint(self.sign, &self.data << rhs)
+            }
+        }
+        impl ShlAssign<$rhs> for BigInt {
+            #[inline]
+            fn shl_assign(&mut self, rhs: $rhs) {
+                self.data <<= rhs
+            }
+        }
+        impl_shift! { @ref Shl::shl, ShlAssign::shl_assign, $rhs }
 
-// Negative values need a rounding adjustment if there are any ones in the
-// bits that are getting shifted out.
-fn shr_round_down(i: &BigInt, rhs: usize) -> bool {
-    i.is_negative() && i.trailing_zeros().map(|n| n < rhs).unwrap_or(false)
-}
+        impl Shr<$rhs> for BigInt {
+            type Output = BigInt;
 
-impl Shr<usize> for BigInt {
-    type Output = BigInt;
+            #[inline]
+            fn shr(self, rhs: $rhs) -> BigInt {
+                let round_down = shr_round_down(&self, rhs);
+                let data = self.data >> rhs;
+                let data = if round_down { data + 1u8 } else { data };
+                BigInt::from_biguint(self.sign, data)
+            }
+        }
+        impl<'a> Shr<$rhs> for &'a BigInt {
+            type Output = BigInt;
 
-    #[inline]
-    fn shr(mut self, rhs: usize) -> BigInt {
-        self >>= rhs;
-        self
-    }
+            #[inline]
+            fn shr(self, rhs: $rhs) -> BigInt {
+                let round_down = shr_round_down(self, rhs);
+                let data = &self.data >> rhs;
+                let data = if round_down { data + 1u8 } else { data };
+                BigInt::from_biguint(self.sign, data)
+            }
+        }
+        impl ShrAssign<$rhs> for BigInt {
+            #[inline]
+            fn shr_assign(&mut self, rhs: $rhs) {
+                let round_down = shr_round_down(self, rhs);
+                self.data >>= rhs;
+                if round_down {
+                    self.data += 1u8;
+                } else if self.data.is_zero() {
+                    self.sign = NoSign;
+                }
+            }
+        }
+        impl_shift! { @ref Shr::shr, ShrAssign::shr_assign, $rhs }
+    )*};
 }
 
-impl<'a> Shr<usize> for &'a BigInt {
-    type Output = BigInt;
-
-    #[inline]
-    fn shr(self, rhs: usize) -> BigInt {
-        let round_down = shr_round_down(self, rhs);
-        let data = &self.data >> rhs;
-        BigInt::from_biguint(self.sign, if round_down { data + 1u8 } else { data })
-    }
-}
+impl_shift! { u8, u16, u32, u64, u128, usize }
+impl_shift! { i8, i16, i32, i64, i128, isize }
 
-impl ShrAssign<usize> for BigInt {
-    #[inline]
-    fn shr_assign(&mut self, rhs: usize) {
-        let round_down = shr_round_down(self, rhs);
-        self.data >>= rhs;
-        if round_down {
-            self.data += 1u8;
-        } else if self.data.is_zero() {
-            self.sign = NoSign;
-        }
+// Negative values need a rounding adjustment if there are any ones in the
+// bits that are getting shifted out.
+fn shr_round_down<T: PrimInt>(i: &BigInt, shift: T) -> bool {
+    if i.is_negative() {
+        let zeros = i.trailing_zeros().expect("negative values are non-zero");
+        shift > T::zero() && shift.to_usize().map(|shift| zeros < shift).unwrap_or(true)
+    } else {
+        false
     }
 }