dashu_int/
shift_ops.rs

1//! Bit shift operators.
2
3use crate::{ibig::IBig, ubig::UBig, Sign::*};
4use core::{
5    mem,
6    ops::{Shl, ShlAssign, Shr, ShrAssign},
7};
8
9macro_rules! impl_shifts {
10    ($t:ty) => {
11        impl Shl<&usize> for $t {
12            type Output = $t;
13
14            #[inline]
15            fn shl(self, rhs: &usize) -> $t {
16                self.shl(*rhs)
17            }
18        }
19
20        impl Shl<&usize> for &$t {
21            type Output = $t;
22
23            #[inline]
24            fn shl(self, rhs: &usize) -> $t {
25                self.shl(*rhs)
26            }
27        }
28
29        impl ShlAssign<usize> for $t {
30            #[inline]
31            fn shl_assign(&mut self, rhs: usize) {
32                *self = mem::take(self) << rhs;
33            }
34        }
35
36        impl ShlAssign<&usize> for $t {
37            #[inline]
38            fn shl_assign(&mut self, rhs: &usize) {
39                *self = mem::take(self) << rhs;
40            }
41        }
42
43        impl Shr<&usize> for $t {
44            type Output = $t;
45
46            #[inline]
47            fn shr(self, rhs: &usize) -> $t {
48                self.shr(*rhs)
49            }
50        }
51
52        impl Shr<&usize> for &$t {
53            type Output = $t;
54
55            #[inline]
56            fn shr(self, rhs: &usize) -> $t {
57                self.shr(*rhs)
58            }
59        }
60
61        impl ShrAssign<usize> for $t {
62            #[inline]
63            fn shr_assign(&mut self, rhs: usize) {
64                *self = mem::take(self).shr(rhs);
65            }
66        }
67
68        impl ShrAssign<&usize> for $t {
69            #[inline]
70            fn shr_assign(&mut self, rhs: &usize) {
71                *self = mem::take(self).shr(rhs);
72            }
73        }
74    };
75}
76
77impl_shifts!(UBig);
78impl_shifts!(IBig);
79
80impl Shl<usize> for UBig {
81    type Output = UBig;
82
83    #[inline]
84    fn shl(self, rhs: usize) -> UBig {
85        UBig(self.into_repr().shl(rhs))
86    }
87}
88
89impl Shl<usize> for &UBig {
90    type Output = UBig;
91
92    #[inline]
93    fn shl(self, rhs: usize) -> UBig {
94        UBig(self.repr().shl(rhs))
95    }
96}
97
98impl Shr<usize> for UBig {
99    type Output = UBig;
100
101    #[inline]
102    fn shr(self, rhs: usize) -> UBig {
103        UBig(self.into_repr().shr(rhs))
104    }
105}
106
107impl Shr<usize> for &UBig {
108    type Output = UBig;
109
110    #[inline]
111    fn shr(self, rhs: usize) -> UBig {
112        UBig(self.repr().shr(rhs))
113    }
114}
115
116impl Shl<usize> for IBig {
117    type Output = IBig;
118
119    #[inline]
120    fn shl(self, rhs: usize) -> IBig {
121        let (sign, mag) = self.into_sign_repr();
122        let repr = mag << rhs;
123        IBig(repr.with_sign(sign))
124    }
125}
126
127impl Shl<usize> for &IBig {
128    type Output = IBig;
129
130    #[inline]
131    fn shl(self, rhs: usize) -> IBig {
132        let (sign, mag) = self.as_sign_repr();
133        let repr = mag << rhs;
134        IBig(repr.with_sign(sign))
135    }
136}
137
138impl Shr<usize> for IBig {
139    type Output = IBig;
140
141    #[inline]
142    fn shr(self, rhs: usize) -> IBig {
143        let (sign, mag) = self.into_sign_repr();
144        match sign {
145            Positive => IBig(mag >> rhs),
146            Negative => {
147                let b = mag.as_ref().are_low_bits_nonzero(rhs);
148                -IBig(mag >> rhs) - IBig::from(b)
149            }
150        }
151    }
152}
153
154impl Shr<usize> for &IBig {
155    type Output = IBig;
156
157    #[inline]
158    fn shr(self, rhs: usize) -> IBig {
159        let (sign, mag) = self.as_sign_repr();
160        match sign {
161            Positive => IBig(mag >> rhs),
162            Negative => {
163                let b = mag.are_low_bits_nonzero(rhs);
164                -IBig(mag >> rhs) - IBig::from(b)
165            }
166        }
167    }
168}
169
170pub(crate) mod repr {
171    use super::*;
172    use crate::{
173        arch::word::{DoubleWord, Word},
174        buffer::Buffer,
175        math,
176        primitive::{double_word, DWORD_BITS_USIZE, WORD_BITS_USIZE},
177        repr::{
178            Repr,
179            TypedRepr::{self, *},
180            TypedReprRef::{self, *},
181        },
182        shift,
183    };
184
185    impl Shl<usize> for TypedRepr {
186        type Output = Repr;
187        #[inline]
188        fn shl(self, rhs: usize) -> Repr {
189            match self {
190                Small(0) => Repr::zero(),
191                Small(dword) => shl_dword(dword, rhs),
192                Large(buffer) => shl_large(buffer, rhs),
193            }
194        }
195    }
196
197    impl Shr<usize> for TypedRepr {
198        type Output = Repr;
199        #[inline]
200        fn shr(self, rhs: usize) -> Repr {
201            match self {
202                Small(dword) => shr_dword(dword, rhs),
203                Large(buffer) => shr_large(buffer, rhs),
204            }
205        }
206    }
207
208    impl<'a> Shl<usize> for TypedReprRef<'a> {
209        type Output = Repr;
210        #[inline]
211        fn shl(self, rhs: usize) -> Repr {
212            match self {
213                RefSmall(0) => Repr::zero(),
214                RefSmall(dword) => shl_dword(dword, rhs),
215                RefLarge(words) => shl_large_ref(words, rhs),
216            }
217        }
218    }
219
220    impl<'a> Shr<usize> for TypedReprRef<'a> {
221        type Output = Repr;
222        #[inline]
223        fn shr(self, rhs: usize) -> Repr {
224            match self {
225                RefSmall(dword) => shr_dword(dword, rhs),
226                RefLarge(words) => shr_large_ref(words, rhs),
227            }
228        }
229    }
230
231    /// Shift left a non-zero `DoubleWord` by `rhs` bits.
232    #[inline]
233    fn shl_dword(dword: DoubleWord, rhs: usize) -> Repr {
234        debug_assert!(dword != 0);
235
236        if rhs <= dword.leading_zeros() as usize {
237            Repr::from_dword(dword << rhs)
238        } else if dword == 1 {
239            shl_one_spilled(rhs)
240        } else {
241            shl_dword_spilled(dword, rhs)
242        }
243    }
244
245    /// Shift left 1 by `rhs` bits
246    fn shl_one_spilled(rhs: usize) -> Repr {
247        debug_assert!(rhs >= DWORD_BITS_USIZE);
248        let idx = rhs / WORD_BITS_USIZE;
249        let mut buffer = Buffer::allocate(idx + 1);
250        buffer.push_zeros(idx);
251        buffer.push(1 << (rhs % WORD_BITS_USIZE));
252        Repr::from_buffer(buffer)
253    }
254
255    /// Shift left a non-zero `DoubleWord` by `rhs` bits.
256    fn shl_dword_spilled(dword: DoubleWord, rhs: usize) -> Repr {
257        let shift_words = rhs / WORD_BITS_USIZE;
258        let shift_bits = (rhs % WORD_BITS_USIZE) as u32;
259
260        let (n0, n1, n2) = math::shl_dword(dword, shift_bits);
261        let mut buffer = Buffer::allocate(shift_words + 3);
262        buffer.push_zeros(shift_words);
263        buffer.push(n0);
264        buffer.push(n1);
265        buffer.push(n2);
266        Repr::from_buffer(buffer)
267    }
268
269    /// Shift left `buffer` by `rhs` bits.
270    fn shl_large(mut buffer: Buffer, rhs: usize) -> Repr {
271        let shift_words = rhs / WORD_BITS_USIZE;
272
273        if buffer.capacity() < buffer.len() + shift_words + 1 {
274            return shl_large_ref(&buffer, rhs);
275        }
276
277        let shift_bits = (rhs % WORD_BITS_USIZE) as u32;
278        let carry = shift::shl_in_place(&mut buffer, shift_bits);
279        buffer.push(carry);
280        buffer.push_zeros_front(shift_words);
281        Repr::from_buffer(buffer)
282    }
283
284    /// Shift left large number of words by `rhs` bits.
285    pub(crate) fn shl_large_ref(words: &[Word], rhs: usize) -> Repr {
286        let shift_words = rhs / WORD_BITS_USIZE;
287        let shift_bits = (rhs % WORD_BITS_USIZE) as u32;
288
289        let mut buffer = Buffer::allocate(shift_words + words.len() + 1);
290        buffer.push_zeros(shift_words);
291        buffer.push_slice(words);
292        let carry = shift::shl_in_place(&mut buffer[shift_words..], shift_bits);
293        buffer.push(carry);
294        Repr::from_buffer(buffer)
295    }
296
297    /// Shift right one `DoubleWord` by `rhs` bits.
298    #[inline]
299    fn shr_dword(dword: DoubleWord, rhs: usize) -> Repr {
300        if rhs < DWORD_BITS_USIZE {
301            Repr::from_dword(dword >> rhs)
302        } else {
303            Repr::zero()
304        }
305    }
306
307    /// Shift right `buffer` by `rhs` bits.
308    fn shr_large(mut buffer: Buffer, rhs: usize) -> Repr {
309        let shift_words = rhs / WORD_BITS_USIZE;
310        if shift_words >= buffer.len() {
311            return Repr::zero();
312        }
313        let shift_bits = (rhs % WORD_BITS_USIZE) as u32;
314        buffer.erase_front(shift_words);
315        shift::shr_in_place(&mut buffer, shift_bits);
316        Repr::from_buffer(buffer)
317    }
318
319    /// Shift right large number of words by `rhs` bits.
320    pub(crate) fn shr_large_ref(words: &[Word], rhs: usize) -> Repr {
321        let shift_words = rhs / WORD_BITS_USIZE;
322        let shift_bits = (rhs % WORD_BITS_USIZE) as u32;
323
324        let words = &words[shift_words.min(words.len())..];
325
326        match words {
327            [] => Repr::zero(),
328            &[w] => Repr::from_word(w >> shift_bits),
329            &[lo, hi] => Repr::from_dword(double_word(lo, hi) >> shift_bits),
330            _ => {
331                let mut buffer = Buffer::allocate(words.len());
332                buffer.push_slice(words);
333                shift::shr_in_place(&mut buffer, shift_bits);
334                Repr::from_buffer(buffer)
335            }
336        }
337    }
338}