Skip to main content

trx_rs/
typed_view.rs

1use bytemuck::Pod;
2
3/// A 2D typed view over a flat slice, providing row-based access.
4///
5/// This is a lightweight wrapper that doesn't own the data — it borrows
6/// a `&[T]` and interprets it as `nrows × ncols`.
7#[derive(Debug, Clone, Copy)]
8pub struct TypedView2D<'a, T: Pod> {
9    data: &'a [T],
10    ncols: usize,
11}
12
13impl<'a, T: Pod> TypedView2D<'a, T> {
14    /// Create a 2D view over `data` with `ncols` columns.
15    ///
16    /// Panics if `data.len()` is not divisible by `ncols`.
17    pub fn new(data: &'a [T], ncols: usize) -> Self {
18        assert!(
19            ncols > 0 && data.len().is_multiple_of(ncols),
20            "data length {} is not divisible by ncols {}",
21            data.len(),
22            ncols,
23        );
24        Self { data, ncols }
25    }
26
27    /// Number of rows.
28    pub fn nrows(&self) -> usize {
29        self.data.len() / self.ncols
30    }
31
32    /// Number of columns.
33    pub fn ncols(&self) -> usize {
34        self.ncols
35    }
36
37    /// Shape as `(nrows, ncols)`.
38    pub fn shape(&self) -> (usize, usize) {
39        (self.nrows(), self.ncols)
40    }
41
42    /// Access row `i` as a slice of `ncols` elements.
43    pub fn row(&self, i: usize) -> &'a [T] {
44        let start = i * self.ncols;
45        &self.data[start..start + self.ncols]
46    }
47
48    /// The underlying flat slice.
49    pub fn as_flat_slice(&self) -> &'a [T] {
50        self.data
51    }
52
53    /// Iterate over rows.
54    pub fn rows(&self) -> impl Iterator<Item = &'a [T]> {
55        self.data.chunks_exact(self.ncols)
56    }
57}
58
59/// Mutable 2D typed view.
60#[derive(Debug)]
61pub struct TypedView2DMut<'a, T: Pod> {
62    data: &'a mut [T],
63    ncols: usize,
64}
65
66impl<'a, T: Pod> TypedView2DMut<'a, T> {
67    pub fn new(data: &'a mut [T], ncols: usize) -> Self {
68        assert!(
69            ncols > 0 && data.len().is_multiple_of(ncols),
70            "data length {} is not divisible by ncols {}",
71            data.len(),
72            ncols,
73        );
74        Self { data, ncols }
75    }
76
77    pub fn nrows(&self) -> usize {
78        self.data.len() / self.ncols
79    }
80
81    pub fn ncols(&self) -> usize {
82        self.ncols
83    }
84
85    pub fn row_mut(&mut self, i: usize) -> &mut [T] {
86        let start = i * self.ncols;
87        let ncols = self.ncols;
88        &mut self.data[start..start + ncols]
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95
96    #[test]
97    fn typed_view_2d_basics() {
98        let data: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
99        let view = TypedView2D::new(&data, 3);
100
101        assert_eq!(view.shape(), (2, 3));
102        assert_eq!(view.row(0), &[1.0, 2.0, 3.0]);
103        assert_eq!(view.row(1), &[4.0, 5.0, 6.0]);
104    }
105
106    #[test]
107    fn typed_view_2d_rows_iter() {
108        let data: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
109        let view = TypedView2D::new(&data, 2);
110        let rows: Vec<_> = view.rows().collect();
111        assert_eq!(rows.len(), 3);
112        assert_eq!(rows[2], &[5.0, 6.0]);
113    }
114
115    #[test]
116    #[should_panic]
117    fn typed_view_2d_bad_shape() {
118        let data: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0];
119        TypedView2D::new(&data, 3);
120    }
121}