num_bigint/biguint/
shift.rs

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::from(big_digit::BITS).unwrap();
20    let digits = (shift / bits).to_usize().expect("capacity overflow");
21    let shift = (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::from(big_digit::BITS).unwrap();
62    let digits = (shift / bits).to_usize().unwrap_or(usize::MAX);
63    let shift = (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 }
num_bigint/biguint/shift.rs

num_bigint/biguint/
shift.rs
