spatial_hash_3d/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use cgmath::Vector3;
4use itertools::{ConsTuples, Product};
5use std::fmt::{Debug, Formatter};
6use std::ops::RangeInclusive;
7
8
9type BackingStorage<D> = Vec<D>;
10
11/// Spatial hash data structure. see crate docs for usage.
12pub struct SpatialHashGrid<D: Sized> {
13    dims: Vector3<usize>,
14    cubes: BackingStorage<D>,
15}
16
17#[inline]
18/// Converts position in a spatial hash of given dimensions into an index. If position is not in dimensions, returns None.
19fn pos_to_index(dims: Vector3<usize>, v: Vector3<u32>) -> Option<usize> {
20    let (x, y, z) = (v.x as usize, v.y as usize, v.z as usize);
21    if (x  >= dims[0] ) || (y  >= dims[1] ) || (z  >= dims[2])
22    {
23        return None;
24    }
25    Some(
26        x  +y  * dims[0] + z  * (dims[0] * dims[1]),
27    )
28}
29
30
31impl<D: Sized> SpatialHashGrid<D> {
32    /// x,y,z set the dimentsions, filler is a function that is used to initialize contents.
33    pub fn new<V>(x: usize, y: usize, z: usize, filler: V) -> Self
34    where V:FnMut()->D {
35        let cap = x * y * z;
36        // allocate memory
37        let mut d = Vec::with_capacity(cap);
38        // initialize elements
39        d.resize_with(cap, filler);
40        Self {
41            dims: Vector3::new(x, y, z),
42            cubes: d,
43        }
44    }
45
46    /// Get the size/bounds of the area under spatial hash.
47    #[inline]
48    pub fn size(&self)-> Vector3<usize>{
49        self.dims
50    }
51
52    /// Safely retrieve element by index, will do runtime OOB checks
53    #[inline(always)]
54    pub fn get_mut(&mut self, idx:usize)->Option<&mut D>{
55        self.cubes.get_mut(idx)
56    }
57
58    /// Safely retrieve element by index, will do runtime OOB checks
59    #[inline(always)]
60    pub fn get(&mut self, idx:usize)->Option<& D>{
61        self.cubes.get(idx)
62    }
63
64    #[inline]
65    /// Convert given position into index in this spatial hash grid
66    pub fn pos_to_index(&self, v: Vector3<u32>) -> Option<usize> {
67        pos_to_index(self.dims, v)
68    }
69    ///Iterate over cube indices in given bounds [min, max]
70    #[inline]
71    pub fn iter_cube_indices(&self, min: Vector3<u32>, max: Vector3<u32>) -> BoxIdxIterator {
72        BoxIdxIterator {
73            dims: self.dims,
74            iter: itertools::iproduct!(min.x..=max.x, min.y..=max.y, min.z..=max.z),
75        }
76    }
77
78    ///Iterate over cubes in given bounds [min, max] inside the main cube in read only state
79    #[inline]
80    pub fn iter_cubes(&self, min: Vector3<u32>, max: Vector3<u32>) -> BoxIterator<D> {
81        BoxIterator {
82            data: &self.cubes,
83            iter: self.iter_cube_indices(min, max)
84        }
85    }
86    ///Iterate over cubes in given bounds [min, max] in read and write state.
87    #[inline]
88    pub fn iter_cubes_mut(&mut self, min: Vector3<u32>, max: Vector3<u32>) -> BoxIteratorMut<D> {
89        let inner_iter = self.iter_cube_indices(min, max);
90        BoxIteratorMut {
91            data: &mut self.cubes,
92            iter: inner_iter,
93        }
94    }
95}
96
97impl<D: Debug> Debug for SpatialHashGrid<D> {
98    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
99        writeln!(
100            f,
101            "<SpatialHashGrid {}x{}x{}:",
102            self.dims[0], self.dims[1],self.dims[2]
103        )?;
104        for z in 0..self.dims[2] {
105            writeln!(f, "#Slice z={}:", z)?;
106            for x in 0..self.dims[0] {
107                for y in 0..self.dims[1] {
108                    let v = Vector3::new(x as u32, y as u32, z as u32);
109                    let idx = self.pos_to_index(v).expect("This can not go wrong");
110                    write!(f, "{:?}, ", &self.cubes[idx])?;
111                }
112                writeln!(f)?; // finish row in a slice
113            }
114        }
115        write!(f, ">")
116    }
117}
118
119
120pub struct BoxIdxIterator {
121    dims: Vector3<usize>,
122    #[allow(clippy::type_complexity)]
123    iter: ConsTuples<
124    Product<Product<RangeInclusive<u32>, RangeInclusive<u32>>, RangeInclusive<u32>>>,
125}
126
127
128
129
130impl Iterator for BoxIdxIterator {
131    type Item = (Vector3<u32>, usize);
132
133    #[inline]
134    fn next(&mut self) -> Option<Self::Item> {
135        loop {
136            let (x, y, z) = self.iter.next()?;
137            let c = Vector3 { x, y, z };
138            let idx = match pos_to_index(self.dims, c) {
139                Some(idx) => idx,
140                None => {
141                    continue;
142                }
143            };
144            return Some((c, idx));
145        }
146    }
147    #[inline]
148    fn size_hint(&self) -> (usize, Option<usize>) {
149        let (_min, max) = self.iter.size_hint();
150        (0, max)
151    }
152}
153
154
155pub struct BoxIterator<'a, D: Sized> {
156    data: &'a BackingStorage<D>,
157    iter: BoxIdxIterator
158}
159impl <'a, D: Sized> BoxIterator<'a, D>{
160    ///Morph into a BoxIterator which also returns index of elements it traverses.
161    pub fn with_index(self)->BoxIteratorWithIndex<'a, D>{
162        BoxIteratorWithIndex{
163            data:self.data,
164            iter:self.iter
165        }
166    }
167}
168
169
170impl<'a, D: Sized> Iterator for BoxIterator<'a, D> {
171    type Item = (Vector3<u32>, &'a D);
172
173    #[inline]
174    fn next(&mut self) -> Option<Self::Item> {
175            let (pos, idx) = self.iter.next()?;
176            let d = self.data.get(idx);
177            Some((pos, d.unwrap()))
178    }
179    #[inline]
180    fn size_hint(&self) -> (usize, Option<usize>) {
181        let (_min, max) = self.iter.size_hint();
182        (0, max)
183    }
184}
185
186pub struct BoxIteratorWithIndex<'a, D: Sized> {
187    data: &'a BackingStorage<D>,
188    iter: BoxIdxIterator
189}
190
191
192impl<'a, D: Sized> Iterator for BoxIteratorWithIndex<'a, D> {
193    type Item = (Vector3<u32>, usize, &'a D);
194
195    #[inline]
196    fn next(&mut self) -> Option<Self::Item> {
197        let (pos, idx) = self.iter.next()?;
198        let d = self.data.get(idx);
199        Some((pos, idx, d.unwrap()))
200    }
201    #[inline]
202    fn size_hint(&self) -> (usize, Option<usize>) {
203        let (_min, max) = self.iter.size_hint();
204        (0, max)
205    }
206}
207
208pub struct BoxIteratorMut<'a, D: Sized> {
209    data: &'a mut BackingStorage<D>,
210    iter: BoxIdxIterator,
211}
212
213impl<'a, D: Sized> Iterator for BoxIteratorMut<'a, D> {
214    type Item = (Vector3<u32>, usize, &'a mut D);
215
216    #[inline]
217    fn next(&mut self) -> Option<Self::Item> {
218        let (pos, idx) = self.iter.next()?;
219        // Need unsafe block to return mutable references here.
220        unsafe {
221            //SAFETY: the index position should be valid since pos_to_index can
222            //not return invalid position.
223            let d = self.data.get_unchecked_mut(idx) as *mut D;
224            // SAFETY: we know this can not possibly point to anything invalid
225            // We also know that the returned reference should not outlive the iterator
226            // unless user does "something really terrible"(tm)
227            return Some((pos,idx, d.as_mut().unwrap_unchecked()));
228        }
229    }
230    #[inline]
231    fn size_hint(&self) -> (usize, Option<usize>) {
232        let (_min, max) = self.iter.size_hint();
233        (0, max)
234    }
235}
236
237
238impl<D: Sized> std::ops::Index<Vector3<u32>> for SpatialHashGrid<D> {
239    type Output = D;
240    /// Retrieve reference to element by position
241    #[inline]
242    fn index(&self, index: Vector3<u32>) -> &Self::Output {
243        let i = self.pos_to_index(index).expect("Index out of bounds");
244        self.cubes.index(i)
245    }
246}
247
248impl<D: Sized> std::ops::IndexMut<Vector3<u32>> for SpatialHashGrid<D> {
249    /// Retrieve mutable reference to element by position
250    #[inline]
251    fn index_mut(&mut self, index: Vector3<u32>) -> &mut Self::Output {
252        let i = self.pos_to_index(index).expect("Index out of bounds");
253        self.cubes.index_mut(i)
254    }
255}
256
257
258
259impl<D: Sized> std::ops::Index<usize> for SpatialHashGrid<D> {
260    type Output = D;
261    /// Retrieve reference to element by index
262    #[inline]
263    fn index(&self, index: usize) -> &Self::Output {
264        self.cubes.index(index)
265    }
266}
267
268impl<D: Sized> std::ops::IndexMut<usize> for SpatialHashGrid<D> {
269    /// Retrieve mutable reference to element by index
270    #[inline]
271    fn index_mut(&mut self, index:usize) -> &mut Self::Output {
272        self.cubes.index_mut(index)
273    }
274}
275
276#[cfg(test)]
277mod tests {
278    use crate::SpatialHashGrid;
279    use cgmath::Vector3;
280
281    #[derive(Debug, Clone)]
282    pub struct Data {
283        x: u32,
284        y: u32,
285        z: u32,
286    }
287
288    impl Default for Data {
289        fn default() -> Self {
290            Data { x: 0, y: 0, z: 0 }
291        }
292    }
293
294    #[test]
295    fn test_iter_cubes_mut() {
296        let mut sh: SpatialHashGrid<u32> = SpatialHashGrid::new(5, 5, 5, u32::default);
297
298        let sx = 3;
299        let sy = 4;
300        let sz = 5;
301        for (i, j, k) in itertools::iproduct!(0..sx, 0..sy, 0..sz) {
302            let pos = Vector3::new(i, j, k);
303            sh[pos] = 1;
304        }
305
306        let min = Vector3::new(0, 0, 0);
307        let b = Vector3::new(2, 1, 2);
308        let max = Vector3::new(sx, sy, sz);
309
310        for (_p, _i, d) in sh.iter_cubes_mut(min, b) {
311            *d += 1;
312        }
313
314        let mut count = 0;
315        for (_p,  d)  in sh.iter_cubes(min, max) {
316            count += *d;
317        }
318        assert_eq!(
319            count,
320            (sx * sy * sz) + (b.x + 1) * (b.y + 1) * (b.z + 1),
321            "Incorrect sum of values"
322        );
323        let mut count = 0;
324        for (p, idx, d)  in sh.iter_cubes(min, max).with_index() {
325            assert_eq!(sh.pos_to_index(p).expect("position should be valid"), idx);
326            count += *d;
327        }
328        assert_eq!(
329            count,
330            (sx * sy * sz) + (b.x + 1) * (b.y + 1) * (b.z + 1),
331                   "Incorrect sum of values"
332        );
333    }
334    #[test]
335    fn test_indexing() {
336        let mut sh: SpatialHashGrid<Option<u32>> = SpatialHashGrid::new(5, 5, 5, || None);
337
338        let p = Vector3 { x: 3, y: 3, z: 3 };
339        sh[p] = Some(42);
340        assert_eq!(sh[p], Some(42));
341        sh[p] = None;
342        assert_eq!(sh[p], None);
343        println!("{:?}", sh);
344    }
345    #[test]
346    fn test_debug_trait() {
347        let mut sh: SpatialHashGrid<Data> = SpatialHashGrid::new(3, 3, 2, Data::default);
348        let p = Vector3 { x: 1, y: 1, z: 1 };
349        sh[p].x = 42;
350        sh[p].y = 42;
351        sh[p].z = 42;
352        println!("{:?}", sh);
353    }
354}
spatial_hash_3d/lib.rs

spatial_hash_3d/
lib.rs
