malachite_base/unsigned_vector/arithmetic/mod_power_of_2_shl.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::arithmetic::traits::{ModPowerOf2IsReduced, ModPowerOf2Shl, ModPowerOf2ShlAssign};
10use crate::num::basic::unsigneds::PrimitiveUnsigned;
11use crate::unsigned_vector::UnsignedVector;
12use alloc::vec;
13
14fn assert_reduced<T: PrimitiveUnsigned>(v: &UnsignedVector<T>, pow: u64) {
15 assert!(
16 pow <= T::WIDTH,
17 "pow must be at most T::WIDTH, but it is {pow}"
18 );
19 assert!(
20 v.mod_power_of_2_is_reduced(pow),
21 "self must be reduced mod 2^pow, but {v} has an element >= 2^{pow}"
22 );
23}
24
25// The shift amount as a `u64`, or `None` if it is at least `pow`, in which case every element
26// becomes zero. `pow` is at most `T::WIDTH`, which is at most 128, so it fits in every unsigned
27// type.
28fn shift_below_pow<U: PrimitiveUnsigned>(bits: U, pow: u64) -> Option<u64> {
29 if bits >= U::exact_from(pow) {
30 None
31 } else {
32 Some(bits.exact_into())
33 }
34}
35
36// Shifts every element left by `bits` modulo 2^pow: only the low `pow - bits` bits of each survive,
37// so the shift cannot overflow.
38fn mod_power_of_2_shl_ref<T: PrimitiveUnsigned, U: PrimitiveUnsigned>(
39 v: &UnsignedVector<T>,
40 bits: U,
41 pow: u64,
42) -> UnsignedVector<T> {
43 assert_reduced(v, pow);
44 let Some(bits) = shift_below_pow(bits, pow) else {
45 return UnsignedVector {
46 elements: vec![T::ZERO; v.elements.len()],
47 };
48 };
49 UnsignedVector {
50 elements: v
51 .elements
52 .iter()
53 .map(|&x| x.mod_power_of_2(pow - bits) << bits)
54 .collect(),
55 }
56}
57
58fn mod_power_of_2_shl_assign<T: PrimitiveUnsigned, U: PrimitiveUnsigned>(
59 v: &mut UnsignedVector<T>,
60 bits: U,
61 pow: u64,
62) {
63 assert_reduced(v, pow);
64 let Some(bits) = shift_below_pow(bits, pow) else {
65 v.elements.fill(T::ZERO);
66 return;
67 };
68 if bits != 0 {
69 for x in &mut v.elements {
70 *x = x.mod_power_of_2(pow - bits) << bits;
71 }
72 }
73}
74
75macro_rules! impl_mod_power_of_2_shl_unsigned {
76 ($t:ident) => {
77 impl<T: PrimitiveUnsigned> ModPowerOf2Shl<$t> for UnsignedVector<T> {
78 type Output = UnsignedVector<T>;
79
80 /// Left-shifts an [`UnsignedVector`] (multiplies it by a power of 2) modulo $2^k$,
81 /// taking the vector by value. The elements must already be reduced modulo $2^k$.
82 ///
83 /// Every element is shifted and reduced, and the dimension is unchanged, so elements
84 /// can become zero; if `bits` is at least `pow`, every element is zero.
85 ///
86 /// $$
87 /// f(v, m, k) = 2^mv \bmod 2^k.
88 /// $$
89 ///
90 /// # Worst-case complexity
91 /// $T(n) = O(n)$
92 ///
93 /// $M(n) = O(1)$
94 ///
95 /// where $T$ is time, $M$ is additional memory, and $n$ is `self.dimension()`.
96 ///
97 /// # Panics
98 /// Panics if `pow` is greater than `T::WIDTH`, or if any element of `self` is greater
99 /// than or equal to $2^k$.
100 ///
101 /// # Examples
102 /// See [here](super::mod_power_of_2_shl#mod_power_of_2_shl).
103 #[inline]
104 fn mod_power_of_2_shl(mut self, bits: $t, pow: u64) -> UnsignedVector<T> {
105 self.mod_power_of_2_shl_assign(bits, pow);
106 self
107 }
108 }
109
110 impl<T: PrimitiveUnsigned> ModPowerOf2Shl<$t> for &UnsignedVector<T> {
111 type Output = UnsignedVector<T>;
112
113 /// Left-shifts an [`UnsignedVector`] (multiplies it by a power of 2) modulo $2^k$,
114 /// taking the vector by reference. The elements must already be reduced modulo $2^k$.
115 ///
116 /// See the documentation for the [`ModPowerOf2Shl`] implementation that takes the
117 /// vector by value for details.
118 ///
119 /// # Worst-case complexity
120 /// $T(n) = O(n)$
121 ///
122 /// $M(n) = O(n)$
123 ///
124 /// where $T$ is time, $M$ is additional memory, and $n$ is `self.dimension()`.
125 ///
126 /// # Panics
127 /// Panics if `pow` is greater than `T::WIDTH`, or if any element of `self` is greater
128 /// than or equal to $2^k$.
129 ///
130 /// # Examples
131 /// See [here](super::mod_power_of_2_shl#mod_power_of_2_shl).
132 #[inline]
133 fn mod_power_of_2_shl(self, bits: $t, pow: u64) -> UnsignedVector<T> {
134 mod_power_of_2_shl_ref(self, bits, pow)
135 }
136 }
137
138 impl<T: PrimitiveUnsigned> ModPowerOf2ShlAssign<$t> for UnsignedVector<T> {
139 /// Left-shifts an [`UnsignedVector`] (multiplies it by a power of 2) modulo $2^k$, in
140 /// place. The elements must already be reduced modulo $2^k$.
141 ///
142 /// See the documentation for the [`ModPowerOf2Shl`] implementation that takes the
143 /// vector by value for details.
144 ///
145 /// # Worst-case complexity
146 /// $T(n) = O(n)$
147 ///
148 /// $M(n) = O(1)$
149 ///
150 /// where $T$ is time, $M$ is additional memory, and $n$ is `self.dimension()`.
151 ///
152 /// # Panics
153 /// Panics if `pow` is greater than `T::WIDTH`, or if any element of `self` is greater
154 /// than or equal to $2^k$.
155 ///
156 /// # Examples
157 /// See [here](super::mod_power_of_2_shl#mod_power_of_2_shl_assign).
158 #[inline]
159 fn mod_power_of_2_shl_assign(&mut self, bits: $t, pow: u64) {
160 mod_power_of_2_shl_assign(self, bits, pow);
161 }
162 }
163 };
164}
165apply_to_unsigneds!(impl_mod_power_of_2_shl_unsigned);