num_bigint/biguint/
shift.rs

Help
1use super::{biguint_from_vec, BigUint};
2
3use crate::big_digit;
4use crate::std_alloc::{Cow, Vec};
5
6use core::mem;
7use core::ops::{Shl, ShlAssign, Shr, ShrAssign};
8use num_traits::{PrimInt, Zero};
9
10#[inline]
11fn biguint_shl<T: PrimInt>(n: Cow<'_, BigUint>, shift: T) -> BigUint {
12    if shift < T::zero() {
13        panic!("attempt to shift left with negative");
14    }
15    if n.is_zero() {
16        return n.into_owned();
17    }
18    let bits = T::from(big_digit::BITS).unwrap();
19    let digits = (shift / bits).to_usize().expect("capacity overflow");
20    let shift = (shift % bits).to_u8().unwrap();
21    biguint_shl2(n, digits, shift)
22}
23
24fn biguint_shl2(n: Cow<'_, BigUint>, digits: usize, shift: u8) -> BigUint {
25    let mut data = match digits {
26        0 => n.into_owned().data,
27        _ => {
28            let len = digits.saturating_add(n.data.len() + 1);
29            let mut data = Vec::with_capacity(len);
30            data.resize(digits, 0);
31            data.extend(n.data.iter());
32            data
33        }
34    };
35
36    if shift > 0 {
37        let mut carry = 0;
38        let carry_shift = big_digit::BITS as u8 - shift;
39        for elem in data[digits..].iter_mut() {
40            let new_carry = *elem >> carry_shift;
41            *elem = (*elem << shift) | carry;
42            carry = new_carry;
43        }
44        if carry != 0 {
45            data.push(carry);
46        }
47    }
48
49    biguint_from_vec(data)
50}
51
52#[inline]
53fn biguint_shr<T: PrimInt>(n: Cow<'_, BigUint>, shift: T) -> BigUint {
54    if shift < T::zero() {
55        panic!("attempt to shift right with negative");
56    }
57    if n.is_zero() {
58        return n.into_owned();
59    }
60    let bits = T::from(big_digit::BITS).unwrap();
61    let digits = (shift / bits).to_usize().unwrap_or(core::usize::MAX);
62    let shift = (shift % bits).to_u8().unwrap();
63    biguint_shr2(n, digits, shift)
64}
65
66fn biguint_shr2(n: Cow<'_, BigUint>, digits: usize, shift: u8) -> BigUint {
67    if digits >= n.data.len() {
68        let mut n = n.into_owned();
69        n.set_zero();
70        return n;
71    }
72    let mut data = match n {
73        Cow::Borrowed(n) => n.data[digits..].to_vec(),
74        Cow::Owned(mut n) => {
75            n.data.drain(..digits);
76            n.data
77        }
78    };
79
80    if shift > 0 {
81        let mut borrow = 0;
82        let borrow_shift = big_digit::BITS as u8 - shift;
83        for elem in data.iter_mut().rev() {
84            let new_borrow = *elem << borrow_shift;
85            *elem = (*elem >> shift) | borrow;
86            borrow = new_borrow;
87        }
88    }
89
90    biguint_from_vec(data)
91}
92
93macro_rules! impl_shift {
94    (@ref $Shx:ident :: $shx:ident, $ShxAssign:ident :: $shx_assign:ident, $rhs:ty) => {
95        impl<'b> $Shx<&'b $rhs> for BigUint {
96            type Output = BigUint;
97
98            #[inline]
99            fn $shx(self, rhs: &'b $rhs) -> BigUint {
100                $Shx::$shx(self, *rhs)
101            }
102        }
103        impl<'a, 'b> $Shx<&'b $rhs> for &'a BigUint {
104            type Output = BigUint;
105
106            #[inline]
107            fn $shx(self, rhs: &'b $rhs) -> BigUint {
108                $Shx::$shx(self, *rhs)
109            }
110        }
111        impl<'b> $ShxAssign<&'b $rhs> for BigUint {
112            #[inline]
113            fn $shx_assign(&mut self, rhs: &'b $rhs) {
114                $ShxAssign::$shx_assign(self, *rhs);
115            }
116        }
117    };
118    ($($rhs:ty),+) => {$(
119        impl Shl<$rhs> for BigUint {
120            type Output = BigUint;
121
122            #[inline]
123            fn shl(self, rhs: $rhs) -> BigUint {
124                biguint_shl(Cow::Owned(self), rhs)
125            }
126        }
127        impl<'a> Shl<$rhs> for &'a BigUint {
128            type Output = BigUint;
129
130            #[inline]
131            fn shl(self, rhs: $rhs) -> BigUint {
132                biguint_shl(Cow::Borrowed(self), rhs)
133            }
134        }
135        impl ShlAssign<$rhs> for BigUint {
136            #[inline]
137            fn shl_assign(&mut self, rhs: $rhs) {
138                let n = mem::replace(self, BigUint::zero());
139                *self = n << rhs;
140            }
141        }
142        impl_shift! { @ref Shl::shl, ShlAssign::shl_assign, $rhs }
143
144        impl Shr<$rhs> for BigUint {
145            type Output = BigUint;
146
147            #[inline]
148            fn shr(self, rhs: $rhs) -> BigUint {
149                biguint_shr(Cow::Owned(self), rhs)
150            }
151        }
152        impl<'a> Shr<$rhs> for &'a BigUint {
153            type Output = BigUint;
154
155            #[inline]
156            fn shr(self, rhs: $rhs) -> BigUint {
157                biguint_shr(Cow::Borrowed(self), rhs)
158            }
159        }
160        impl ShrAssign<$rhs> for BigUint {
161            #[inline]
162            fn shr_assign(&mut self, rhs: $rhs) {
163                let n = mem::replace(self, BigUint::zero());
164                *self = n >> rhs;
165            }
166        }
167        impl_shift! { @ref Shr::shr, ShrAssign::shr_assign, $rhs }
168    )*};
169}
170
171impl_shift! { u8, u16, u32, u64, u128, usize }
172impl_shift! { i8, i16, i32, i64, i128, isize }
num_bigint/biguint/shift.rs

num_bigint/biguint/
shift.rs
