Skip to content

Commit

Permalink
Merge #142
Browse files Browse the repository at this point in the history 
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]>
  • Loading branch information
@bors @cuviper
bors[bot] and cuviper authored Apr 4, 2020
2 parents c076bb8 + e827fd5 commit 4239ddd
Show file tree
Hide file tree
Showing 7 changed files with 244 additions and 149 deletions.
10 changes: 6 additions & 4 deletions benches/bigint.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand All @@ -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;
}
Expand Down
74 changes: 50 additions & 24 deletions src/algorithms.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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 {
Expand All @@ -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;
}
}
Expand Down
144 changes: 90 additions & 54 deletions src/bigint.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};

Expand Down Expand Up @@ -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
}
}

Expand Down
Loading

0 comments on commit 4239ddd

Please sign in to comment.