use hashbrown::HashMap;
use rustsim_geometry::vec2::{self, Vec2};
#[derive(Debug, Clone)]
pub struct UniformGrid2 {
pub cell_size: f64,
inv_cell: f64,
cells: HashMap<(i32, i32), Vec<(u64, Vec2)>>,
}
impl UniformGrid2 {
pub fn new(cell_size: f64) -> Self {
assert!(cell_size > 0.0, "cell_size must be positive");
Self {
cell_size,
inv_cell: 1.0 / cell_size,
cells: HashMap::new(),
}
}
pub fn clear(&mut self) {
self.cells.clear();
}
pub fn insert(&mut self, key: u64, pos: Vec2) {
let cell = self.cell_of(pos);
self.cells.entry(cell).or_default().push((key, pos));
}
pub fn insert_all<I: IntoIterator<Item = (u64, Vec2)>>(&mut self, iter: I) {
for (k, p) in iter {
self.insert(k, p);
}
}
pub fn query_radius(&self, center: Vec2, radius: f64) -> Vec<u64> {
let r2 = radius * radius;
let cells = self.cells_within(center, radius);
let mut out = Vec::new();
for cell in cells {
if let Some(bucket) = self.cells.get(&cell) {
for (k, p) in bucket {
if vec2::norm_squared(vec2::sub(*p, center)) <= r2 {
out.push(*k);
}
}
}
}
out
}
pub fn for_each_within<F: FnMut(u64, Vec2)>(&self, center: Vec2, radius: f64, mut f: F) {
let r2 = radius * radius;
for cell in self.cells_within(center, radius) {
if let Some(bucket) = self.cells.get(&cell) {
for (k, p) in bucket {
if vec2::norm_squared(vec2::sub(*p, center)) <= r2 {
f(*k, *p);
}
}
}
}
}
pub fn len(&self) -> usize {
self.cells.values().map(|v| v.len()).sum()
}
pub fn is_empty(&self) -> bool {
self.cells.values().all(|v| v.is_empty())
}
#[inline]
fn cell_of(&self, p: Vec2) -> (i32, i32) {
(
(p[0] * self.inv_cell).floor() as i32,
(p[1] * self.inv_cell).floor() as i32,
)
}
fn cells_within(&self, center: Vec2, radius: f64) -> Vec<(i32, i32)> {
let (cx, cy) = self.cell_of(center);
let r = (radius * self.inv_cell).ceil() as i32;
let mut out = Vec::with_capacity(((2 * r + 1) * (2 * r + 1)) as usize);
for dx in -r..=r {
for dy in -r..=r {
out.push((cx + dx, cy + dy));
}
}
out
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn query_radius_returns_only_nearby_keys() {
let mut g = UniformGrid2::new(1.0);
g.insert(1, [0.0, 0.0]);
g.insert(2, [0.5, 0.5]);
g.insert(3, [10.0, 10.0]);
let mut hits = g.query_radius([0.0, 0.0], 1.0);
hits.sort();
assert_eq!(hits, vec![1, 2]);
}
#[test]
fn negative_coords_bucket_correctly() {
let mut g = UniformGrid2::new(1.0);
g.insert(1, [-0.5, -0.5]);
g.insert(2, [-2.0, -2.0]);
let hits = g.query_radius([-0.5, -0.5], 0.1);
assert_eq!(hits, vec![1]);
}
#[test]
fn len_and_clear() {
let mut g = UniformGrid2::new(2.0);
g.insert_all([(1, [0.0, 0.0]), (2, [5.0, 5.0])]);
assert_eq!(g.len(), 2);
g.clear();
assert!(g.is_empty());
}
}