1 use super::{biguint_from_vec, BigUint};
2 
3 use crate::big_digit;
4 use crate::std_alloc::{Cow, Vec};
5 
6 use core::mem;
7 use core::ops::{Shl, ShlAssign, Shr, ShrAssign};
8 use num_traits::{PrimInt, Zero};
9 
10 #[inline]
biguint_shl<T: PrimInt>(n: Cow<'_, BigUint>, shift: T) -> BigUint11 fn 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 
biguint_shl2(n: Cow<'_, BigUint>, digits: usize, shift: u8) -> BigUint24 fn 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 - 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]
biguint_shr<T: PrimInt>(n: Cow<'_, BigUint>, shift: T) -> BigUint53 fn 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 
biguint_shr2(n: Cow<'_, BigUint>, digits: usize, shift: u8) -> BigUint66 fn 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 - 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 
93 macro_rules! impl_shift {
94     (@ref $Shx:ident :: $shx:ident, $ShxAssign:ident :: $shx_assign:ident, $rhs:ty) => {
95         impl $Shx<&$rhs> for BigUint {
96             type Output = BigUint;
97 
98             #[inline]
99             fn $shx(self, rhs: &$rhs) -> BigUint {
100                 $Shx::$shx(self, *rhs)
101             }
102         }
103         impl $Shx<&$rhs> for &BigUint {
104             type Output = BigUint;
105 
106             #[inline]
107             fn $shx(self, rhs: &$rhs) -> BigUint {
108                 $Shx::$shx(self, *rhs)
109             }
110         }
111         impl $ShxAssign<&$rhs> for BigUint {
112             #[inline]
113             fn $shx_assign(&mut self, rhs: &$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 Shl<$rhs> for &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 Shr<$rhs> for &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 
171 impl_shift! { u8, u16, u32, u64, u128, usize }
172 impl_shift! { i8, i16, i32, i64, i128, isize }
173