1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107

use super::BigInt;
use super::Sign::NoSign;

use core::ops::{Shl, ShlAssign, Shr, ShrAssign};
use num_traits::{PrimInt, Signed, Zero};

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 $shx(self, rhs: &'b $rhs) -> BigInt {
                $Shx::$shx(self, *rhs)
            }
        }
        impl<'a, 'b> $Shx<&'b $rhs> 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: $rhs) -> BigInt {
                BigInt::from_biguint(self.sign, self.data << rhs)
            }
        }
        impl<'a> Shl<$rhs> for &'a BigInt {
            type Output = BigInt;

            #[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 }

        impl Shr<$rhs> 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(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_shift! { u8, u16, u32, u64, u128, usize }
impl_shift! { i8, i16, i32, i64, i128, isize }

// 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_u64().map(|shift| zeros < shift).unwrap_or(true)
    } else {
        false
    }
}