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

primitives/types/heap_array/
matrix.rs
