Skip to main content

rustsim_broadphase/
uniform_grid_3d.rs

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