Skip to main content

rustsim_broadphase/
uniform_grid_2d.rs

1//! 2-D uniform hash-grid broadphase.
2
3use hashbrown::HashMap;
4use rustsim_geometry::vec2::{self, Vec2};
5
6/// Uniform hash-grid indexing points in 2-D by cell coordinate.
7///
8/// Each key is stored under the cell containing its point. Queries
9/// inspect only the cells within the requested radius.
10#[derive(Debug, Clone)]
11pub struct UniformGrid2 {
12    /// Cell edge length (metres).
13    pub cell_size: f64,
14    inv_cell: f64,
15    cells: HashMap<(i32, i32), Vec<(u64, Vec2)>>,
16}
17
18impl UniformGrid2 {
19    /// Construct a grid with the given cell size (metres).
20    ///
21    /// # Panics
22    /// Panics if `cell_size` is not strictly positive.
23    pub fn new(cell_size: f64) -> Self {
24        assert!(cell_size > 0.0, "cell_size must be positive");
25        Self {
26            cell_size,
27            inv_cell: 1.0 / cell_size,
28            cells: HashMap::new(),
29        }
30    }
31
32    /// Remove every entry.
33    pub fn clear(&mut self) {
34        self.cells.clear();
35    }
36
37    /// Insert `key` at `pos`.
38    pub fn insert(&mut self, key: u64, pos: Vec2) {
39        let cell = self.cell_of(pos);
40        self.cells.entry(cell).or_default().push((key, pos));
41    }
42
43    /// Insert many `(key, pos)` pairs.
44    pub fn insert_all<I: IntoIterator<Item = (u64, Vec2)>>(&mut self, iter: I) {
45        for (k, p) in iter {
46            self.insert(k, p);
47        }
48    }
49
50    /// Return all keys within Euclidean distance `radius` of `center`.
51    pub fn query_radius(&self, center: Vec2, radius: f64) -> Vec<u64> {
52        let r2 = radius * radius;
53        let cells = self.cells_within(center, radius);
54        let mut out = Vec::new();
55        for cell in cells {
56            if let Some(bucket) = self.cells.get(&cell) {
57                for (k, p) in bucket {
58                    if vec2::norm_squared(vec2::sub(*p, center)) <= r2 {
59                        out.push(*k);
60                    }
61                }
62            }
63        }
64        out
65    }
66
67    /// Invoke `f(key, pos)` for every key within `radius` of `center`.
68    pub fn for_each_within<F: FnMut(u64, Vec2)>(&self, center: Vec2, radius: f64, mut f: F) {
69        let r2 = radius * radius;
70        for cell in self.cells_within(center, radius) {
71            if let Some(bucket) = self.cells.get(&cell) {
72                for (k, p) in bucket {
73                    if vec2::norm_squared(vec2::sub(*p, center)) <= r2 {
74                        f(*k, *p);
75                    }
76                }
77            }
78        }
79    }
80
81    /// Number of entries currently stored.
82    pub fn len(&self) -> usize {
83        self.cells.values().map(|v| v.len()).sum()
84    }
85
86    /// True if the grid is empty.
87    pub fn is_empty(&self) -> bool {
88        self.cells.values().all(|v| v.is_empty())
89    }
90
91    #[inline]
92    fn cell_of(&self, p: Vec2) -> (i32, i32) {
93        (
94            (p[0] * self.inv_cell).floor() as i32,
95            (p[1] * self.inv_cell).floor() as i32,
96        )
97    }
98
99    fn cells_within(&self, center: Vec2, radius: f64) -> Vec<(i32, i32)> {
100        let (cx, cy) = self.cell_of(center);
101        let r = (radius * self.inv_cell).ceil() as i32;
102        let mut out = Vec::with_capacity(((2 * r + 1) * (2 * r + 1)) as usize);
103        for dx in -r..=r {
104            for dy in -r..=r {
105                out.push((cx + dx, cy + dy));
106            }
107        }
108        out
109    }
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115
116    #[test]
117    fn query_radius_returns_only_nearby_keys() {
118        let mut g = UniformGrid2::new(1.0);
119        g.insert(1, [0.0, 0.0]);
120        g.insert(2, [0.5, 0.5]);
121        g.insert(3, [10.0, 10.0]);
122        let mut hits = g.query_radius([0.0, 0.0], 1.0);
123        hits.sort();
124        assert_eq!(hits, vec![1, 2]);
125    }
126
127    #[test]
128    fn negative_coords_bucket_correctly() {
129        let mut g = UniformGrid2::new(1.0);
130        g.insert(1, [-0.5, -0.5]);
131        g.insert(2, [-2.0, -2.0]);
132        let hits = g.query_radius([-0.5, -0.5], 0.1);
133        assert_eq!(hits, vec![1]);
134    }
135
136    #[test]
137    fn len_and_clear() {
138        let mut g = UniformGrid2::new(2.0);
139        g.insert_all([(1, [0.0, 0.0]), (2, [5.0, 5.0])]);
140        assert_eq!(g.len(), 2);
141        g.clear();
142        assert!(g.is_empty());
143    }
144}