primitives/types/heap_array/

matrix.rs

use std::{fmt::Debug, marker::PhantomData};

use itertools::Itertools;
use serde::{Deserialize, Serialize};
use typenum::Prod;

use super::HeapArray;
use crate::{types::Positive, utils::IntoExactSizeIterator};

/// A matrix with M rows and N columns on the heap that encodes its shape in the type system.
///
/// The matrix is stored as a contiguous memory chunk in column-major order.
#[derive(Clone, PartialEq, Eq)]
pub struct HeapMatrix<T: Sized, M: Positive, N: Positive> {
    pub(super) data: Box<[T]>,

    // `fn() -> (M, N)` is used instead of `(M, N)` so `HeapMatrix<T, M, N>` doesn't need `(M, N)`
    // to implement `Send + Sync` to be `Send + Sync` itself. This would be the case if `(M, N)`
    // was used directly.
    pub(super) _len: PhantomData<fn() -> (M, N)>,
}

impl<T: Sized, M: Positive, N: Positive> HeapMatrix<T, M, N> {
    fn new(data: Box<[T]>) -> Self {
        Self {
            data,
            _len: PhantomData,
        }
    }

    /// All matrix elements iterator in column-major order
    pub fn flat_iter(&self) -> impl ExactSizeIterator<Item = &T> {
        self.data.iter()
    }

    /// All matrix elements mutable iterator in column-major order
    pub fn flat_iter_mut(&mut self) -> impl ExactSizeIterator<Item = &mut T> {
        self.data.iter_mut()
    }

    /// Transform all matrix elements into owned-values iterator in column-major order
    pub fn into_flat_iter(self) -> impl ExactSizeIterator<Item = T> {
        self.data.into_vec().into_iter()
    }

    /// Matrix column iterator
    pub fn col_iter(&self) -> impl ExactSizeIterator<Item = &[T]> {
        self.data.chunks_exact(M::USIZE)
    }

    /// Matrix column mutable iterator
    pub fn col_iter_mut(&mut self) -> impl ExactSizeIterator<Item = &mut [T]> {
        self.data.chunks_exact_mut(M::USIZE)
    }

    /// Build a matrix from a elements in column-major order
    pub fn from_flat_iter(it: impl IntoExactSizeIterator<Item = T>) -> Self {
        let it = it.into_iter();
        assert_eq!(it.len(), M::USIZE * N::USIZE);
        Self::new(it.collect_vec().into_boxed_slice())
    }
}

impl<T: Sized, M: Positive, N: Positive> HeapMatrix<T, M, N>
where
    M: std::ops::Mul<N, Output: Positive>,
{
    /// Flatten matrix into an heap array in column-major order
    pub fn into_flat_array(self) -> HeapArray<T, Prod<M, N>> {
        HeapArray {
            data: self.data,
            _len: PhantomData,
        }
    }

    /// Build matrix from an heap array in column-major order
    pub fn from_flat_array(value: HeapArray<T, Prod<M, N>>) -> Self {
        Self::new(value.data)
    }

    /// Build a matrix from an array of columns
    pub fn from_cols_array(val: HeapArray<HeapArray<T, M>, N>) -> Self {
        Self::new(val.into_iter().flatten().collect())
    }
}

impl<T: Sized + Default, M: Positive, N: Positive> Default for HeapMatrix<T, M, N> {
    fn default() -> Self {
        Self::new((0..M::USIZE * N::USIZE).map(|_| T::default()).collect())
    }
}

impl<T: Sized + Debug, M: Positive, N: Positive> Debug for HeapMatrix<T, M, N> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct(format!("HeapMatrix<{}, {}>", M::USIZE, N::USIZE).as_str())
            .field("data", &self.data)
            .finish()
    }
}

impl<T: Sized + Serialize, M: Positive, N: Positive> Serialize for HeapMatrix<T, M, N> {
    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        self.data.serialize(serializer)
    }
}

impl<'de, T: Sized + Deserialize<'de>, M: Positive, N: Positive> Deserialize<'de>
    for HeapMatrix<T, M, N>
{
    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        let data = Box::<[T]>::deserialize(deserializer)?;

        if data.len() != M::USIZE * N::USIZE {
            return Err(serde::de::Error::custom(format!(
                "Expected array of length {}, got {}",
                M::USIZE,
                data.len()
            )));
        }

        Ok(Self {
            data,
            _len: PhantomData,
        })
    }
}

#[cfg(test)]
pub mod tests {
    use itertools::Itertools;
    use typenum::{U15, U3, U5};

    use crate::types::{heap_array::matrix::HeapMatrix, HeapArray};

    #[test]
    fn test_default() {
        let mut m = HeapMatrix::<usize, U3, U5>::default();
        assert_eq!(m.col_iter().len(), 5);
        for col in m.col_iter() {
            assert_eq!(col, &[0, 0, 0]);
        }

        // Change second column
        {
            let mut it = m.col_iter_mut();
            let _ = it.next();
            let c = it.next().unwrap();
            c[0] = 3;
            c[1] = 2;
            c[2] = 1;
        }

        assert_eq!(
            m.into_flat_iter().collect_vec(),
            vec![0, 0, 0, 3, 2, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0]
        );
    }

    #[test]
    fn test_into_from_flat_array() {
        let a = HeapArray::<usize, U15>::from_fn(|k| k);
        let m: HeapMatrix<usize, U3, U5> = HeapMatrix::from_flat_iter(a.iter().copied());
        let a1 = m.into_flat_array();
        assert_eq!(a, a1);
    }

    #[test]
    fn test_from_flat_array() {
        let a = HeapArray::<usize, U15>::from_fn(|k| k);
        let m: HeapMatrix<usize, U3, U5> = HeapMatrix::from_flat_array(a);
        let mut it = m.col_iter();
        assert_eq!(it.next().unwrap(), &[0, 1, 2]);
        assert_eq!(it.next().unwrap(), &[3, 4, 5]);
        assert_eq!(it.next().unwrap(), &[6, 7, 8]);
        assert_eq!(it.next().unwrap(), &[9, 10, 11]);
        assert_eq!(it.next().unwrap(), &[12, 13, 14]);
    }
}