Skip to main content

fpm_rs/
array.rs

1use serde::{Deserialize, Deserializer, Serialize, de::Error as _};
2use std::ops::{Index, IndexMut};
3
4use crate::error::{Error, Result};
5
6/// A small, row-major, contiguous 2-D array used by the computational core.
7#[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}