1use serde::{Deserialize, Deserializer, Serialize, de::Error as _};
2use std::ops::{Index, IndexMut};
3
4use crate::error::{Error, Result};
5
6#[derive(Clone, Debug, PartialEq, Serialize)]
8pub struct Array2<T> {
9 height: usize,
10 width: usize,
11 data: Vec<T>,
12}
13
14impl<'de, T> Deserialize<'de> for Array2<T>
15where
16 T: Deserialize<'de>,
17{
18 fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
19 where
20 D: Deserializer<'de>,
21 {
22 #[derive(Deserialize)]
23 struct Representation<T> {
24 height: usize,
25 width: usize,
26 data: Vec<T>,
27 }
28
29 let representation = Representation::deserialize(deserializer)?;
30 Self::from_vec(
31 (representation.height, representation.width),
32 representation.data,
33 )
34 .map_err(D::Error::custom)
35 }
36}
37
38impl<T> Array2<T> {
39 pub fn from_vec(shape: (usize, usize), data: Vec<T>) -> Result<Self> {
40 let expected = shape
41 .0
42 .checked_mul(shape.1)
43 .ok_or_else(|| Error::InvalidShape(format!("shape {shape:?} overflows")))?;
44 if shape.0 == 0 || shape.1 == 0 {
45 return Err(Error::InvalidShape(format!(
46 "dimensions must be non-zero, got {shape:?}"
47 )));
48 }
49 if data.len() != expected {
50 return Err(Error::LengthMismatch {
51 actual: data.len(),
52 expected,
53 shape,
54 });
55 }
56 Ok(Self {
57 height: shape.0,
58 width: shape.1,
59 data,
60 })
61 }
62
63 pub fn shape(&self) -> (usize, usize) {
64 (self.height, self.width)
65 }
66
67 pub fn height(&self) -> usize {
68 self.height
69 }
70
71 pub fn width(&self) -> usize {
72 self.width
73 }
74
75 pub fn len(&self) -> usize {
76 self.data.len()
77 }
78
79 pub fn is_empty(&self) -> bool {
80 self.data.is_empty()
81 }
82
83 pub fn as_slice(&self) -> &[T] {
84 &self.data
85 }
86
87 pub fn as_mut_slice(&mut self) -> &mut [T] {
88 &mut self.data
89 }
90
91 pub fn into_vec(self) -> Vec<T> {
92 self.data
93 }
94
95 pub fn get(&self, row: usize, col: usize) -> Option<&T> {
96 if row < self.height && col < self.width {
97 self.data.get(row * self.width + col)
98 } else {
99 None
100 }
101 }
102
103 pub fn get_mut(&mut self, row: usize, col: usize) -> Option<&mut T> {
104 if row < self.height && col < self.width {
105 self.data.get_mut(row * self.width + col)
106 } else {
107 None
108 }
109 }
110
111 pub fn map<U>(&self, mut function: impl FnMut(&T) -> U) -> Array2<U> {
112 Array2 {
113 height: self.height,
114 width: self.width,
115 data: self.data.iter().map(&mut function).collect(),
116 }
117 }
118}
119
120impl<T: Clone> Array2<T> {
121 pub fn filled(shape: (usize, usize), value: T) -> Result<Self> {
122 let len = shape
123 .0
124 .checked_mul(shape.1)
125 .ok_or_else(|| Error::InvalidShape(format!("shape {shape:?} overflows")))?;
126 Self::from_vec(shape, vec![value; len])
127 }
128}
129
130impl<T: Default + Clone> Array2<T> {
131 pub fn zeros(shape: (usize, usize)) -> Result<Self> {
132 Self::filled(shape, T::default())
133 }
134}
135
136impl<T> Index<(usize, usize)> for Array2<T> {
137 type Output = T;
138
139 fn index(&self, index: (usize, usize)) -> &Self::Output {
140 &self.data[index.0 * self.width + index.1]
141 }
142}
143
144impl<T> IndexMut<(usize, usize)> for Array2<T> {
145 fn index_mut(&mut self, index: (usize, usize)) -> &mut Self::Output {
146 &mut self.data[index.0 * self.width + index.1]
147 }
148}