1use bytemuck::Pod;
2
3#[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 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 pub fn nrows(&self) -> usize {
29 self.data.len() / self.ncols
30 }
31
32 pub fn ncols(&self) -> usize {
34 self.ncols
35 }
36
37 pub fn shape(&self) -> (usize, usize) {
39 (self.nrows(), self.ncols)
40 }
41
42 pub fn row(&self, i: usize) -> &'a [T] {
44 let start = i * self.ncols;
45 &self.data[start..start + self.ncols]
46 }
47
48 pub fn as_flat_slice(&self) -> &'a [T] {
50 self.data
51 }
52
53 pub fn rows(&self) -> impl Iterator<Item = &'a [T]> {
55 self.data.chunks_exact(self.ncols)
56 }
57}
58
59#[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}