Skip to main content

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}