num_bigint/biguint/
shift.rs

1use super::BigUint;
2
3use crate::big_digit::{self, BigDigits};
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]
12pub(super) fn 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            BigDigits::from_vec(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 { 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        return match n {
70            Cow::Borrowed(_) => BigUint::ZERO,
71            Cow::Owned(mut n) => {
72                n.set_zero();
73                n
74            }
75        };
76    }
77    let mut data = match n {
78        Cow::Borrowed(n) => BigDigits::from_slice(&n.data[digits..]),
79        Cow::Owned(mut n) => {
80            n.data.drain_front(digits);
81            n.data.shrink();
82            n.data
83        }
84    };
85
86    if shift > 0 {
87        let mut borrow = 0;
88        let borrow_shift = big_digit::BITS - shift;
89        for elem in data.iter_mut().rev() {
90            let new_borrow = *elem << borrow_shift;
91            *elem = (*elem >> shift) | borrow;
92            borrow = new_borrow;
93        }
94        // Assuming we were normal before, only one
95        // most-significant digit might be off now.
96        if !data.is_normal() {
97            data.pop();
98        }
99    }
100
101    BigUint { data }
102}
103
104macro_rules! impl_shift {
105    (@ref $Shx:ident :: $shx:ident, $ShxAssign:ident :: $shx_assign:ident, $rhs:ty) => {
106        impl $Shx<&$rhs> for BigUint {
107            type Output = BigUint;
108
109            #[inline]
110            fn $shx(self, rhs: &$rhs) -> BigUint {
111                $Shx::$shx(self, *rhs)
112            }
113        }
114        impl $Shx<&$rhs> for &BigUint {
115            type Output = BigUint;
116
117            #[inline]
118            fn $shx(self, rhs: &$rhs) -> BigUint {
119                $Shx::$shx(self, *rhs)
120            }
121        }
122        impl $ShxAssign<&$rhs> for BigUint {
123            #[inline]
124            fn $shx_assign(&mut self, rhs: &$rhs) {
125                $ShxAssign::$shx_assign(self, *rhs);
126            }
127        }
128    };
129    ($($rhs:ty),+) => {$(
130        impl Shl<$rhs> for BigUint {
131            type Output = BigUint;
132
133            #[inline]
134            fn shl(self, rhs: $rhs) -> BigUint {
135                biguint_shl(Cow::Owned(self), rhs)
136            }
137        }
138        impl Shl<$rhs> for &BigUint {
139            type Output = BigUint;
140
141            #[inline]
142            fn shl(self, rhs: $rhs) -> BigUint {
143                biguint_shl(Cow::Borrowed(self), rhs)
144            }
145        }
146        impl ShlAssign<$rhs> for BigUint {
147            #[inline]
148            fn shl_assign(&mut self, rhs: $rhs) {
149                let n = mem::replace(self, Self::ZERO);
150                *self = n << rhs;
151            }
152        }
153        impl_shift! { @ref Shl::shl, ShlAssign::shl_assign, $rhs }
154
155        impl Shr<$rhs> for BigUint {
156            type Output = BigUint;
157
158            #[inline]
159            fn shr(self, rhs: $rhs) -> BigUint {
160                biguint_shr(Cow::Owned(self), rhs)
161            }
162        }
163        impl Shr<$rhs> for &BigUint {
164            type Output = BigUint;
165
166            #[inline]
167            fn shr(self, rhs: $rhs) -> BigUint {
168                biguint_shr(Cow::Borrowed(self), rhs)
169            }
170        }
171        impl ShrAssign<$rhs> for BigUint {
172            #[inline]
173            fn shr_assign(&mut self, rhs: $rhs) {
174                let n = mem::replace(self, Self::ZERO);
175                *self = n >> rhs;
176            }
177        }
178        impl_shift! { @ref Shr::shr, ShrAssign::shr_assign, $rhs }
179    )*};
180}
181
182impl_shift! { u8, u16, u32, u64, u128, usize }
183impl_shift! { i8, i16, i32, i64, i128, isize }
num_bigint/biguint/shift.rs

num_bigint/biguint/
shift.rs
