malachite-base 0.14.0

// Copyright © 2026 Mikhail Hogrefe
//
// This file is part of Malachite.
//
// Malachite is free software: you can redistribute it and/or modify it under the terms of the GNU
// Lesser General Public License (LGPL) as published by the Free Software Foundation; either version
// 3 of the License, or (at your option) any later version. See <https://www.gnu.org/licenses/>.

use crate::num::arithmetic::mod_mul::{limbs_invert_limb_u64, mod_preinverted_double};
use crate::num::arithmetic::traits::{ModDot, ModIsReduced};
use crate::num::basic::integers::PrimitiveInt;
use crate::num::basic::unsigneds::PrimitiveUnsigned;
use crate::num::logic::traits::LeadingZeros;
use crate::unsigned_vector::UnsignedVector;

// Dot products of vectors whose elements are words reduced modulo a word $m$. As in FLINT's
// `nmod_vec` dot products, the sum of products is accumulated exactly, in one, two, or three words,
// depending on how large it can get, and reduced once.

// What summing products of vector elements and reducing the sums modulo $m$ needs: $m$, data for
// reducing two-word values modulo $m$, and the number of words in which the sums are accumulated.
// As in FLINT's `_nmod_vec_dot_bound_limbs`, a sum of at most `terms` products of values less than
// $m$ is bounded by `terms` $(m - 1)^2$, and when the bound fits in one or two words, fewer words
// are used.
pub(crate) struct ModData<T: PrimitiveUnsigned> {
    pub(crate) m: T,
    // When `T` has 64 bits: the inverse of $m$, shifted left until its top bit is set, for
    // `mod_preinverted_double`.
    inv: u64,
    // When `T` has more than 64 bits: $2^\text{W} \bmod m$, where W is `T::WIDTH`.
    radix: T,
    // 1, 2, or 3.
    pub(crate) words: u8,
    // When three words are used: the number of products that can be added to a reduced value
    // without overflowing three words, with room to double the sum: $2^{\text{W} - 1}$, or
    // `usize::MAX` if that does not fit.
    pub(crate) max_terms: usize,
}

impl<T: PrimitiveUnsigned> ModData<T> {
    // Sums will have at most `terms` products, which must be positive.
    pub(crate) fn new(m: T, terms: usize) -> Self {
        assert_ne!(m, T::ZERO);
        let words = match T::try_from(terms) {
            Ok(terms) => {
                // (m - 1)^2 terms, in three words.
                let (hi, lo) = T::x_mul_y_to_zz(m - T::ONE, m - T::ONE);
                let carry = T::x_mul_y_to_zz(lo, terms).0;
                let (top, mid) = T::x_mul_y_to_zz(hi, terms);
                let (top, mid) = T::xx_add_yy_to_zz(top, mid, T::ZERO, carry);
                if top != T::ZERO {
                    3
                } else if mid != T::ZERO {
                    2
                } else {
                    1
                }
            }
            Err(_) => 3,
        };
        let inv = if T::WIDTH == u64::WIDTH {
            let m: u64 = m.wrapping_into();
            limbs_invert_limb_u64(m << LeadingZeros::leading_zeros(m))
        } else {
            0
        };
        Self {
            m,
            inv,
            // 2^W mod m = (2^W - m) mod m, computed without overflow.
            radix: if T::WIDTH > u64::WIDTH {
                m.wrapping_neg() % m
            } else {
                T::ZERO
            },
            words,
            max_terms: if T::WIDTH > usize::WIDTH {
                usize::MAX
            } else {
                1 << (T::WIDTH - 1)
            },
        }
    }

    // The value $x_1 2^\text{W} + x_0$ modulo $m$. Words of at most 32 bits are combined into a
    // `u64` and divided once, and 64-bit words are reduced with a precomputed inverse, as by
    // FLINT's `NMOD2_RED2`.
    #[inline]
    pub(crate) fn reduce_2(&self, x1: T, x0: T) -> T {
        let m = self.m;
        if T::WIDTH <= u32::WIDTH {
            let x: u64 = x1.wrapping_into();
            let x0: u64 = x0.wrapping_into();
            let m: u64 = m.wrapping_into();
            T::wrapping_from(((x << T::WIDTH) | x0) % m)
        } else if T::WIDTH == u64::WIDTH {
            T::wrapping_from(mod_preinverted_double::<u64, u128>(
                x1.wrapping_into(),
                x0.wrapping_into(),
                m.wrapping_into(),
                self.inv,
            ))
        } else {
            let data = T::precompute_mod_mul_data(&m);
            (x1 % m)
                .mod_mul_precomputed(self.radix, m, &data)
                .mod_add(x0 % m, m)
        }
    }

    // The value $x_2 2^{2\text{W}} + x_1 2^\text{W} + x_0$ modulo $m$, by Horner's rule, as by
    // FLINT's `NMOD_RED3`.
    #[inline]
    pub(crate) fn reduce(&self, x2: T, x1: T, x0: T) -> T {
        self.reduce_2(self.reduce_2(x2, x1), x0)
    }

    // A sum from `column_sum`, modulo $m$.
    #[inline]
    pub(crate) fn reduce_sum(&self, (x2, x1, x0): (T, T, T)) -> T {
        match self.words {
            1 => x0 % self.m,
            2 => self.reduce_2(x1, x0),
            _ => self.reduce(x2, x1, x0),
        }
    }
}

// Adds the product of `x` and `y` to the three-word accumulator `(a2, a1, a0)`.
#[inline]
pub(crate) fn accumulate<T: PrimitiveUnsigned>(acc: &mut (T, T, T), x: T, y: T) {
    let (hi, lo) = T::x_mul_y_to_zz(x, y);
    let (a2, a1, a0) = *acc;
    *acc = T::xxx_add_yyy_to_zzz(a2, a1, a0, T::ZERO, hi, lo);
}

// The sum of the products of `xs[i]` and `ys[ys.len() - 1 - i]`, for `i` less than `xs.len()`, as a
// three-word value, accumulated in `d.words` words; `xs` and `ys` must have the same length, at
// most the number of terms `d` was made for. When three words are used, whenever the number of
// products would exceed `d.max_terms` the sum so far is reduced modulo $m$, so the result is
// congruent to the true sum and, with room to spare, can be doubled without overflowing.
#[inline]
pub(crate) fn column_sum<T: PrimitiveUnsigned>(xs: &[T], ys: &[T], d: &ModData<T>) -> (T, T, T) {
    match d.words {
        1 => {
            let mut sum = T::ZERO;
            for (&x, &y) in xs.iter().zip(ys.iter().rev()) {
                sum.wrapping_add_assign(x.wrapping_mul(y));
            }
            (T::ZERO, T::ZERO, sum)
        }
        2 => {
            let (mut hi, mut lo) = (T::ZERO, T::ZERO);
            for (&x, &y) in xs.iter().zip(ys.iter().rev()) {
                let (p_hi, p_lo) = T::x_mul_y_to_zz(x, y);
                (hi, lo) = T::xx_add_yy_to_zz(hi, lo, p_hi, p_lo);
            }
            (T::ZERO, hi, lo)
        }
        _ => {
            let mut acc = (T::ZERO, T::ZERO, T::ZERO);
            for (i, (x_chunk, y_chunk)) in xs
                .chunks(d.max_terms)
                .zip(ys.rchunks(d.max_terms))
                .enumerate()
            {
                if i != 0 {
                    acc = (T::ZERO, T::ZERO, d.reduce(acc.0, acc.1, acc.2));
                }
                for (&x, &y) in x_chunk.iter().zip(y_chunk.iter().rev()) {
                    accumulate(&mut acc, x, y);
                }
            }
            acc
        }
    }
}

// The sum of the products of `xs[i]` and `ys[i]`, as a three-word value: `column_sum` without the
// reversal of `ys`.
#[inline]
pub(crate) fn dot_sum<T: PrimitiveUnsigned>(xs: &[T], ys: &[T], d: &ModData<T>) -> (T, T, T) {
    match d.words {
        1 => {
            let mut sum = T::ZERO;
            for (&x, &y) in xs.iter().zip(ys) {
                sum.wrapping_add_assign(x.wrapping_mul(y));
            }
            (T::ZERO, T::ZERO, sum)
        }
        2 => {
            let (mut hi, mut lo) = (T::ZERO, T::ZERO);
            for (&x, &y) in xs.iter().zip(ys) {
                let (p_hi, p_lo) = T::x_mul_y_to_zz(x, y);
                (hi, lo) = T::xx_add_yy_to_zz(hi, lo, p_hi, p_lo);
            }
            (T::ZERO, hi, lo)
        }
        _ => {
            let mut acc = (T::ZERO, T::ZERO, T::ZERO);
            for (i, (x_chunk, y_chunk)) in xs
                .chunks(d.max_terms)
                .zip(ys.chunks(d.max_terms))
                .enumerate()
            {
                if i != 0 {
                    acc = (T::ZERO, T::ZERO, d.reduce(acc.0, acc.1, acc.2));
                }
                for (&x, &y) in x_chunk.iter().zip(y_chunk) {
                    accumulate(&mut acc, x, y);
                }
            }
            acc
        }
    }
}

// The dot product of `xs` and `ys` modulo `m`, which must be positive; `xs` and `ys` must have the
// same length and be reduced. The products are summed exactly in one, two, or three words and the
// sum is reduced once.
//
// This is equivalent to `_nmod_vec_dot` from `nmod_vec/dot.c`, FLINT 3.6.0.
pub(crate) fn mod_dot<T: PrimitiveUnsigned>(xs: &[T], ys: &[T], m: T) -> T {
    if xs.is_empty() {
        return T::ZERO;
    }
    let d = ModData::new(m, xs.len());
    d.reduce_sum(dot_sum(xs, ys, &d))
}

fn mod_dot_vectors<T: PrimitiveUnsigned>(v: &UnsignedVector<T>, w: &UnsignedVector<T>, m: T) -> T {
    assert_ne!(m, T::ZERO, "m cannot be zero");
    assert_eq!(
        v.elements.len(),
        w.elements.len(),
        "the vectors must have the same dimension"
    );
    assert!(
        v.mod_is_reduced(&m),
        "self must be reduced mod m, but {v} has an element >= {m}"
    );
    assert!(
        w.mod_is_reduced(&m),
        "other must be reduced mod m, but {w} has an element >= {m}"
    );
    mod_dot(&v.elements, &w.elements, m)
}

impl<T: PrimitiveUnsigned> ModDot<Self, T> for UnsignedVector<T> {
    type Output = T;

    /// Computes the dot product of two [`UnsignedVector`]s modulo $m$, taking the first by value
    /// and the second by value. The elements of both must already be reduced.
    ///
    /// The dot product of two 0-dimensional vectors is 0.
    ///
    /// $$
    /// f(v, w, m) = \sum_{i=0}^{n-1} v_iw_i \bmod m,
    /// $$
    ///
    /// where $n$ is the dimension of $v$ and $w$.
    ///
    /// # Worst-case complexity
    /// $T(n) = O(n)$
    ///
    /// $M(n) = O(1)$
    ///
    /// where $T$ is time, $M$ is additional memory, and $n$ is the dimension of the vectors.
    ///
    /// # Panics
    /// Panics if `m` is zero, if the vectors have different dimensions, or if any element is
    /// greater than or equal to `m`.
    ///
    /// # Examples
    /// ```
    /// use core::str::FromStr;
    /// use malachite_base::num::arithmetic::traits::ModDot;
    /// use malachite_base::unsigned_vector::UnsignedVector;
    ///
    /// let v = UnsignedVector::<u8>::from_str("(1, 2, 3)").unwrap();
    /// let w = UnsignedVector::<u8>::from_str("(4, 5, 7)").unwrap();
    /// assert_eq!(v.clone().mod_dot(w.clone(), 10), 5);
    /// ```
    ///
    /// This is equivalent to `_nmod_vec_dot` from `nmod_vec/dot.c`, FLINT 3.6.0.
    #[inline]
    fn mod_dot(self, other: Self, m: T) -> T {
        mod_dot_vectors(&self, &other, m)
    }
}

impl<T: PrimitiveUnsigned> ModDot<&Self, T> for UnsignedVector<T> {
    type Output = T;

    /// Computes the dot product of two [`UnsignedVector`]s modulo $m$, taking the first by value
    /// and the second by reference. The elements of both must already be reduced.
    ///
    /// The dot product of two 0-dimensional vectors is 0.
    ///
    /// $$
    /// f(v, w, m) = \sum_{i=0}^{n-1} v_iw_i \bmod m,
    /// $$
    ///
    /// where $n$ is the dimension of $v$ and $w$.
    ///
    /// # Worst-case complexity
    /// $T(n) = O(n)$
    ///
    /// $M(n) = O(1)$
    ///
    /// where $T$ is time, $M$ is additional memory, and $n$ is the dimension of the vectors.
    ///
    /// # Panics
    /// Panics if `m` is zero, if the vectors have different dimensions, or if any element is
    /// greater than or equal to `m`.
    ///
    /// # Examples
    /// ```
    /// use core::str::FromStr;
    /// use malachite_base::num::arithmetic::traits::ModDot;
    /// use malachite_base::unsigned_vector::UnsignedVector;
    ///
    /// let v = UnsignedVector::<u8>::from_str("(1, 2, 3)").unwrap();
    /// let w = UnsignedVector::<u8>::from_str("(4, 5, 7)").unwrap();
    /// assert_eq!(v.clone().mod_dot(&w, 10), 5);
    /// ```
    ///
    /// This is equivalent to `_nmod_vec_dot` from `nmod_vec/dot.c`, FLINT 3.6.0.
    #[inline]
    fn mod_dot(self, other: &Self, m: T) -> T {
        mod_dot_vectors(&self, other, m)
    }
}

impl<T: PrimitiveUnsigned> ModDot<UnsignedVector<T>, T> for &UnsignedVector<T> {
    type Output = T;

    /// Computes the dot product of two [`UnsignedVector`]s modulo $m$, taking the first by
    /// reference and the second by value. The elements of both must already be reduced.
    ///
    /// The dot product of two 0-dimensional vectors is 0.
    ///
    /// $$
    /// f(v, w, m) = \sum_{i=0}^{n-1} v_iw_i \bmod m,
    /// $$
    ///
    /// where $n$ is the dimension of $v$ and $w$.
    ///
    /// # Worst-case complexity
    /// $T(n) = O(n)$
    ///
    /// $M(n) = O(1)$
    ///
    /// where $T$ is time, $M$ is additional memory, and $n$ is the dimension of the vectors.
    ///
    /// # Panics
    /// Panics if `m` is zero, if the vectors have different dimensions, or if any element is
    /// greater than or equal to `m`.
    ///
    /// # Examples
    /// ```
    /// use core::str::FromStr;
    /// use malachite_base::num::arithmetic::traits::ModDot;
    /// use malachite_base::unsigned_vector::UnsignedVector;
    ///
    /// let v = UnsignedVector::<u8>::from_str("(1, 2, 3)").unwrap();
    /// let w = UnsignedVector::<u8>::from_str("(4, 5, 7)").unwrap();
    /// assert_eq!((&v).mod_dot(w.clone(), 10), 5);
    /// ```
    ///
    /// This is equivalent to `_nmod_vec_dot` from `nmod_vec/dot.c`, FLINT 3.6.0.
    #[inline]
    fn mod_dot(self, other: UnsignedVector<T>, m: T) -> T {
        mod_dot_vectors(self, &other, m)
    }
}

impl<T: PrimitiveUnsigned> ModDot<&UnsignedVector<T>, T> for &UnsignedVector<T> {
    type Output = T;

    /// Computes the dot product of two [`UnsignedVector`]s modulo $m$, taking the first by
    /// reference and the second by reference. The elements of both must already be reduced.
    ///
    /// The dot product of two 0-dimensional vectors is 0.
    ///
    /// $$
    /// f(v, w, m) = \sum_{i=0}^{n-1} v_iw_i \bmod m,
    /// $$
    ///
    /// where $n$ is the dimension of $v$ and $w$.
    ///
    /// # Worst-case complexity
    /// $T(n) = O(n)$
    ///
    /// $M(n) = O(1)$
    ///
    /// where $T$ is time, $M$ is additional memory, and $n$ is the dimension of the vectors.
    ///
    /// # Panics
    /// Panics if `m` is zero, if the vectors have different dimensions, or if any element is
    /// greater than or equal to `m`.
    ///
    /// # Examples
    /// ```
    /// use core::str::FromStr;
    /// use malachite_base::num::arithmetic::traits::ModDot;
    /// use malachite_base::unsigned_vector::UnsignedVector;
    ///
    /// let v = UnsignedVector::<u8>::from_str("(1, 2, 3)").unwrap();
    /// let w = UnsignedVector::<u8>::from_str("(4, 5, 7)").unwrap();
    /// assert_eq!((&v).mod_dot(&w, 10), 5);
    /// ```
    ///
    /// This is equivalent to `_nmod_vec_dot` from `nmod_vec/dot.c`, FLINT 3.6.0.
    #[inline]
    fn mod_dot(self, other: &UnsignedVector<T>, m: T) -> T {
        mod_dot_vectors(self, other, m)
    }
}