voxelize/libs/
ndarray.rs

1use std::ops::{Index, IndexMut};
2
3use num::Num;
4
5/// N-dimensional array stored in a 1D array.
6#[derive(Debug, Clone, Default)]
7pub struct Ndarray<T>
8where
9    T: Num + Clone,
10{
11    /// Internal data of a n-dimensional array, represented in 1-dimension.
12    pub data: Vec<T>,
13
14    /// Shape of the n-dimensional array.
15    pub shape: Vec<usize>,
16
17    /// Stride of the n-dimensional array, generated from the shape.
18    pub stride: Vec<usize>,
19}
20
21impl<T> Ndarray<T>
22where
23    T: Num + Clone,
24{
25    /// Create a new n-dimensional array.
26    pub fn new(shape: &[usize], default: T) -> Self {
27        let d = shape.len();
28
29        let mut size = 1;
30        shape.iter().for_each(|x| size *= x);
31
32        let data = vec![default; size];
33
34        let mut stride = vec![0; d];
35
36        let mut s = 1;
37        for i in (0..d).rev() {
38            stride[i] = s;
39            s *= shape[i];
40        }
41
42        Self {
43            data,
44            shape: shape.to_vec(),
45            stride,
46        }
47    }
48
49    /// Obtain the index of the n-dimensional array
50    pub fn index(&self, coords: &[usize]) -> usize {
51        coords
52            .iter()
53            .zip(self.stride.iter())
54            .map(|(a, b)| a * b)
55            .sum()
56    }
57
58    /// Check to see if index is within the n-dimensional array's bounds
59    pub fn contains(&self, coords: &[usize]) -> bool {
60        !coords.iter().zip(self.shape.iter()).any(|(&a, &b)| a >= b)
61    }
62}
63
64impl<T: Num + Clone> Index<&[usize]> for Ndarray<T> {
65    type Output = T;
66
67    fn index(&self, index: &[usize]) -> &Self::Output {
68        &self.data.get(self.index(index)).unwrap()
69    }
70}
71
72impl<T: Num + Clone> IndexMut<&[usize]> for Ndarray<T> {
73    fn index_mut(&mut self, index: &[usize]) -> &mut Self::Output {
74        let index = self.index(index);
75        self.data.get_mut(index).unwrap()
76    }
77}
78
79/// Create a new n-dimensional array.
80pub fn ndarray<T: Num + Clone>(shape: &[usize], default: T) -> Ndarray<T> {
81    Ndarray::new(shape, default)
82}