Skip to main content

u_geometry/
spatial_index.rs

1//! Spatial indexing for efficient geometric queries.
2//!
3//! Provides a simple linear-scan spatial index suitable for small to medium
4//! object counts (<1000). For larger datasets, consumers should implement
5//! tree-based structures (BVH, R-tree, Octree) using the same API patterns.
6//!
7//! # Design
8//!
9//! The spatial index is deliberately simple: a flat `Vec` of (key, AABB) pairs
10//! with linear-scan queries. This outperforms tree-based structures for small N
11//! due to cache locality and zero overhead.
12//!
13//! For larger N (>1000), tree-based indices amortize O(log n) queries, but
14//! for typical bin-packing/nesting instances (10-500 items), linear scan wins.
15//!
16//! # References
17//!
18//! - Ericson (2005), "Real-Time Collision Detection", Ch. 6 (BVH)
19//! - Akenine-Möller et al. (2018), "Real-Time Rendering", Ch. 25.1 (Spatial Indexing)
20
21use crate::primitives::{AABB2, AABB3};
22
23/// A 2D spatial index entry.
24#[derive(Debug, Clone)]
25struct Entry2D {
26    key: usize,
27    bounds: AABB2,
28}
29
30/// A linear-scan 2D spatial index.
31///
32/// Stores (key, AABB2) pairs and answers overlap queries via brute-force scan.
33///
34/// # Complexity
35/// - Insert: O(1) amortized
36/// - Query: O(n)
37/// - Remove: O(n)
38///
39/// Optimal for N < 1000 due to cache locality.
40#[derive(Debug, Clone, Default)]
41pub struct SpatialIndex2D {
42    entries: Vec<Entry2D>,
43}
44
45impl SpatialIndex2D {
46    /// Creates an empty spatial index.
47    pub fn new() -> Self {
48        Self {
49            entries: Vec::new(),
50        }
51    }
52
53    /// Creates with pre-allocated capacity.
54    pub fn with_capacity(capacity: usize) -> Self {
55        Self {
56            entries: Vec::with_capacity(capacity),
57        }
58    }
59
60    /// Returns the number of entries.
61    #[inline]
62    pub fn len(&self) -> usize {
63        self.entries.len()
64    }
65
66    /// Whether the index is empty.
67    #[inline]
68    pub fn is_empty(&self) -> bool {
69        self.entries.is_empty()
70    }
71
72    /// Inserts an entry.
73    ///
74    /// # Complexity
75    /// O(1) amortized
76    pub fn insert(&mut self, key: usize, bounds: AABB2) {
77        self.entries.push(Entry2D { key, bounds });
78    }
79
80    /// Removes all entries with the given key.
81    ///
82    /// # Complexity
83    /// O(n)
84    pub fn remove(&mut self, key: usize) {
85        self.entries.retain(|e| e.key != key);
86    }
87
88    /// Returns all keys whose AABB overlaps the query region.
89    ///
90    /// # Complexity
91    /// O(n)
92    pub fn query(&self, region: &AABB2) -> Vec<usize> {
93        self.entries
94            .iter()
95            .filter(|e| e.bounds.intersects(region))
96            .map(|e| e.key)
97            .collect()
98    }
99
100    /// Returns all keys whose AABB overlaps the query region, excluding a key.
101    ///
102    /// Useful for "find all neighbors except self" queries.
103    ///
104    /// # Complexity
105    /// O(n)
106    pub fn query_except(&self, region: &AABB2, exclude: usize) -> Vec<usize> {
107        self.entries
108            .iter()
109            .filter(|e| e.key != exclude && e.bounds.intersects(region))
110            .map(|e| e.key)
111            .collect()
112    }
113
114    /// Checks if any entry overlaps the query region.
115    ///
116    /// # Complexity
117    /// O(n) worst case, but short-circuits on first hit.
118    pub fn has_collision(&self, region: &AABB2) -> bool {
119        self.entries.iter().any(|e| e.bounds.intersects(region))
120    }
121
122    /// Checks if any entry (except the excluded key) overlaps the query region.
123    pub fn has_collision_except(&self, region: &AABB2, exclude: usize) -> bool {
124        self.entries
125            .iter()
126            .any(|e| e.key != exclude && e.bounds.intersects(region))
127    }
128
129    /// Removes all entries.
130    pub fn clear(&mut self) {
131        self.entries.clear();
132    }
133}
134
135/// A 3D spatial index entry.
136#[derive(Debug, Clone)]
137struct Entry3D {
138    key: usize,
139    bounds: AABB3,
140}
141
142/// A linear-scan 3D spatial index.
143///
144/// Stores (key, AABB3) pairs and answers overlap queries via brute-force scan.
145///
146/// # Complexity
147/// - Insert: O(1) amortized
148/// - Query: O(n)
149/// - Remove: O(n)
150///
151/// Optimal for N < 1000 due to cache locality.
152///
153/// # Reference
154/// Ericson (2005), "Real-Time Collision Detection", Ch. 6.1
155#[derive(Debug, Clone, Default)]
156pub struct SpatialIndex3D {
157    entries: Vec<Entry3D>,
158}
159
160impl SpatialIndex3D {
161    /// Creates an empty spatial index.
162    pub fn new() -> Self {
163        Self {
164            entries: Vec::new(),
165        }
166    }
167
168    /// Creates with pre-allocated capacity.
169    pub fn with_capacity(capacity: usize) -> Self {
170        Self {
171            entries: Vec::with_capacity(capacity),
172        }
173    }
174
175    /// Returns the number of entries.
176    #[inline]
177    pub fn len(&self) -> usize {
178        self.entries.len()
179    }
180
181    /// Whether the index is empty.
182    #[inline]
183    pub fn is_empty(&self) -> bool {
184        self.entries.is_empty()
185    }
186
187    /// Inserts an entry.
188    ///
189    /// # Complexity
190    /// O(1) amortized
191    pub fn insert(&mut self, key: usize, bounds: AABB3) {
192        self.entries.push(Entry3D { key, bounds });
193    }
194
195    /// Removes all entries with the given key.
196    ///
197    /// # Complexity
198    /// O(n)
199    pub fn remove(&mut self, key: usize) {
200        self.entries.retain(|e| e.key != key);
201    }
202
203    /// Returns all keys whose AABB overlaps the query region.
204    ///
205    /// # Complexity
206    /// O(n)
207    pub fn query(&self, region: &AABB3) -> Vec<usize> {
208        self.entries
209            .iter()
210            .filter(|e| e.bounds.intersects(region))
211            .map(|e| e.key)
212            .collect()
213    }
214
215    /// Returns all keys whose AABB overlaps the query region, excluding a key.
216    ///
217    /// # Complexity
218    /// O(n)
219    pub fn query_except(&self, region: &AABB3, exclude: usize) -> Vec<usize> {
220        self.entries
221            .iter()
222            .filter(|e| e.key != exclude && e.bounds.intersects(region))
223            .map(|e| e.key)
224            .collect()
225    }
226
227    /// Checks if any entry overlaps the query region.
228    ///
229    /// # Complexity
230    /// O(n) worst case, but short-circuits on first hit.
231    pub fn has_collision(&self, region: &AABB3) -> bool {
232        self.entries.iter().any(|e| e.bounds.intersects(region))
233    }
234
235    /// Checks if any entry (except the excluded key) overlaps the query region.
236    pub fn has_collision_except(&self, region: &AABB3, exclude: usize) -> bool {
237        self.entries
238            .iter()
239            .any(|e| e.key != exclude && e.bounds.intersects(region))
240    }
241
242    /// Returns the AABB for a given key, if found.
243    pub fn get_bounds(&self, key: usize) -> Option<&AABB3> {
244        self.entries
245            .iter()
246            .find(|e| e.key == key)
247            .map(|e| &e.bounds)
248    }
249
250    /// Removes all entries.
251    pub fn clear(&mut self) {
252        self.entries.clear();
253    }
254}
255
256#[cfg(test)]
257mod tests {
258    use super::*;
259
260    // ======================== 2D Spatial Index Tests ========================
261
262    #[test]
263    fn test_2d_insert_and_query() {
264        let mut idx = SpatialIndex2D::new();
265        idx.insert(0, AABB2::new(0.0, 0.0, 10.0, 10.0));
266        idx.insert(1, AABB2::new(20.0, 20.0, 30.0, 30.0));
267        idx.insert(2, AABB2::new(5.0, 5.0, 15.0, 15.0));
268
269        let hits = idx.query(&AABB2::new(8.0, 8.0, 12.0, 12.0));
270        assert!(hits.contains(&0)); // overlaps
271        assert!(hits.contains(&2)); // overlaps
272        assert!(!hits.contains(&1)); // no overlap
273    }
274
275    #[test]
276    fn test_2d_query_except() {
277        let mut idx = SpatialIndex2D::new();
278        idx.insert(0, AABB2::new(0.0, 0.0, 10.0, 10.0));
279        idx.insert(1, AABB2::new(5.0, 5.0, 15.0, 15.0));
280
281        let hits = idx.query_except(&AABB2::new(8.0, 8.0, 12.0, 12.0), 0);
282        assert!(!hits.contains(&0));
283        assert!(hits.contains(&1));
284    }
285
286    #[test]
287    fn test_2d_has_collision() {
288        let mut idx = SpatialIndex2D::new();
289        idx.insert(0, AABB2::new(0.0, 0.0, 10.0, 10.0));
290
291        assert!(idx.has_collision(&AABB2::new(5.0, 5.0, 15.0, 15.0)));
292        assert!(!idx.has_collision(&AABB2::new(20.0, 20.0, 30.0, 30.0)));
293    }
294
295    #[test]
296    fn test_2d_remove() {
297        let mut idx = SpatialIndex2D::new();
298        idx.insert(0, AABB2::new(0.0, 0.0, 10.0, 10.0));
299        idx.insert(1, AABB2::new(5.0, 5.0, 15.0, 15.0));
300
301        assert_eq!(idx.len(), 2);
302        idx.remove(0);
303        assert_eq!(idx.len(), 1);
304        assert!(!idx.has_collision(&AABB2::new(0.0, 0.0, 4.0, 4.0)));
305    }
306
307    #[test]
308    fn test_2d_clear() {
309        let mut idx = SpatialIndex2D::new();
310        idx.insert(0, AABB2::new(0.0, 0.0, 10.0, 10.0));
311        idx.clear();
312        assert!(idx.is_empty());
313    }
314
315    #[test]
316    fn test_2d_empty_query() {
317        let idx = SpatialIndex2D::new();
318        assert!(idx.query(&AABB2::new(0.0, 0.0, 10.0, 10.0)).is_empty());
319        assert!(!idx.has_collision(&AABB2::new(0.0, 0.0, 10.0, 10.0)));
320    }
321
322    // ======================== 3D Spatial Index Tests ========================
323
324    fn box3d(x: f64, y: f64, z: f64, w: f64, d: f64, h: f64) -> AABB3 {
325        AABB3::new(x, y, z, x + w, y + d, z + h)
326    }
327
328    #[test]
329    fn test_3d_insert_and_query() {
330        let mut idx = SpatialIndex3D::new();
331        idx.insert(0, box3d(0.0, 0.0, 0.0, 10.0, 10.0, 10.0));
332        idx.insert(1, box3d(20.0, 20.0, 20.0, 10.0, 10.0, 10.0));
333        idx.insert(2, box3d(5.0, 5.0, 5.0, 10.0, 10.0, 10.0));
334
335        let hits = idx.query(&box3d(8.0, 8.0, 8.0, 4.0, 4.0, 4.0));
336        assert!(hits.contains(&0));
337        assert!(hits.contains(&2));
338        assert!(!hits.contains(&1));
339    }
340
341    #[test]
342    fn test_3d_query_except() {
343        let mut idx = SpatialIndex3D::new();
344        idx.insert(0, box3d(0.0, 0.0, 0.0, 10.0, 10.0, 10.0));
345        idx.insert(1, box3d(5.0, 5.0, 5.0, 10.0, 10.0, 10.0));
346
347        let hits = idx.query_except(&box3d(8.0, 8.0, 8.0, 4.0, 4.0, 4.0), 0);
348        assert!(!hits.contains(&0));
349        assert!(hits.contains(&1));
350    }
351
352    #[test]
353    fn test_3d_has_collision() {
354        let mut idx = SpatialIndex3D::new();
355        idx.insert(0, box3d(0.0, 0.0, 0.0, 10.0, 10.0, 10.0));
356
357        assert!(idx.has_collision(&box3d(5.0, 5.0, 5.0, 10.0, 10.0, 10.0)));
358        assert!(!idx.has_collision(&box3d(20.0, 20.0, 20.0, 10.0, 10.0, 10.0)));
359    }
360
361    #[test]
362    fn test_3d_has_collision_except() {
363        let mut idx = SpatialIndex3D::new();
364        idx.insert(0, box3d(0.0, 0.0, 0.0, 10.0, 10.0, 10.0));
365        idx.insert(1, box3d(5.0, 5.0, 5.0, 10.0, 10.0, 10.0));
366
367        // Query overlaps both, but exclude 0 → only 1 counts
368        assert!(idx.has_collision_except(&box3d(8.0, 8.0, 8.0, 4.0, 4.0, 4.0), 0));
369        // Exclude both matches
370        assert!(!idx.has_collision_except(&box3d(0.0, 0.0, 0.0, 2.0, 2.0, 2.0), 0));
371    }
372
373    #[test]
374    fn test_3d_remove() {
375        let mut idx = SpatialIndex3D::new();
376        idx.insert(0, box3d(0.0, 0.0, 0.0, 10.0, 10.0, 10.0));
377        idx.insert(1, box3d(5.0, 5.0, 5.0, 10.0, 10.0, 10.0));
378
379        assert_eq!(idx.len(), 2);
380        idx.remove(0);
381        assert_eq!(idx.len(), 1);
382        assert!(!idx.has_collision(&box3d(0.0, 0.0, 0.0, 2.0, 2.0, 2.0)));
383    }
384
385    #[test]
386    fn test_3d_get_bounds() {
387        let mut idx = SpatialIndex3D::new();
388        let b = box3d(1.0, 2.0, 3.0, 4.0, 5.0, 6.0);
389        idx.insert(42, b);
390
391        let found = idx.get_bounds(42).unwrap();
392        assert!((found.min.x - 1.0).abs() < 1e-10);
393        assert!((found.max.z - 9.0).abs() < 1e-10);
394        assert!(idx.get_bounds(99).is_none());
395    }
396
397    #[test]
398    fn test_3d_clear() {
399        let mut idx = SpatialIndex3D::new();
400        idx.insert(0, box3d(0.0, 0.0, 0.0, 10.0, 10.0, 10.0));
401        idx.clear();
402        assert!(idx.is_empty());
403    }
404
405    #[test]
406    fn test_3d_with_capacity() {
407        let idx = SpatialIndex3D::with_capacity(100);
408        assert_eq!(idx.len(), 0);
409    }
410
411    #[test]
412    fn test_3d_z_axis_separation() {
413        let mut idx = SpatialIndex3D::new();
414        // Stack at same x,y but different z
415        idx.insert(0, box3d(0.0, 0.0, 0.0, 10.0, 10.0, 10.0));
416        idx.insert(1, box3d(0.0, 0.0, 20.0, 10.0, 10.0, 10.0));
417
418        // Query in the gap between them
419        let hits = idx.query(&box3d(0.0, 0.0, 12.0, 10.0, 10.0, 5.0));
420        assert!(hits.is_empty());
421
422        // Query overlapping the top box
423        let hits = idx.query(&box3d(0.0, 0.0, 18.0, 10.0, 10.0, 5.0));
424        assert_eq!(hits.len(), 1);
425        assert!(hits.contains(&1));
426    }
427}