primitives/types/heap_array/
matrix.rs

1use std::{
2    fmt::{Debug, Display},
3    marker::PhantomData,
4    ops::Mul,
5};
6
7use derive_more::derive::Display;
8use serde::{Deserialize, Serialize};
9use typenum::Prod;
10
11use super::HeapArray;
12use crate::{
13    errors::PrimitiveError,
14    random::{CryptoRngCore, Random},
15    types::Positive,
16};
17
18/// Indicates that the matrix is stored in column-major order (Fortran order).
19#[derive(Display, Default, Clone, Copy, PartialEq, Eq)]
20pub struct ColumnMajor;
21/// Indicates that the matrix is stored in row-major order (C order).
22#[derive(Display, Default, Clone, Copy, PartialEq, Eq)]
23pub struct RowMajor;
24
25pub type RowMajorHeapMatrix<T, M, N> = HeapMatrix<T, M, N, RowMajor>;
26
27/// A matrix with M rows and N columns on the heap that encodes its shape in the type system.
28///
29/// The matrix is stored as a contiguous memory chunk.
30#[derive(Clone, PartialEq, Eq)]
31#[repr(C)]
32pub struct HeapMatrix<T: Sized, M: Positive, N: Positive, O = ColumnMajor> {
33    pub(super) data: Box<[T]>,
34
35    // `fn() -> (M, N)` is used instead of `(M, N)` so `HeapMatrix<T, M, N>` doesn't need `(M, N)`
36    // to implement `Send + Sync` to be `Send + Sync` itself. This would be the case if `(M, N)`
37    // was used directly.
38    #[allow(clippy::type_complexity)]
39    pub(super) _len: PhantomData<fn() -> (M, N, O)>,
40}
41impl<T: Sized, M: Positive, N: Positive, O> HeapMatrix<T, M, N, O> {
42    fn new(data: Box<[T]>) -> Self {
43        Self {
44            data,
45            _len: PhantomData,
46        }
47    }
48
49    /// All matrix elements iterator in order.
50    pub fn flat_iter(&self) -> impl ExactSizeIterator<Item = &T> {
51        self.data.iter()
52    }
53
54    /// All matrix elements mutable iterator in order.
55    pub fn flat_iter_mut(&mut self) -> impl ExactSizeIterator<Item = &mut T> {
56        self.data.iter_mut()
57    }
58
59    /// Convert matrix into an iterator over all elements in order.
60    pub fn into_flat_iter(self) -> impl ExactSizeIterator<Item = T> {
61        self.data.into_vec().into_iter()
62    }
63
64    /// Length of total number of elements in the matrix
65    pub const fn len(&self) -> usize {
66        M::USIZE * N::USIZE
67    }
68
69    /// Check if the matrix is empty
70    pub const fn is_empty(&self) -> bool {
71        self.len() == 0
72    }
73
74    /// Number of rows in the matrix
75    pub const fn rows(&self) -> usize {
76        M::USIZE
77    }
78
79    /// Number of columns in the matrix
80    pub const fn cols(&self) -> usize {
81        N::USIZE
82    }
83}
84
85impl<T: Sized, M: Positive, N: Positive, O> HeapMatrix<T, M, N, O> {
86    pub fn map<F, U>(self, f: F) -> HeapMatrix<U, M, N, O>
87    where
88        F: FnMut(T) -> U,
89    {
90        HeapMatrix::new(
91            self.data
92                .into_vec()
93                .into_iter()
94                .map(f)
95                .collect::<Box<[U]>>(),
96        )
97    }
98}
99
100impl<T: Sized, M: Positive + Mul<N, Output: Positive>, N: Positive, O> HeapMatrix<T, M, N, O> {
101    /// Flatten matrix into an heap array.
102    pub fn into_flat_array(self) -> HeapArray<T, Prod<M, N>> {
103        HeapArray::new(self.data)
104    }
105
106    /// Build matrix from an heap array in column-major order
107    pub fn from_flat_array(value: HeapArray<T, Prod<M, N>>) -> Self {
108        Self::new(value.data)
109    }
110}
111
112// --------------------- Column Major (Fortran-style) ----------------------- //
113
114impl<T: Sized, M: Positive, N: Positive> HeapMatrix<T, M, N, ColumnMajor> {
115    /// Matrix column iterator
116    pub fn col_iter(&self) -> impl ExactSizeIterator<Item = &[T]> {
117        self.data.chunks_exact(M::USIZE)
118    }
119
120    /// Matrix column mutable iterator
121    pub fn col_iter_mut(&mut self) -> impl ExactSizeIterator<Item = &mut [T]> {
122        self.data.chunks_exact_mut(M::USIZE)
123    }
124
125    /// Get a reference to an element at position (row, col)
126    pub fn get(&self, row: usize, col: usize) -> Option<&T> {
127        (row < M::USIZE && col < N::USIZE)
128            .then(|| unsafe { self.data.get_unchecked(col * M::USIZE + row) })
129    }
130
131    /// Get a mutable reference to an element at position (row, col)
132    pub fn get_mut(&mut self, row: usize, col: usize) -> Option<&mut T> {
133        (row < M::USIZE && col < N::USIZE)
134            .then(|| unsafe { self.data.get_unchecked_mut(col * M::USIZE + row) })
135    }
136}
137
138impl<T: Sized, M: Positive + Mul<N, Output: Positive>, N: Positive>
139    HeapMatrix<T, M, N, ColumnMajor>
140{
141    /// Build a matrix from an array of columns
142    pub fn from_cols(val: HeapArray<HeapArray<T, M>, N>) -> Self {
143        Self::try_from(val.into_iter().flatten().collect::<Box<[T]>>()).unwrap()
144    }
145}
146
147// --------------------- Row Major (C-style) ----------------------- //
148
149impl<T: Sized, M: Positive, N: Positive> HeapMatrix<T, M, N, RowMajor> {
150    /// Matrix row iterator
151    pub fn row_iter(&self) -> impl ExactSizeIterator<Item = &[T]> {
152        self.data.chunks_exact(N::USIZE)
153    }
154
155    /// Matrix row mutable iterator
156    pub fn row_iter_mut(&mut self) -> impl ExactSizeIterator<Item = &mut [T]> {
157        self.data.chunks_exact_mut(N::USIZE)
158    }
159
160    /// Get a reference to an element at position (row, col)
161    pub fn get(&self, row: usize, col: usize) -> Option<&T> {
162        (row < M::USIZE && col < N::USIZE)
163            .then(|| unsafe { self.data.get_unchecked(row * N::USIZE + col) })
164    }
165
166    /// Get a mutable reference to an element at position (row, col)
167    pub fn get_mut(&mut self, row: usize, col: usize) -> Option<&mut T> {
168        (row < M::USIZE && col < N::USIZE)
169            .then(|| unsafe { self.data.get_unchecked_mut(row * N::USIZE + col) })
170    }
171}
172
173impl<T: Sized, M: Positive + Mul<N, Output: Positive>, N: Positive>
174    HeapMatrix<T, M, N, ColumnMajor>
175{
176    /// Build a matrix from an array of rows
177    pub fn from_rows(val: HeapArray<HeapArray<T, N>, M>) -> Self {
178        Self::try_from(val.into_iter().flatten().collect::<Box<[T]>>()).unwrap()
179    }
180}
181
182// ------------------------ Common Implementations -------------------------- //
183
184impl<T: Sized + Default, M: Positive, N: Positive, O> Default for HeapMatrix<T, M, N, O> {
185    fn default() -> Self {
186        Self::new(
187            (0..M::USIZE * N::USIZE)
188                .map(|_| T::default())
189                .collect::<Box<[T]>>(),
190        )
191    }
192}
193
194impl<T: Sized + Debug, M: Positive, N: Positive, O: Display + Default> Debug
195    for HeapMatrix<T, M, N, O>
196{
197    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
198        f.debug_struct(format!("Matrix[{}]<{}, {}>", O::default(), M::USIZE, N::USIZE).as_str())
199            .field("data", &self.data)
200            .finish()
201    }
202}
203
204impl<T: Sized + Serialize, M: Positive, N: Positive, O> Serialize for HeapMatrix<T, M, N, O> {
205    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
206        self.data.serialize(serializer)
207    }
208}
209
210impl<T: Random + Sized, M: Positive, N: Positive, O: 'static> Random for HeapMatrix<T, M, N, O> {
211    fn random(mut rng: impl CryptoRngCore) -> Self {
212        Self::new(T::random_n(&mut rng, M::USIZE * N::USIZE))
213    }
214}
215
216impl<'de, T: Sized + Deserialize<'de>, M: Positive, N: Positive, O> Deserialize<'de>
217    for HeapMatrix<T, M, N, O>
218{
219    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
220        let data = Box::<[T]>::deserialize(deserializer)?;
221        if data.len() != M::USIZE * N::USIZE {
222            return Err(serde::de::Error::custom(format!(
223                "Expected matrix of length {}, got {}",
224                M::USIZE * N::USIZE,
225                data.len()
226            )));
227        }
228        Ok(Self::new(data))
229    }
230}
231
232impl<T: Sized, M: Positive, N: Positive, O> AsRef<[T]> for HeapMatrix<T, M, N, O> {
233    fn as_ref(&self) -> &[T] {
234        &self.data
235    }
236}
237
238impl<T: Sized, M: Positive, N: Positive, O> AsMut<[T]> for HeapMatrix<T, M, N, O> {
239    fn as_mut(&mut self) -> &mut [T] {
240        &mut self.data
241    }
242}
243
244impl<T: Sized, M: Positive, N: Positive, O> From<HeapMatrix<T, M, N, O>> for Vec<T> {
245    fn from(matrix: HeapMatrix<T, M, N, O>) -> Self {
246        matrix.data.into_vec()
247    }
248}
249
250impl<T: Sized, M: Positive, N: Positive, O> TryFrom<Vec<T>> for HeapMatrix<T, M, N, O> {
251    type Error = PrimitiveError;
252
253    fn try_from(matrix: Vec<T>) -> Result<Self, Self::Error> {
254        if matrix.len() != M::USIZE * N::USIZE {
255            return Err(PrimitiveError::InvalidSize(
256                M::USIZE * N::USIZE,
257                matrix.len(),
258            ));
259        }
260        Ok(Self::new(matrix.into_boxed_slice()))
261    }
262}
263
264impl<T: Sized, M: Positive, N: Positive, O> From<HeapMatrix<T, M, N, O>> for Box<[T]> {
265    fn from(matrix: HeapMatrix<T, M, N, O>) -> Self {
266        matrix.data
267    }
268}
269
270impl<T: Sized, M: Positive, N: Positive, O> TryFrom<Box<[T]>> for HeapMatrix<T, M, N, O> {
271    type Error = PrimitiveError;
272
273    fn try_from(matrix: Box<[T]>) -> Result<Self, Self::Error> {
274        if matrix.len() != M::USIZE * N::USIZE {
275            return Err(PrimitiveError::InvalidSize(
276                M::USIZE * N::USIZE,
277                matrix.len(),
278            ));
279        }
280        Ok(Self::new(matrix))
281    }
282}
283
284#[cfg(test)]
285pub mod tests {
286    use itertools::Itertools;
287    use typenum::{U2, U3, U4, U6};
288
289    use super::{ColumnMajor, RowMajor};
290    use crate::types::{HeapArray, HeapMatrix};
291
292    #[test]
293    fn test_flat_operations() {
294        let data: Vec<u32> = (0..6).collect();
295        let matrix: HeapMatrix<u32, U2, U3> = HeapMatrix::new(data.clone().into_boxed_slice());
296
297        // flat_iter
298        assert_eq!(matrix.flat_iter().copied().collect_vec(), data);
299
300        // flat_iter_mut
301        let mut matrix3: HeapMatrix<u32, U2, U3> = HeapMatrix::new(data.clone().into_boxed_slice());
302        for x in matrix3.flat_iter_mut() {
303            *x *= 2;
304        }
305        assert_eq!(matrix3.get(0, 0), Some(&0));
306        assert_eq!(matrix3.get(1, 0), Some(&2));
307
308        // into_flat_array / from_flat_array
309        let array: HeapArray<u32, U6> = matrix.into_flat_array();
310        assert_eq!(array.as_ref(), data.as_slice());
311        let matrix4: HeapMatrix<u32, U2, U3> = HeapMatrix::from_flat_array(array);
312        assert_eq!(matrix4.as_ref(), data.as_slice());
313    }
314
315    #[test]
316    fn test_column_first_operations() {
317        // 3x2 matrix in column-major: col0=[0,1,2], col1=[3,4,5]
318        let data: Vec<u32> = vec![0, 1, 2, 3, 4, 5];
319        let mut matrix: HeapMatrix<u32, U3, U2, ColumnMajor> =
320            HeapMatrix::new(data.into_boxed_slice());
321
322        // get - column-major layout
323        assert_eq!(matrix.get(0, 0), Some(&0));
324        assert_eq!(matrix.get(1, 0), Some(&1));
325        assert_eq!(matrix.get(2, 0), Some(&2));
326        assert_eq!(matrix.get(0, 1), Some(&3));
327        assert_eq!(matrix.get(1, 1), Some(&4));
328        assert_eq!(matrix.get(2, 1), Some(&5));
329        assert_eq!(matrix.get(3, 0), None);
330        assert_eq!(matrix.get(0, 2), None);
331
332        // get_mut
333        *matrix.get_mut(1, 1).unwrap() = 42;
334        assert_eq!(matrix.get(1, 1), Some(&42));
335
336        // col_iter
337        let data2: Vec<u32> = (0..12).collect();
338        let matrix2: HeapMatrix<u32, U4, U3, ColumnMajor> =
339            HeapMatrix::new(data2.into_boxed_slice());
340        let cols: Vec<Vec<u32>> = matrix2.col_iter().map(|col| col.to_vec()).collect();
341        assert_eq!(cols.len(), 3);
342        assert_eq!(cols[0], vec![0, 1, 2, 3]);
343        assert_eq!(cols[1], vec![4, 5, 6, 7]);
344        assert_eq!(cols[2], vec![8, 9, 10, 11]);
345
346        // col_iter_mut
347        let data3: Vec<u32> = (0..12).collect();
348        let mut matrix3: HeapMatrix<u32, U4, U3, ColumnMajor> =
349            HeapMatrix::new(data3.into_boxed_slice());
350        for col in matrix3.col_iter_mut() {
351            col[0] = 99;
352        }
353        assert_eq!(matrix3.get(0, 0), Some(&99));
354        assert_eq!(matrix3.get(0, 1), Some(&99));
355        assert_eq!(matrix3.get(0, 2), Some(&99));
356    }
357
358    #[test]
359    fn test_row_first_operations() {
360        // 3x2 matrix in row-major: row0=[0,1], row1=[2,3], row2=[4,5]
361        let data: Vec<u32> = vec![0, 1, 2, 3, 4, 5];
362        let mut matrix: HeapMatrix<u32, U3, U2, RowMajor> =
363            HeapMatrix::new(data.into_boxed_slice());
364
365        // get - row-major layout
366        assert_eq!(matrix.get(0, 0), Some(&0));
367        assert_eq!(matrix.get(0, 1), Some(&1));
368        assert_eq!(matrix.get(1, 0), Some(&2));
369        assert_eq!(matrix.get(1, 1), Some(&3));
370        assert_eq!(matrix.get(2, 0), Some(&4));
371        assert_eq!(matrix.get(2, 1), Some(&5));
372        assert_eq!(matrix.get(3, 0), None);
373        assert_eq!(matrix.get(0, 2), None);
374
375        // get_mut
376        *matrix.get_mut(1, 1).unwrap() = 42;
377        assert_eq!(matrix.get(1, 1), Some(&42));
378
379        // row_iter
380        let data2: Vec<u32> = (0..12).collect();
381        let matrix2: HeapMatrix<u32, U3, U4, RowMajor> = HeapMatrix::new(data2.into_boxed_slice());
382        let rows: Vec<Vec<u32>> = matrix2.row_iter().map(|row| row.to_vec()).collect();
383        assert_eq!(rows.len(), 3);
384        assert_eq!(rows[0], vec![0, 1, 2, 3]);
385        assert_eq!(rows[1], vec![4, 5, 6, 7]);
386        assert_eq!(rows[2], vec![8, 9, 10, 11]);
387
388        // row_iter_mut
389        let data3: Vec<u32> = (0..12).collect();
390        let mut matrix3: HeapMatrix<u32, U3, U4, RowMajor> =
391            HeapMatrix::new(data3.into_boxed_slice());
392        for row in matrix3.row_iter_mut() {
393            row[0] = 99;
394        }
395        assert_eq!(matrix3.get(0, 0), Some(&99));
396        assert_eq!(matrix3.get(1, 0), Some(&99));
397        assert_eq!(matrix3.get(2, 0), Some(&99));
398    }
399
400    #[test]
401    fn test_try_from_vec() {
402        // Success case
403        let data: Vec<u32> = (0..12).collect();
404        let result: Result<HeapMatrix<u32, U3, U4>, _> = data.try_into();
405        assert!(result.is_ok());
406
407        // Failure case - wrong size
408        let data: Vec<u32> = (0..10).collect();
409        let result: Result<HeapMatrix<u32, U3, U4>, _> = data.try_into();
410        assert!(result.is_err());
411    }
412}
primitives/types/heap_array/matrix.rs

primitives/types/heap_array/
matrix.rs
