rustbatch/entities/detection/
spatial_hash.rs

1use crate::math::rect::Rect;
2use crate::math::vect::Vect;
3use std::collections::HashSet;
4use std::hash::BuildHasherDefault;
5use hashers::fnv::FNV1aHasher32;
6use crate::entities::{FastHash};
7
8use std::hash::Hash;
9
10pub struct SpatialHasher<T: Hash + Eq + Copy + Clone> {
11    pub map: Vec<Vec<HashSet<T, FastHash>>>,
12    tile_size: Vect,
13    w: usize,
14    h: usize,
15}
16
17impl<T: Hash + Eq + Copy + Clone> SpatialHasher<T> {
18    pub fn new(w: usize, h: usize, tile_size: Vect) -> Self {
19        Self {map: vec![vec![HashSet::with_hasher(BuildHasherDefault::<FNV1aHasher32>::default()); w]; h], tile_size, w, h}
20    }
21
22    #[inline]
23    pub fn get_coord(&self, pos: Vect) -> (usize, usize) {
24        (clamp!((pos.x / self.tile_size.x) as i32, 0, self.w as i32-1) as usize,
25         clamp!((pos.y/self.tile_size.y) as i32, 0, self.h as i32 -1) as usize)
26    }
27
28    #[inline]
29    pub fn insert(&mut self, pos: Vect, id: T) -> (usize, usize) {
30        let (x, y) = self.get_coord(pos);
31        self.map[y][x].insert(id);
32        (x, y)
33    }
34
35    #[inline]
36    pub fn remove(&mut self, pos: (usize, usize), id: T) -> bool {
37        self.map[pos.1][pos.0].remove(&id)
38
39    }
40
41    pub fn slow_remove(&mut self, id: T) -> bool {
42        for row in self.map.iter_mut() {
43            for tile in row.iter_mut() {
44                if tile.remove(&id) {
45                    return true;
46                }
47            }
48        }
49
50        false
51    }
52
53
54    #[inline]
55    pub fn update(&mut self, old: (usize, usize), new: Vect, id: T) -> (usize, usize) {
56        let new = self.get_coord(new);
57
58        if old == new {
59            return new;
60        }
61
62        if !self.map[old.1][old.0].remove(&id) {
63            return (usize::MAX, usize::MAX);
64        }
65
66        self.map[new.1][new.0].insert(id);
67        
68        new
69    }
70
71    #[inline]
72    pub fn query(&self, rect: &Rect, collector: &mut Vec<T>) {
73        let mut min = self.get_coord(rect.min);
74        let mut max = self.get_coord(rect.max);
75        min = (
76            if min.0 == 0 {0} else { (min.0-1).min(self.w)},
77            if min.1 == 0 {0} else {(min.1-1).min(self.h)}
78        );
79        max = ((max.0+2).min(self.w), (max.1+2).min(self.h));
80        for y in min.1..max.1 {
81            for x in min.0..max.0 {
82                collector.extend(&self.map[y][x]);
83            }
84        }
85    }
86
87    #[inline]
88    pub fn query_point(&self, pos: Vect, collector: &mut Vec<T>) {
89        let pos = self.get_coord(pos);
90        let min = (
91            if pos.0 == 0 {0} else { (pos.0-1).min(self.w)},
92            if pos.1 == 0 {0} else {(pos.1-1).min(self.h)}
93        );
94        let max = ((pos.0+2).min(self.w),(pos.1+2).min(self.h));
95
96        for y in min.1..max.1 {
97            for x in min.0..max.0 {
98                collector.extend(&self.map[y][x]);
99            }
100        }
101    }
102
103    pub fn get_shape_count(&self) -> usize {
104        let mut count = 0;
105        for row in self.map.iter() {
106            for tile in row.iter() {
107                count += tile.len();
108            }
109        }
110
111        count
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use crate::math::vect::Vect;
118    use rand::Rng;
119    use crate::entities::detection::spatial_hash::SpatialHasher;
120
121    #[test]
122    fn insert_test() {
123        let mut map: SpatialHasher<usize> = SpatialHasher::new(10, 10, Vect::new(100f32, 100f32));
124        let mut rng = rand::thread_rng();
125        for i in 0..100 {
126            map.insert(Vect::new(rng.gen::<f32>()*1000f32, rng.gen::<f32>()*1000f32), i);
127        }
128
129        assert_eq!(100 as usize, map.get_shape_count());
130    }
131
132    #[test]
133    fn remove_test() {
134        let mut map: SpatialHasher<usize> = SpatialHasher::new(10, 10, Vect::new(100f32, 100f32));
135        let mut rng = rand::thread_rng();
136        let mut poss = Vec::new();
137        let mut addresses = Vec::new();
138        for _ in 0..100 {
139            poss.push(Vect::new(rng.gen::<f32>()*1000f32, rng.gen::<f32>()*1000f32));
140        }
141
142        for i in 0..100 {
143            addresses.push(map.insert(poss[i], i));
144        }
145
146        for i in 0..100 {
147            map.remove(addresses[i], i);
148        }
149
150        assert_eq!(0 as usize, map.get_shape_count());
151    }
152}