malachite_base/unsigned_vector/arithmetic/
shr.rs

1// Copyright © 2026 Mikhail Hogrefe
2//
3// This file is part of Malachite.
4//
5// Malachite is free software: you can redistribute it and/or modify it under the terms of the GNU
6// Lesser General Public License (LGPL) as published by the Free Software Foundation; either version
7// 3 of the License, or (at your option) any later version. See <https://www.gnu.org/licenses/>.
8
9use crate::num::basic::unsigneds::PrimitiveUnsigned;
10use crate::unsigned_vector::UnsignedVector;
11use alloc::vec;
12use core::ops::{Shr, ShrAssign};
13
14// The shift amount as a `u64`, or `None` if it is at least `T::WIDTH`, in which case every element
15// becomes zero. `T::WIDTH` is at most 128, so it fits in every unsigned type.
16fn shift_below_width<T: PrimitiveUnsigned, U: PrimitiveUnsigned>(bits: U) -> Option<u64> {
17    if bits >= U::exact_from(T::WIDTH) {
18        None
19    } else {
20        Some(bits.exact_into())
21    }
22}
23
24// Shifts every element right by `bits`, taking the floor.
25fn shr_ref<T: PrimitiveUnsigned, U: PrimitiveUnsigned>(
26    v: &UnsignedVector<T>,
27    bits: U,
28) -> UnsignedVector<T> {
29    let Some(bits) = shift_below_width::<T, U>(bits) else {
30        return UnsignedVector {
31            elements: vec![T::ZERO; v.elements.len()],
32        };
33    };
34    UnsignedVector {
35        elements: v.elements.iter().map(|&x| x >> bits).collect(),
36    }
37}
38
39// Shifts every element right by `bits`, taking the floor, in place.
40fn shr_assign<T: PrimitiveUnsigned, U: PrimitiveUnsigned>(v: &mut UnsignedVector<T>, bits: U) {
41    let Some(bits) = shift_below_width::<T, U>(bits) else {
42        v.elements.fill(T::ZERO);
43        return;
44    };
45    if bits != 0 {
46        for x in &mut v.elements {
47            *x >>= bits;
48        }
49    }
50}
51
52macro_rules! impl_shr_unsigned {
53    ($t:ident) => {
54        impl<T: PrimitiveUnsigned> Shr<$t> for UnsignedVector<T> {
55            type Output = UnsignedVector<T>;
56
57            /// Right-shifts an [`UnsignedVector`] (divides it by a power of 2 and takes the floor),
58            /// taking it by value. Every element is shifted, and the dimension is unchanged; if
59            /// `bits` is at least `T::WIDTH`, every element is zero.
60            ///
61            /// $f(v, k) = \lfloor v/2^k \rfloor$. The floor is taken entrywise; to round
62            /// differently, use
63            /// [`entrywise_shr_round`](crate::num::arithmetic::traits::EntrywiseShrRound).
64            ///
65            /// There is no corresponding non-modular left shift, since it could overflow; use
66            /// [`mod_power_of_2_shl`](crate::num::arithmetic::traits::ModPowerOf2Shl) or
67            /// [`mod_shl`](crate::num::arithmetic::traits::ModShl) instead.
68            ///
69            /// # Worst-case complexity
70            /// $T(n) = O(n)$
71            ///
72            /// $M(n) = O(1)$
73            ///
74            /// where $T$ is time, $M$ is additional memory, and $n$ is `self.dimension()`.
75            ///
76            /// # Examples
77            /// See [here](super::shr#shr).
78            #[inline]
79            fn shr(mut self, bits: $t) -> UnsignedVector<T> {
80                self >>= bits;
81                self
82            }
83        }
84
85        impl<T: PrimitiveUnsigned> Shr<$t> for &UnsignedVector<T> {
86            type Output = UnsignedVector<T>;
87
88            /// Right-shifts an [`UnsignedVector`] (divides it by a power of 2 and takes the floor),
89            /// taking it by reference. Every element is shifted, and the dimension is unchanged; if
90            /// `bits` is at least `T::WIDTH`, every element is zero.
91            ///
92            /// $f(v, k) = \lfloor v/2^k \rfloor$. The floor is taken entrywise; to round
93            /// differently, use
94            /// [`entrywise_shr_round`](crate::num::arithmetic::traits::EntrywiseShrRound).
95            ///
96            /// # Worst-case complexity
97            /// $T(n) = O(n)$
98            ///
99            /// $M(n) = O(n)$
100            ///
101            /// where $T$ is time, $M$ is additional memory, and $n$ is `self.dimension()`.
102            ///
103            /// # Examples
104            /// See [here](super::shr#shr).
105            #[inline]
106            fn shr(self, bits: $t) -> UnsignedVector<T> {
107                shr_ref(self, bits)
108            }
109        }
110
111        impl<T: PrimitiveUnsigned> ShrAssign<$t> for UnsignedVector<T> {
112            /// Right-shifts an [`UnsignedVector`] (divides it by a power of 2 and takes the floor),
113            /// in place. Every element is shifted, and the dimension is unchanged; if `bits` is at
114            /// least `T::WIDTH`, every element becomes zero.
115            ///
116            /// $v \gets \lfloor v/2^k \rfloor$. The floor is taken entrywise; to round differently,
117            /// use [`entrywise_shr_round`](crate::num::arithmetic::traits::EntrywiseShrRound).
118            ///
119            /// # Worst-case complexity
120            /// $T(n) = O(n)$
121            ///
122            /// $M(n) = O(1)$
123            ///
124            /// where $T$ is time, $M$ is additional memory, and $n$ is `self.dimension()`.
125            ///
126            /// # Examples
127            /// See [here](super::shr#shr_assign).
128            #[inline]
129            fn shr_assign(&mut self, bits: $t) {
130                shr_assign(self, bits);
131            }
132        }
133    };
134}
135apply_to_unsigneds!(impl_shr_unsigned);
malachite_base/unsigned_vector/arithmetic/shr.rs

malachite_base/unsigned_vector/arithmetic/
shr.rs
