dotloom_geometry/
spatial.rs

1//! Spatial index over bounding boxes (R*-tree via `rstar`).
2//!
3//! Pointer hit-testing and snapping query only nearby candidates instead of
4//! scanning every entity on each pointer move.
5
6use std::collections::BTreeMap;
7
8use rstar::{AABB, PointDistance, RTree, RTreeObject};
9
10use crate::{Aabb, Point};
11
12#[derive(Debug, Clone, PartialEq)]
13struct Item<K> {
14    key: K,
15    bbox: Aabb,
16}
17
18impl<K> RTreeObject for Item<K> {
19    type Envelope = AABB<[f64; 2]>;
20    fn envelope(&self) -> Self::Envelope {
21        AABB::from_corners([self.bbox.min.x, self.bbox.min.y], [self.bbox.max.x, self.bbox.max.y])
22    }
23}
24
25impl<K> PointDistance for Item<K> {
26    fn distance_2(&self, p: &[f64; 2]) -> f64 {
27        let d = self.bbox.distance_to_point(Point::new(p[0], p[1]));
28        d * d
29    }
30}
31
32/// Bounding-box index keyed by `K`.
33#[derive(Debug, Clone)]
34pub struct SpatialIndex<K: Ord + Copy> {
35    tree: RTree<Item<K>>,
36    boxes: BTreeMap<K, Aabb>,
37}
38
39impl<K: Ord + Copy> Default for SpatialIndex<K> {
40    fn default() -> Self {
41        Self { tree: RTree::new(), boxes: BTreeMap::new() }
42    }
43}
44
45fn usable(b: Aabb) -> bool {
46    !b.is_empty() && b.min.is_finite() && b.max.is_finite()
47}
48
49impl<K: Ord + Copy + PartialEq> SpatialIndex<K> {
50    /// Empty index.
51    #[must_use]
52    pub fn new() -> Self {
53        Self::default()
54    }
55
56    /// Bulk-load (much faster than repeated inserts). Empty boxes are skipped.
57    #[must_use]
58    pub fn bulk_load(items: Vec<(K, Aabb)>) -> Self {
59        let mut boxes = BTreeMap::new();
60        let mut list = Vec::with_capacity(items.len());
61        for (k, b) in items {
62            if usable(b) {
63                boxes.insert(k, b);
64            }
65        }
66        for (k, b) in &boxes {
67            list.push(Item { key: *k, bbox: *b });
68        }
69        Self { tree: RTree::bulk_load(list), boxes }
70    }
71
72    /// Number of indexed keys.
73    #[must_use]
74    pub fn len(&self) -> usize {
75        self.boxes.len()
76    }
77
78    /// Whether the index is empty.
79    #[must_use]
80    pub fn is_empty(&self) -> bool {
81        self.boxes.is_empty()
82    }
83
84    /// Insert or replace the box of `key`. Empty boxes remove the key.
85    pub fn upsert(&mut self, key: K, bbox: Aabb) {
86        self.remove(key);
87        if usable(bbox) {
88            self.tree.insert(Item { key, bbox });
89            self.boxes.insert(key, bbox);
90        }
91    }
92
93    /// Remove `key` if present.
94    pub fn remove(&mut self, key: K) {
95        if let Some(b) = self.boxes.remove(&key) {
96            self.tree.remove(&Item { key, bbox: b });
97        }
98    }
99
100    /// Stored box of `key`.
101    #[must_use]
102    pub fn get(&self, key: K) -> Option<Aabb> {
103        self.boxes.get(&key).copied()
104    }
105
106    /// Keys whose boxes intersect `r`, sorted.
107    #[must_use]
108    pub fn query_rect(&self, r: Aabb) -> Vec<K> {
109        if r.is_empty() {
110            return Vec::new();
111        }
112        let env = AABB::from_corners([r.min.x, r.min.y], [r.max.x, r.max.y]);
113        let mut out: Vec<K> = self.tree.locate_in_envelope_intersecting(env).map(|i| i.key).collect();
114        out.sort();
115        out
116    }
117
118    /// Keys whose boxes lie within `radius` of `p`, nearest first.
119    #[must_use]
120    pub fn query_point(&self, p: Point, radius: f64) -> Vec<K> {
121        if !p.is_finite() || radius.is_nan() || radius < 0.0 {
122            return Vec::new();
123        }
124        self.tree
125            .nearest_neighbor_iter_with_distance_2([p.x, p.y])
126            .take_while(|(_, d2)| *d2 <= radius * radius)
127            .map(|(i, _)| i.key)
128            .collect()
129    }
130
131    /// Union of all boxes.
132    #[must_use]
133    pub fn extent(&self) -> Aabb {
134        self.boxes.values().fold(Aabb::EMPTY, |a, b| a.union(*b))
135    }
136}
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141
142    fn bx(x: f64, y: f64) -> Aabb {
143        Aabb::from_corners(Point::new(x, y), Point::new(x + 1.0, y + 1.0))
144    }
145
146    #[test]
147    fn queries_match_bruteforce() {
148        let items: Vec<(u32, Aabb)> =
149            (0..400).map(|i| (i, bx(f64::from(i % 20) * 3.0, f64::from(i / 20) * 3.0))).collect();
150        let mut idx = SpatialIndex::bulk_load(items.clone());
151        let q = Aabb::from_corners(Point::new(10.0, 10.0), Point::new(20.5, 14.0));
152        let mut brute: Vec<u32> = items.iter().filter(|(_, b)| b.intersects(q)).map(|(k, _)| *k).collect();
153        brute.sort();
154        assert_eq!(idx.query_rect(q), brute);
155
156        // Moving key 5 far away removes it from its old neighbourhood.
157        let old = idx.get(5).unwrap();
158        assert!(idx.query_rect(old).contains(&5));
159        idx.upsert(5, bx(1000.0, 1000.0));
160        assert!(!idx.query_rect(old).contains(&5));
161        assert_eq!(idx.query_point(Point::new(1000.5, 1000.5), 0.0), vec![5]);
162        idx.remove(5);
163        assert!(idx.query_point(Point::new(1000.5, 1000.5), 0.0).is_empty());
164        assert_eq!(idx.len(), 399);
165    }
166
167    #[test]
168    fn point_query_sorted_by_distance() {
169        let idx = SpatialIndex::bulk_load(vec![(1u8, bx(0.0, 0.0)), (2, bx(5.0, 0.0)), (3, bx(50.0, 0.0))]);
170        assert_eq!(idx.query_point(Point::new(3.0, 0.5), 3.0), vec![1, 2]);
171        assert!(idx.query_point(Point::new(f64::NAN, 0.0), 3.0).is_empty());
172    }
173
174    #[test]
175    fn empty_boxes_are_not_indexed() {
176        let mut idx = SpatialIndex::new();
177        idx.upsert(1u8, Aabb::EMPTY);
178        assert!(idx.is_empty());
179    }
180}
dotloom_geometry/spatial.rs

dotloom_geometry/
spatial.rs
