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