malachite_base/unsigned_vector/arithmetic/mod_power_of_2_sum.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, ModPowerOf2Sum};
10use crate::num::basic::unsigneds::PrimitiveUnsigned;
11use crate::unsigned_vector::UnsignedVector;
12use core::borrow::Borrow;
13
14fn assert_reduced<T: PrimitiveUnsigned>(v: &UnsignedVector<T>, pow: u64) {
15 assert!(
16 v.mod_power_of_2_is_reduced(pow),
17 "every vector must be reduced mod 2^pow, but {v} has an element >= 2^{pow}"
18 );
19}
20
21// The first vector fixes the dimension. Since $2^k$ divides $2^W$, where $W$ is the width of `T`,
22// wrapping additions followed by a single reduction of each entry give the sum modulo $2^k$.
23fn mod_power_of_2_sum<T: PrimitiveUnsigned, B: Borrow<UnsignedVector<T>>, I: Iterator<Item = B>>(
24 mut xs: I,
25 pow: u64,
26) -> UnsignedVector<T> {
27 assert!(
28 pow <= T::WIDTH,
29 "pow must be at most T::WIDTH, but it is {pow}"
30 );
31 let first = xs
32 .next()
33 .expect("cannot sum an empty iterator of vectors: the dimension is unknown");
34 let first = first.borrow();
35 assert_reduced(first, pow);
36 let mut sum = first.clone();
37 for v in xs {
38 let v = v.borrow();
39 assert_eq!(
40 v.elements.len(),
41 sum.elements.len(),
42 "cannot sum vectors of different dimensions"
43 );
44 assert_reduced(v, pow);
45 for (s, &x) in sum.elements.iter_mut().zip(&v.elements) {
46 s.wrapping_add_assign(x);
47 }
48 }
49 for s in &mut sum.elements {
50 s.mod_power_of_2_assign(pow);
51 }
52 sum
53}
54
55impl<T: PrimitiveUnsigned> ModPowerOf2Sum for UnsignedVector<T> {
56 /// Adds up all the [`UnsignedVector`]s in an iterator, entry by entry, modulo $2^k$. The
57 /// vectors must all have the same dimension, and their elements must already be reduced modulo
58 /// $2^k$.
59 ///
60 /// $f((v_i)_ {i=0}^{n-1}, k) = w$, where every element of $w$ is less than $2^k$ and $\sum_
61 /// {i=0}^{n-1} v_i \equiv w \mod 2^k$, entry by entry.
62 ///
63 /// # Worst-case complexity
64 /// $T(n) = O(n)$
65 ///
66 /// $M(m) = O(m)$
67 ///
68 /// where $T$ is time, $M$ is additional memory, $n$ is the total number of elements of all the
69 /// vectors, and $m$ is the dimension of the vectors.
70 ///
71 /// # Panics
72 /// Panics if the iterator is empty, if the vectors do not all have the same dimension, if `pow`
73 /// is greater than `T::WIDTH`, or if any element is greater than or equal to $2^k$.
74 ///
75 /// # Examples
76 /// ```
77 /// use core::str::FromStr;
78 /// use malachite_base::num::arithmetic::traits::ModPowerOf2Sum;
79 /// use malachite_base::unsigned_vector::UnsignedVector;
80 ///
81 /// let vs = ["(5, 1, 3)", "(2, 7, 4)", "(6, 6, 6)"]
82 /// .map(|s| UnsignedVector::<u8>::from_str(s).unwrap());
83 /// assert_eq!(
84 /// UnsignedVector::mod_power_of_2_sum(vs.into_iter(), 3).to_string(),
85 /// "(5, 6, 5)"
86 /// );
87 /// ```
88 #[inline]
89 fn mod_power_of_2_sum<I: Iterator<Item = Self>>(xs: I, pow: u64) -> Self {
90 mod_power_of_2_sum(xs, pow)
91 }
92}
93
94impl<'a, T: PrimitiveUnsigned> ModPowerOf2Sum<&'a Self> for UnsignedVector<T> {
95 /// Adds up all the [`UnsignedVector`]s in an iterator of references, entry by entry, modulo
96 /// $2^k$. The vectors must all have the same dimension, and their elements must already be
97 /// reduced modulo $2^k$.
98 ///
99 /// $f((v_i)_ {i=0}^{n-1}, k) = w$, where every element of $w$ is less than $2^k$ and $\sum_
100 /// {i=0}^{n-1} v_i \equiv w \mod 2^k$, entry by entry.
101 ///
102 /// # Worst-case complexity
103 /// $T(n) = O(n)$
104 ///
105 /// $M(m) = O(m)$
106 ///
107 /// where $T$ is time, $M$ is additional memory, $n$ is the total number of elements of all the
108 /// vectors, and $m$ is the dimension of the vectors.
109 ///
110 /// # Panics
111 /// Panics if the iterator is empty, if the vectors do not all have the same dimension, if `pow`
112 /// is greater than `T::WIDTH`, or if any element is greater than or equal to $2^k$.
113 ///
114 /// # Examples
115 /// ```
116 /// use core::str::FromStr;
117 /// use malachite_base::num::arithmetic::traits::ModPowerOf2Sum;
118 /// use malachite_base::unsigned_vector::UnsignedVector;
119 ///
120 /// let vs = ["(5, 1, 3)", "(2, 7, 4)", "(6, 6, 6)"]
121 /// .map(|s| UnsignedVector::<u8>::from_str(s).unwrap());
122 /// assert_eq!(
123 /// UnsignedVector::mod_power_of_2_sum(vs.iter(), 3).to_string(),
124 /// "(5, 6, 5)"
125 /// );
126 /// ```
127 #[inline]
128 fn mod_power_of_2_sum<I: Iterator<Item = &'a Self>>(xs: I, pow: u64) -> Self {
129 mod_power_of_2_sum(xs, pow)
130 }
131}