rustsim_broadphase/
uniform_grid_2d.rs1use hashbrown::HashMap;
4use rustsim_geometry::vec2::{self, Vec2};
5
6#[derive(Debug, Clone)]
11pub struct UniformGrid2 {
12 pub cell_size: f64,
14 inv_cell: f64,
15 cells: HashMap<(i32, i32), Vec<(u64, Vec2)>>,
16}
17
18impl UniformGrid2 {
19 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 pub fn clear(&mut self) {
34 self.cells.clear();
35 }
36
37 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 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 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 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 pub fn len(&self) -> usize {
83 self.cells.values().map(|v| v.len()).sum()
84 }
85
86 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}