rustsim_broadphase/
uniform_grid_3d.rs1use hashbrown::HashMap;
4use rustsim_geometry::vec3::{self, Vec3};
5
6#[derive(Debug, Clone)]
8pub struct UniformGrid3 {
9 pub cell_size: f64,
11 inv_cell: f64,
12 cells: HashMap<(i32, i32, i32), Vec<(u64, Vec3)>>,
13}
14
15impl UniformGrid3 {
16 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 pub fn clear(&mut self) {
31 self.cells.clear();
32 }
33
34 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 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 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 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 pub fn len(&self) -> usize {
79 self.cells.values().map(|v| v.len()).sum()
80 }
81
82 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]); g.insert(3, [0.0, 0.0, 5.0]); 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}