p3-matrix 0.8.0

#![doc = include_str!("../README.md")]
#![no_std]

extern crate alloc;

use alloc::vec::Vec;
use core::fmt::{Debug, Display, Formatter};
use core::ops::Deref;

use itertools::Itertools;
use p3_field::{
    BasedVectorSpace, ExtensionField, Field, FieldArray, PackedField, PackedFieldExtension,
    PackedValue, PrimeCharacteristicRing,
};
use p3_maybe_rayon::PARALLEL_ENABLED;
use p3_maybe_rayon::prelude::*;
use strided::{VerticallyStridedMatrixView, VerticallyStridedRowIndexMap};
use tracing::instrument;

use crate::dense::RowMajorMatrix;

pub mod bitrev;
pub mod dense;
pub mod extension;
pub mod horizontally_truncated;
pub mod interpolation;
pub mod row_index_mapped;
pub mod stack;
pub mod strided;
pub mod util;

/// A simple struct representing the shape of a matrix.
///
/// The `Dimensions` type stores the number of columns (`width`) and rows (`height`)
/// of a matrix. It is commonly used for querying and displaying matrix shapes.
#[derive(Copy, Clone, PartialEq, Eq)]
pub struct Dimensions {
    /// Number of columns in the matrix.
    pub width: usize,
    /// Number of rows in the matrix.
    pub height: usize,
}

impl Debug for Dimensions {
    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
        write!(f, "{}x{}", self.width, self.height)
    }
}

impl Display for Dimensions {
    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
        write!(f, "{}x{}", self.width, self.height)
    }
}

/// A generic trait for two-dimensional matrix-like data structures.
///
/// The `Matrix` trait provides a uniform interface for accessing rows, elements,
/// and computing with matrices in both sequential and parallel contexts. It supports
/// packing strategies for SIMD optimizations and interaction with extension fields.
pub trait Matrix<T: Send + Sync + Clone>: Send + Sync {
    /// Returns the number of columns in the matrix.
    fn width(&self) -> usize;

    /// Returns the number of rows in the matrix.
    fn height(&self) -> usize;

    /// Returns the dimensions (width, height) of the matrix.
    fn dimensions(&self) -> Dimensions {
        Dimensions {
            width: self.width(),
            height: self.height(),
        }
    }

    // The methods:
    // get, get_unchecked, row, row_unchecked, row_subseq_unchecked, row_slice, row_slice_unchecked, row_subslice_unchecked
    // are all defined in a circular manner so you only need to implement a subset of them.
    // In particular is is enough to implement just one of: row_unchecked, row_subseq_unchecked
    //
    // That being said, most implementations will want to implement several methods for performance reasons.

    /// Returns the element at the given row and column.
    ///
    /// Returns `None` if either `r >= height()` or `c >= width()`.
    #[inline]
    fn get(&self, r: usize, c: usize) -> Option<T> {
        (r < self.height() && c < self.width()).then(|| unsafe {
            // Safety: Clearly `r < self.height()` and `c < self.width()`.
            self.get_unchecked(r, c)
        })
    }

    /// Returns the element at the given row and column.
    ///
    /// For a safe alternative, see [`Self::get`].
    ///
    /// # Safety
    /// The caller must ensure that `r < self.height()` and `c < self.width()`.
    /// Breaking any of these assumptions is considered undefined behaviour.
    #[inline]
    unsafe fn get_unchecked(&self, r: usize, c: usize) -> T {
        unsafe { self.row_slice_unchecked(r)[c].clone() }
    }

    /// Returns an iterator over the elements of the `r`-th row.
    ///
    /// The iterator will have `self.width()` elements.
    ///
    /// Returns `None` if `r >= height()`.
    #[inline]
    fn row(
        &self,
        r: usize,
    ) -> Option<impl IntoIterator<Item = T, IntoIter = impl Iterator<Item = T> + Send + Sync>> {
        (r < self.height()).then(|| unsafe {
            // Safety: Clearly `r < self.height()`.
            self.row_unchecked(r)
        })
    }

    /// Returns an iterator over the elements of the `r`-th row.
    ///
    /// The iterator will have `self.width()` elements.
    ///
    /// For a safe alternative, see [`Self::row`].
    ///
    /// # Safety
    /// The caller must ensure that `r < self.height()`.
    /// Breaking this assumption is considered undefined behaviour.
    #[inline]
    unsafe fn row_unchecked(
        &self,
        r: usize,
    ) -> impl IntoIterator<Item = T, IntoIter = impl Iterator<Item = T> + Send + Sync> {
        unsafe { self.row_subseq_unchecked(r, 0, self.width()) }
    }

    /// Returns an iterator over the elements of the `r`-th row from position `start` to `end`.
    ///
    /// When `start = 0` and `end = width()`, this is equivalent to [`Self::row_unchecked`].
    ///
    /// For a safe alternative, use [`Self::row`], along with the `skip` and `take` iterator methods.
    ///
    /// # Safety
    /// The caller must ensure that `r < self.height()` and `start <= end <= self.width()`.
    /// Breaking any of these assumptions is considered undefined behaviour.
    #[inline]
    unsafe fn row_subseq_unchecked(
        &self,
        r: usize,
        start: usize,
        end: usize,
    ) -> impl IntoIterator<Item = T, IntoIter = impl Iterator<Item = T> + Send + Sync> {
        unsafe {
            self.row_unchecked(r)
                .into_iter()
                .skip(start)
                .take(end - start)
        }
    }

    /// Returns the elements of the `r`-th row as something which can be coerced to a slice.
    ///
    /// Returns `None` if `r >= height()`.
    #[inline]
    fn row_slice(&self, r: usize) -> Option<impl Deref<Target = [T]>> {
        (r < self.height()).then(|| unsafe {
            // Safety: Clearly `r < self.height()`.
            self.row_slice_unchecked(r)
        })
    }

    /// Returns the elements of the `r`-th row as something which can be coerced to a slice.
    ///
    /// For a safe alternative, see [`Self::row_slice`].
    ///
    /// # Safety
    /// The caller must ensure that `r < self.height()`.
    /// Breaking this assumption is considered undefined behaviour.
    #[inline]
    unsafe fn row_slice_unchecked(&self, r: usize) -> impl Deref<Target = [T]> {
        unsafe { self.row_subslice_unchecked(r, 0, self.width()) }
    }

    /// Returns a subset of elements of the `r`-th row as something which can be coerced to a slice.
    ///
    /// When `start = 0` and `end = width()`, this is equivalent to [`Self::row_slice_unchecked`].
    ///
    /// For a safe alternative, see [`Self::row_slice`].
    ///
    /// # Safety
    /// The caller must ensure that `r < self.height()` and `start <= end <= self.width()`.
    /// Breaking any of these assumptions is considered undefined behaviour.
    #[inline]
    unsafe fn row_subslice_unchecked(
        &self,
        r: usize,
        start: usize,
        end: usize,
    ) -> impl Deref<Target = [T]> {
        unsafe {
            self.row_subseq_unchecked(r, start, end)
                .into_iter()
                .collect_vec()
        }
    }

    /// Returns an iterator over all rows in the matrix.
    #[inline]
    fn rows(&self) -> impl Iterator<Item = impl Iterator<Item = T>> + Send + Sync {
        unsafe {
            // Safety: `r` always satisfies `r < self.height()`.
            (0..self.height()).map(move |r| self.row_unchecked(r).into_iter())
        }
    }

    /// Returns a parallel iterator over all rows in the matrix.
    #[inline]
    fn par_rows(
        &self,
    ) -> impl IndexedParallelIterator<Item = impl Iterator<Item = T>> + Send + Sync {
        unsafe {
            // Safety: `r` always satisfies `r < self.height()`.
            (0..self.height())
                .into_par_iter()
                .map(move |r| self.row_unchecked(r).into_iter())
        }
    }

    /// Collect the elements of the rows `r` through `r + c`. If anything is larger than `self.height()`
    /// simply wrap around to the beginning of the matrix.
    fn wrapping_row_slices(&self, r: usize, c: usize) -> Vec<impl Deref<Target = [T]>> {
        unsafe {
            // Safety: Thank to the `%`, the rows index is always less than `self.height()`.
            (0..c)
                .map(|i| self.row_slice_unchecked((r + i) % self.height()))
                .collect_vec()
        }
    }

    /// Returns an iterator over the first row of the matrix.
    ///
    /// Returns None if `height() == 0`.
    #[inline]
    fn first_row(
        &self,
    ) -> Option<impl IntoIterator<Item = T, IntoIter = impl Iterator<Item = T> + Send + Sync>> {
        self.row(0)
    }

    /// Returns an iterator over the last row of the matrix.
    ///
    /// Returns None if `height() == 0`.
    #[inline]
    fn last_row(
        &self,
    ) -> Option<impl IntoIterator<Item = T, IntoIter = impl Iterator<Item = T> + Send + Sync>> {
        if self.height() == 0 {
            None
        } else {
            // Safety: Clearly `self.height() - 1 < self.height()`.
            unsafe { Some(self.row_unchecked(self.height() - 1)) }
        }
    }

    /// Converts the matrix into a `RowMajorMatrix` by collecting all rows into a single vector.
    fn to_row_major_matrix(self) -> RowMajorMatrix<T>
    where
        Self: Sized,
        T: Clone,
    {
        RowMajorMatrix::new(self.rows().flatten().collect(), self.width())
    }

    /// Get a packed iterator over the `r`-th row.
    ///
    /// If the row length is not divisible by the packing width, the final elements
    /// are returned as a base iterator with length `<= P::WIDTH - 1`.
    ///
    /// # Panics
    /// Panics if `r >= height()`.
    fn horizontally_packed_row<'a, P>(
        &'a self,
        r: usize,
    ) -> (
        impl Iterator<Item = P> + Send + Sync,
        impl Iterator<Item = T> + Send + Sync,
    )
    where
        P: PackedValue<Value = T>,
        T: Clone + 'a,
    {
        assert!(r < self.height(), "Row index out of bounds.");
        let num_packed = self.width() / P::WIDTH;
        unsafe {
            // Safety: We have already checked that `r < height()`.
            let mut iter = self
                .row_subseq_unchecked(r, 0, num_packed * P::WIDTH)
                .into_iter();

            // array::from_fn is guaranteed to always call in order.
            let packed =
                (0..num_packed).map(move |_| P::from_fn(|_| iter.next().unwrap_unchecked()));

            let sfx = self
                .row_subseq_unchecked(r, num_packed * P::WIDTH, self.width())
                .into_iter();
            (packed, sfx)
        }
    }

    /// Get a packed iterator over the `r`-th row.
    ///
    /// If the row length is not divisible by the packing width, the final entry will be zero-padded.
    ///
    /// # Panics
    /// Panics if `r >= height()`.
    fn padded_horizontally_packed_row<'a, P>(
        &'a self,
        r: usize,
    ) -> impl Iterator<Item = P> + Send + Sync
    where
        P: PackedValue<Value = T>,
        T: Clone + Default + 'a,
    {
        let mut row_iter = self.row(r).expect("Row index out of bounds.").into_iter();
        let num_elems = self.width().div_ceil(P::WIDTH);
        // array::from_fn is guaranteed to always call in order.
        (0..num_elems).map(move |_| P::from_fn(|_| row_iter.next().unwrap_or_default()))
    }

    /// Get a parallel iterator over all packed rows of the matrix.
    ///
    /// If the matrix width is not divisible by the packing width, the final elements
    /// of each row are returned as a base iterator with length `<= P::WIDTH - 1`.
    fn par_horizontally_packed_rows<'a, P>(
        &'a self,
    ) -> impl IndexedParallelIterator<
        Item = (
            impl Iterator<Item = P> + Send + Sync,
            impl Iterator<Item = T> + Send + Sync,
        ),
    >
    where
        P: PackedValue<Value = T>,
        T: Clone + 'a,
    {
        (0..self.height())
            .into_par_iter()
            .map(|r| self.horizontally_packed_row(r))
    }

    /// Get a parallel iterator over all packed rows of the matrix.
    ///
    /// If the matrix width is not divisible by the packing width, the final entry of each row will be zero-padded.
    fn par_padded_horizontally_packed_rows<'a, P>(
        &'a self,
    ) -> impl IndexedParallelIterator<Item = impl Iterator<Item = P> + Send + Sync>
    where
        P: PackedValue<Value = T>,
        T: Clone + Default + 'a,
    {
        (0..self.height())
            .into_par_iter()
            .map(|r| self.padded_horizontally_packed_row(r))
    }

    /// Pack together a collection of adjacent rows from the matrix.
    ///
    /// Returns an iterator whose i'th element is packing of the i'th element of the
    /// rows r through r + P::WIDTH - 1. If we exceed the height of the matrix,
    /// wrap around and include initial rows.
    #[inline]
    fn vertically_packed_row<P>(&self, r: usize) -> impl Iterator<Item = P>
    where
        T: Copy,
        P: PackedValue<Value = T>,
    {
        // Precompute row slices once to minimize redundant calls and improve performance.
        let rows = self.wrapping_row_slices(r, P::WIDTH);

        // Using precomputed rows avoids repeatedly calling `row_slice`, which is costly.
        (0..self.width()).map(move |c| P::from_fn(|i| rows[i][c]))
    }

    /// Pack together a collection of rows and "next" rows from the matrix.
    ///
    /// Returns a vector corresponding to 2 packed rows. The i'th element of the first
    /// row contains the packing of the i'th element of the rows r through r + P::WIDTH - 1.
    /// The i'th element of the second row contains the packing of the i'th element of the
    /// rows r + step through r + step + P::WIDTH - 1. If at some point we exceed the
    /// height of the matrix, wrap around and include initial rows.
    #[inline]
    fn vertically_packed_row_pair<P>(&self, r: usize, step: usize) -> Vec<P>
    where
        T: Copy,
        P: PackedValue<Value = T>,
    {
        // Whilst it would appear that this can be replaced by two calls to vertically_packed_row
        // tests seem to indicate that combining them in the same function is slightly faster.
        // It's probably allowing the compiler to make some optimizations on the fly.

        let rows = self.wrapping_row_slices(r, P::WIDTH);
        let next_rows = self.wrapping_row_slices(r + step, P::WIDTH);

        (0..self.width())
            .map(|c| P::from_fn(|i| rows[i][c]))
            .chain((0..self.width()).map(|c| P::from_fn(|i| next_rows[i][c])))
            .collect_vec()
    }

    /// Returns a view over a vertically strided submatrix.
    ///
    /// The view selects rows using `r = offset + i * stride` for each `i`.
    fn vertically_strided(self, stride: usize, offset: usize) -> VerticallyStridedMatrixView<Self>
    where
        Self: Sized,
    {
        VerticallyStridedRowIndexMap::new_view(self, stride, offset)
    }

    /// Compute Mᵀv, aka premultiply this matrix by the given vector,
    /// aka scale each row by the corresponding entry in `v` and take the sum across rows.
    /// `v` can be a vector of extension elements.
    #[instrument(level = "debug", skip_all, fields(dims = %self.dimensions()))]
    fn columnwise_dot_product<EF>(&self, v: &[EF]) -> Vec<EF>
    where
        T: Field,
        EF: ExtensionField<T>,
    {
        assert_eq!(
            v.len(),
            self.height(),
            "weight count must match matrix height"
        );

        // Below this many total elements, the rayon fork-join and SIMD-packing machinery
        // costs more than the dot product itself; fall back to a plain scalar accumulation.
        // Gating on total elements (rather than height alone) also covers wide-but-short
        // matrices, where a per-row cost proportional to width still adds up.
        const SMALL_ELEMS: usize = 256;
        if self.height().saturating_mul(self.width()) <= SMALL_ELEMS {
            let mut acc = EF::zero_vec(self.width());
            for (row, &scale) in self.rows().zip(v) {
                for (l, r) in acc.iter_mut().zip(row) {
                    *l += scale * r;
                }
            }
            return acc;
        }

        let packed_width = self.width().div_ceil(T::Packing::WIDTH);

        // Avoid multi-worker scheduling for modest products while bounding both input
        // traffic and per-call extension-coefficient work. Larger products, serial
        // builds, and one-worker pools retain the existing reduction.
        //
        // The task-size model cannot place this gate, since the two arms are different
        // kernels rather than one kernel cut two ways.
        //
        // The serial arm defers modular reductions across rows; the reduction below cannot.
        //
        // So on Zen 5 at 32 workers the serial arm stays ahead to between 80 and 340 us of
        // priced work, where the model asks for 20.
        const SERIAL_PACKED_ELEMS: usize = 4096;
        const SERIAL_PACKED_COEFF_OPS: usize = 512;
        if T::Packing::WIDTH > 1
            && self.height() > 1
            && self.height().saturating_mul(self.width()) <= SERIAL_PACKED_ELEMS
            && self
                .height()
                .saturating_mul(packed_width)
                .saturating_mul(EF::DIMENSION)
                <= SERIAL_PACKED_COEFF_OPS
            && PARALLEL_ENABLED
            && current_num_threads() > 1
        {
            return serial_packed_columnwise_dot_product(self, v);
        }

        let packed_result = self
            .par_padded_horizontally_packed_rows::<T::Packing>()
            .zip(v)
            .par_fold_reduce(
                || EF::ExtensionPacking::zero_vec(packed_width),
                |mut acc, (row, &scale)| {
                    let scale: EF::ExtensionPacking = scale.into();
                    acc.iter_mut().zip(row).for_each(|(l, r)| *l += scale * r);
                    acc
                },
                |mut acc_l, acc_r| {
                    acc_l.iter_mut().zip(&acc_r).for_each(|(l, r)| *l += *r);
                    acc_l
                },
            );

        EF::ExtensionPacking::to_ext_iter(packed_result)
            .take(self.width())
            .collect()
    }

    /// Compute Mᵀ · [v₀, v₁, ..., vₙ₋₁] for N weight vectors simultaneously.
    ///
    /// Computes `result[col][j] = Σᵣ M[r, col] · vⱼ[r]` for all columns and all j ∈ [0, N).
    ///
    /// Batching N weight vectors reduces memory bandwidth: each matrix row is loaded once
    /// instead of N times. Uses SIMD packing (width W) to process W columns in parallel.
    #[instrument(level = "debug", skip_all, fields(dims = %self.dimensions()))]
    fn columnwise_dot_product_batched<EF, const N: usize>(
        &self,
        vs: &[FieldArray<EF, N>],
    ) -> Vec<FieldArray<EF, N>>
    where
        T: Field,
        EF: ExtensionField<T>,
    {
        assert_eq!(vs.len(), self.height());

        let packed_width = self.width().div_ceil(T::Packing::WIDTH);
        let height = self.height();

        // Split the rows into a bounded number of contiguous chunks; each task runs the
        // field's columnwise kernel serially over its chunk (letting it defer modular
        // reductions across rows) and the per-task accumulators are summed at the end.
        //
        // The floor collapses the split to a single chunk on a small matrix.
        //
        // Below that size the work is not worth handing to another worker.
        let row_bytes = columnwise_row_bytes::<T, EF>(self.width(), N);
        // The floor is never zero, so the chunk length is always usable as a divisor.
        let chunk_rows = height
            .div_ceil((4 * current_num_threads()).clamp(1, height.max(1)))
            .max(min_task_len(height, row_bytes));
        let num_chunks = height.div_ceil(chunk_rows);

        let packed_results: Vec<EF::ExtensionPacking> =
            (0..num_chunks).into_par_iter().par_fold_reduce(
                || EF::ExtensionPacking::zero_vec(packed_width * N),
                |mut acc, chunk| {
                    let rows = chunk * chunk_rows..((chunk + 1) * chunk_rows).min(height);
                    T::batched_columnwise_dot_product::<EF, _, _, N>(
                        &mut acc,
                        rows.map(|r| {
                            (
                                self.padded_horizontally_packed_row::<T::Packing>(r),
                                vs[r].0,
                            )
                        }),
                    );
                    acc
                },
                |mut acc_l, acc_r| {
                    acc_l.iter_mut().zip(&acc_r).for_each(|(lj, rj)| *lj += *rj);
                    acc_l
                },
            );

        // Unpack: chunk[j].lane(i) → result[c·W + i][j] for column batch c
        packed_results
            .chunks(N)
            .flat_map(|chunk| {
                (0..T::Packing::WIDTH)
                    .map(move |lane| FieldArray::from_fn(|j| chunk[j].extract(lane)))
            })
            .take(self.width())
            .collect()
    }

    /// Compute the matrix vector product `M . vec`, aka take the dot product of each
    /// row of `M` by `vec`. If the length of `vec` is longer than the width of `M`,
    /// `vec` is truncated to the first `width()` elements.
    ///
    /// We make use of `PackedFieldExtension` to speed up computations. Thus `vec` is passed in as
    /// a slice of `PackedFieldExtension` elements.
    ///
    /// # Panics
    /// This function panics if the length of `vec` is less than `self.width().div_ceil(T::Packing::WIDTH)`.
    fn rowwise_packed_dot_product<EF>(
        &self,
        vec: &[EF::ExtensionPacking],
    ) -> impl IndexedParallelIterator<Item = EF>
    where
        T: Field,
        EF: ExtensionField<T>,
    {
        // The length of a `padded_horizontally_packed_row` is `self.width().div_ceil(T::Packing::WIDTH)`.
        assert!(vec.len() >= self.width().div_ceil(T::Packing::WIDTH));

        // Instead of creating N intermediate ExtPacking products and summing them,
        // we track D separate BasePacking accumulators (one per extension coefficient).
        self.par_padded_horizontally_packed_rows::<T::Packing>()
            .map(move |row_packed| {
                // Get the extension dimension from the first vec element's coefficients
                let d = <EF::ExtensionPacking as BasedVectorSpace<T::Packing>>::DIMENSION;

                // Accumulate coefficient-wise: for each (v, r) pair, acc[i] += v.coefficient(i) * r
                let coeff_accs = T::Packing::coeffwise_dot_product(
                    d,
                    vec.iter()
                        .zip(row_packed)
                        .map(|(v, r)| (v.as_basis_coefficients_slice(), r)),
                );

                // Construct the result ExtPacking from the accumulators and sum the coefficients.
                let packed_result =
                    EF::ExtensionPacking::from_basis_coefficients_fn(|i| coeff_accs[i]);
                EF::ExtensionPacking::to_ext_iter([packed_result]).sum()
            })
    }
}

/// Extension widths one [`COLUMNWISE_MAC_LANES`]-column multiply-accumulate step is charged.
///
/// Measured on Zen 5 at four weight vectors, [`COLUMNWISE_MAC_LANES`] columns per step:
///
/// ```text
///     columns   real bytes per row   widths per step
///          16                  602               7.4
///          64                 2235               7.5
///         256                 8716               7.4
///         512                17424               7.5
/// ```
const COLUMNWISE_MAC_WIDTHS: usize = 7;

/// Columns one packed step carried on the host the figure above was measured on.
///
/// Seven widths over sixteen lanes is 0.74 ns per column.
///
/// A four-lane NEON host reproduces that per column, not per packed step.
///
/// So the arithmetic is charged per column; per packed element it would scale with the
/// host's lane count and overcharge that build by about 3.5x.
const COLUMNWISE_MAC_LANES: usize = 16;

/// Bytes one matrix row moves when weighted into `weight_vectors` accumulators.
///
/// ```text
///     traffic : one packed row, plus the weights it scales
///     arith   : one multiply-accumulate per column, per weight vector
/// ```
///
/// Pricing the row by traffic alone puts the gate seven times too high, leaving a
/// compute-bound matrix serial: 64 columns by 512 rows costs 114 us against 19 us split.
const fn columnwise_row_bytes<T, EF>(width: usize, weight_vectors: usize) -> usize
where
    T: Field,
    EF: ExtensionField<T>,
{
    // The kernel walks whole packed words, so padding lanes are multiplied like any other.
    //
    // That rounding is the only place the host's lane count enters the price.
    let columns = width.div_ceil(T::Packing::WIDTH) * T::Packing::WIDTH;

    columns * size_of::<T>()
        + weight_vectors * size_of::<EF>()
        + (COLUMNWISE_MAC_WIDTHS * weight_vectors * columns * size_of::<EF>())
            .div_ceil(COLUMNWISE_MAC_LANES)
}

#[inline(never)]
fn serial_packed_columnwise_dot_product<F, EF, M>(matrix: &M, weights: &[EF]) -> Vec<EF>
where
    F: Field,
    EF: ExtensionField<F>,
    M: Matrix<F> + ?Sized,
{
    let mut acc = EF::ExtensionPacking::zero_vec(matrix.width().div_ceil(F::Packing::WIDTH));
    F::batched_columnwise_dot_product::<EF, _, _, 1>(
        &mut acc,
        (0..matrix.height()).map(|r| {
            (
                matrix.padded_horizontally_packed_row::<F::Packing>(r),
                [weights[r]],
            )
        }),
    );
    EF::ExtensionPacking::to_ext_iter(acc)
        .take(matrix.width())
        .collect()
}

#[cfg(test)]
mod tests {
    use alloc::vec::Vec;
    use alloc::{format, vec};

    use itertools::izip;
    use p3_baby_bear::BabyBear;
    use p3_field::PrimeCharacteristicRing;
    use p3_field::extension::{BinomialExtensionField, CubicTrinomialExtensionField};
    use p3_goldilocks::Goldilocks;
    use p3_mersenne_31::{Mersenne31, QM31};
    use rand::SeedableRng;
    use rand::rngs::SmallRng;

    use super::*;
    use crate::bitrev::BitReversibleMatrix;
    use crate::extension::FlatMatrixView;

    fn patterned_matrix<F: Field>(height: usize, width: usize) -> RowMajorMatrix<F> {
        RowMajorMatrix::new(
            (0..height * width)
                .map(|i| F::from_usize((i * 17 + 3) % 127))
                .collect(),
            width,
        )
    }

    fn patterned_extension_matrix<F, EF>(height: usize, width: usize) -> RowMajorMatrix<EF>
    where
        F: Field,
        EF: ExtensionField<F>,
    {
        RowMajorMatrix::new(
            (0..height * width)
                .map(|i| {
                    EF::from_basis_coefficients_fn(|d| {
                        F::from_usize((i * EF::DIMENSION + d + 1) % 127)
                    })
                })
                .collect(),
            width,
        )
    }

    fn assert_columnwise_dot_product_matches_scalar<F, EF, M>(mat: &M)
    where
        F: Field,
        EF: ExtensionField<F>,
        M: Matrix<F>,
    {
        let weights: Vec<EF> = (0..mat.height())
            .map(|r| {
                EF::from_basis_coefficients_fn(|d| F::from_usize((r * EF::DIMENSION + d + 5) % 127))
            })
            .collect();
        let expected: Vec<EF> = (0..mat.width())
            .map(|c| {
                (0..mat.height())
                    .map(|r| weights[r] * mat.get(r, c).unwrap())
                    .sum()
            })
            .collect();

        assert_eq!(mat.columnwise_dot_product(&weights), expected);
    }

    fn assert_columnwise_dot_product_grid<F, EF>()
    where
        F: Field,
        EF: ExtensionField<F>,
    {
        for height in [0, 1, 17, 32, 128, 1024] {
            for width in [1, 3, 8, 17, 65] {
                let mat = patterned_matrix::<F>(height, width);
                assert_columnwise_dot_product_matches_scalar::<F, EF, _>(&mat);
            }
        }
    }

    #[test]
    fn test_columnwise_dot_product() {
        type F = BabyBear;
        type EF = BinomialExtensionField<BabyBear, 4>;

        let mut rng = SmallRng::seed_from_u64(1);
        let m = RowMajorMatrix::<F>::rand(&mut rng, 1 << 8, 1 << 4);
        let v = RowMajorMatrix::<EF>::rand(&mut rng, 1 << 8, 1).values;

        let mut expected = EF::zero_vec(m.width());
        for (row, &scale) in izip!(m.rows(), &v) {
            for (l, r) in izip!(&mut expected, row) {
                *l += scale * r;
            }
        }

        assert_eq!(m.columnwise_dot_product(&v), expected);
    }

    #[test]
    fn test_columnwise_dot_product_small_height() {
        type F = BabyBear;
        type EF = BinomialExtensionField<BabyBear, 4>;

        let mut rng = SmallRng::seed_from_u64(2);

        // Cover heights below, at, and just above the small-height serial threshold.
        for height in [0, 1, 3, 16, 17] {
            let m = RowMajorMatrix::<F>::rand(&mut rng, height, 1 << 4);
            let v = RowMajorMatrix::<EF>::rand(&mut rng, height, 1).values;

            let mut expected = EF::zero_vec(m.width());
            for (row, &scale) in izip!(m.rows(), &v) {
                for (l, r) in izip!(&mut expected, row) {
                    *l += scale * r;
                }
            }

            assert_eq!(m.columnwise_dot_product(&v), expected, "height = {height}");
        }
    }

    #[test]
    fn test_columnwise_dot_product_matches_scalar_across_extension_fields() {
        type BabyBear4 = BinomialExtensionField<BabyBear, 4>;
        type BabyBear5 = BinomialExtensionField<BabyBear, 5>;
        type Goldilocks2 = BinomialExtensionField<Goldilocks, 2>;
        type Goldilocks3 = CubicTrinomialExtensionField<Goldilocks>;
        type Mersenne31_3 = BinomialExtensionField<Mersenne31, 3>;

        assert_columnwise_dot_product_grid::<BabyBear, BabyBear4>();
        assert_columnwise_dot_product_grid::<BabyBear, BabyBear5>();
        assert_columnwise_dot_product_grid::<Goldilocks, Goldilocks2>();
        assert_columnwise_dot_product_grid::<Goldilocks, Goldilocks3>();
        assert_columnwise_dot_product_grid::<Mersenne31, QM31>();
        assert_columnwise_dot_product_grid::<Mersenne31, Mersenne31_3>();
    }

    #[test]
    fn test_columnwise_dot_product_matches_scalar_for_matrix_views() {
        type BabyBear4 = BinomialExtensionField<BabyBear, 4>;
        type Goldilocks3 = CubicTrinomialExtensionField<Goldilocks>;
        type Mersenne31_3 = BinomialExtensionField<Mersenne31, 3>;

        let mapped = patterned_matrix::<BabyBear>(32, 17).bit_reverse_rows();
        assert_columnwise_dot_product_matches_scalar::<BabyBear, BabyBear4, _>(&mapped);

        let flat = FlatMatrixView::<Goldilocks, Goldilocks3, _>::new(patterned_extension_matrix::<
            Goldilocks,
            Goldilocks3,
        >(32, 3));
        assert_columnwise_dot_product_matches_scalar::<Goldilocks, Goldilocks3, _>(&flat);

        let mapped_flat = FlatMatrixView::<Mersenne31, Mersenne31_3, _>::new(
            patterned_extension_matrix::<Mersenne31, Mersenne31_3>(128, 3).bit_reverse_rows(),
        );
        assert_columnwise_dot_product_matches_scalar::<Mersenne31, Mersenne31_3, _>(&mapped_flat);
    }

    #[test]
    #[should_panic(expected = "weight count must match matrix height")]
    fn test_columnwise_dot_product_rejects_short_weights() {
        let mat = patterned_matrix::<BabyBear>(17, 17);
        let weights = BinomialExtensionField::<BabyBear, 4>::zero_vec(16);
        let _ = mat.columnwise_dot_product(&weights);
    }

    #[test]
    #[should_panic(expected = "weight count must match matrix height")]
    fn test_columnwise_dot_product_rejects_long_weights() {
        let mat = patterned_matrix::<BabyBear>(17, 17);
        let weights = BinomialExtensionField::<BabyBear, 4>::zero_vec(18);
        let _ = mat.columnwise_dot_product(&weights);
    }

    #[test]
    fn test_columnwise_dot_product_batched() {
        type F = BabyBear;
        type EF = BinomialExtensionField<BabyBear, 4>;

        let mut rng = SmallRng::seed_from_u64(1);
        let m = RowMajorMatrix::<F>::rand(&mut rng, 1 << 8, 1 << 4);
        let v1 = RowMajorMatrix::<EF>::rand(&mut rng, 1 << 8, 1).values;
        let v2 = RowMajorMatrix::<EF>::rand(&mut rng, 1 << 8, 1).values;

        // Compute expected via two separate calls
        let expected1 = m.columnwise_dot_product(&v1);
        let expected2 = m.columnwise_dot_product(&v2);

        // Compute via batched call - returns Vec<[EF; 2]> where result[col] = [dot1, dot2]
        let vs: Vec<FieldArray<EF, 2>> = v1
            .into_iter()
            .zip(v2)
            .map(|(a, b)| FieldArray([a, b]))
            .collect();
        let results = m.columnwise_dot_product_batched::<EF, 2>(&vs);

        // Extract each point's results
        let result1: Vec<EF> = results.iter().map(|r| r[0]).collect();
        let result2: Vec<EF> = results.iter().map(|r| r[1]).collect();

        assert_eq!(result1, expected1);
        assert_eq!(result2, expected2);
    }

    // Mock implementation for testing purposes
    struct MockMatrix {
        data: Vec<Vec<u32>>,
        width: usize,
        height: usize,
    }

    impl Matrix<u32> for MockMatrix {
        fn width(&self) -> usize {
            self.width
        }

        fn height(&self) -> usize {
            self.height
        }

        unsafe fn row_unchecked(
            &self,
            r: usize,
        ) -> impl IntoIterator<Item = u32, IntoIter = impl Iterator<Item = u32> + Send + Sync>
        {
            // Just a mock implementation so we just do the easy safe thing.
            self.data[r].clone()
        }
    }

    #[test]
    fn test_dimensions() {
        let dims = Dimensions {
            width: 3,
            height: 5,
        };
        assert_eq!(dims.width, 3);
        assert_eq!(dims.height, 5);
        assert_eq!(format!("{dims:?}"), "3x5");
        assert_eq!(format!("{dims}"), "3x5");
    }

    #[test]
    fn test_mock_matrix_dimensions() {
        let matrix = MockMatrix {
            data: vec![vec![1, 2, 3], vec![4, 5, 6], vec![7, 8, 9]],
            width: 3,
            height: 3,
        };
        assert_eq!(matrix.width(), 3);
        assert_eq!(matrix.height(), 3);
        assert_eq!(
            matrix.dimensions(),
            Dimensions {
                width: 3,
                height: 3
            }
        );
    }

    #[test]
    fn test_first_row() {
        let matrix = MockMatrix {
            data: vec![vec![1, 2, 3], vec![4, 5, 6], vec![7, 8, 9]],
            width: 3,
            height: 3,
        };
        let mut first_row = matrix.first_row().unwrap().into_iter();
        assert_eq!(first_row.next(), Some(1));
        assert_eq!(first_row.next(), Some(2));
        assert_eq!(first_row.next(), Some(3));
    }

    #[test]
    fn test_last_row() {
        let matrix = MockMatrix {
            data: vec![vec![1, 2, 3], vec![4, 5, 6], vec![7, 8, 9]],
            width: 3,
            height: 3,
        };
        let mut last_row = matrix.last_row().unwrap().into_iter();
        assert_eq!(last_row.next(), Some(7));
        assert_eq!(last_row.next(), Some(8));
        assert_eq!(last_row.next(), Some(9));
    }

    #[test]
    fn test_first_last_row_empty_matrix() {
        let matrix = MockMatrix {
            data: vec![],
            width: 3,
            height: 0,
        };
        let first_row = matrix.first_row();
        let last_row = matrix.last_row();
        assert!(first_row.is_none());
        assert!(last_row.is_none());
    }

    #[test]
    fn test_to_row_major_matrix() {
        let matrix = MockMatrix {
            data: vec![vec![1, 2], vec![3, 4]],
            width: 2,
            height: 2,
        };
        let row_major = matrix.to_row_major_matrix();
        assert_eq!(row_major.values, vec![1, 2, 3, 4]);
        assert_eq!(row_major.width, 2);
    }

    #[test]
    fn test_matrix_get_methods() {
        let matrix = MockMatrix {
            data: vec![vec![1, 2, 3], vec![4, 5, 6], vec![7, 8, 9]],
            width: 3,
            height: 3,
        };
        assert_eq!(matrix.get(0, 0), Some(1));
        assert_eq!(matrix.get(1, 2), Some(6));
        assert_eq!(matrix.get(2, 1), Some(8));

        unsafe {
            assert_eq!(matrix.get_unchecked(0, 1), 2);
            assert_eq!(matrix.get_unchecked(1, 0), 4);
            assert_eq!(matrix.get_unchecked(2, 2), 9);
        }

        assert_eq!(matrix.get(3, 0), None); // Height out of bounds
        assert_eq!(matrix.get(0, 3), None); // Width out of bounds
    }

    #[test]
    fn test_matrix_row_methods_iteration() {
        let matrix = MockMatrix {
            data: vec![vec![1, 2, 3], vec![4, 5, 6], vec![7, 8, 9]],
            width: 3,
            height: 3,
        };

        let mut row_iter = matrix.row(1).unwrap().into_iter();
        assert_eq!(row_iter.next(), Some(4));
        assert_eq!(row_iter.next(), Some(5));
        assert_eq!(row_iter.next(), Some(6));
        assert_eq!(row_iter.next(), None);

        unsafe {
            let mut row_iter_unchecked = matrix.row_unchecked(2).into_iter();
            assert_eq!(row_iter_unchecked.next(), Some(7));
            assert_eq!(row_iter_unchecked.next(), Some(8));
            assert_eq!(row_iter_unchecked.next(), Some(9));
            assert_eq!(row_iter_unchecked.next(), None);

            let mut row_iter_subset = matrix.row_subseq_unchecked(0, 1, 3).into_iter();
            assert_eq!(row_iter_subset.next(), Some(2));
            assert_eq!(row_iter_subset.next(), Some(3));
            assert_eq!(row_iter_subset.next(), None);
        }

        assert!(matrix.row(3).is_none()); // Height out of bounds
    }

    #[test]
    fn test_row_slice_methods() {
        let matrix = MockMatrix {
            data: vec![vec![1, 2, 3], vec![4, 5, 6], vec![7, 8, 9]],
            width: 3,
            height: 3,
        };
        let row_slice = matrix.row_slice(1).unwrap();
        assert_eq!(*row_slice, [4, 5, 6]);
        unsafe {
            let row_slice_unchecked = matrix.row_slice_unchecked(2);
            assert_eq!(*row_slice_unchecked, [7, 8, 9]);

            let row_subslice = matrix.row_subslice_unchecked(0, 1, 2);
            assert_eq!(*row_subslice, [2]);
        }

        assert!(matrix.row_slice(3).is_none()); // Height out of bounds
    }

    #[test]
    fn test_matrix_rows() {
        let matrix = MockMatrix {
            data: vec![vec![1, 2, 3], vec![4, 5, 6], vec![7, 8, 9]],
            width: 3,
            height: 3,
        };

        let all_rows: Vec<Vec<u32>> = matrix.rows().map(|row| row.collect()).collect();
        assert_eq!(all_rows, vec![vec![1, 2, 3], vec![4, 5, 6], vec![7, 8, 9]]);
    }

    #[test]
    fn test_rowwise_packed_dot_product() {
        use p3_field::PackedFieldExtension;

        type F = BabyBear;
        type EF = BinomialExtensionField<BabyBear, 4>;
        type PF = <F as p3_field::Field>::Packing;
        type EFPacked = <EF as p3_field::ExtensionField<F>>::ExtensionPacking;

        let mut rng = SmallRng::seed_from_u64(42);

        // Test with various matrix dimensions to cover edge cases.
        for (height, width) in [(32, 16), (64, 128), (128, 17), (256, 255)] {
            let m = RowMajorMatrix::<F>::rand(&mut rng, height, width);
            let v = RowMajorMatrix::<EF>::rand(&mut rng, width, 1).values;

            // Compute expected result naively: for each row, compute dot product with v.
            let expected: Vec<EF> = m
                .rows()
                .map(|row| {
                    row.into_iter()
                        .zip(v.iter())
                        .map(|(r, &ve)| ve * r)
                        .sum::<EF>()
                })
                .collect();

            // Pack the vector for the optimized function.
            let packed_v: Vec<EFPacked> = v
                .chunks(<PF as PackedValue>::WIDTH)
                .map(|chunk| {
                    let mut padded = EF::zero_vec(<PF as PackedValue>::WIDTH);
                    padded[..chunk.len()].copy_from_slice(chunk);
                    EFPacked::from_ext_slice(&padded)
                })
                .collect();

            // Compute using the optimized function.
            let result: Vec<EF> = m.rowwise_packed_dot_product::<EF>(&packed_v).collect();

            assert_eq!(result, expected, "Mismatch for matrix {}x{}", height, width);
        }
    }
}