nodedb_spatial/rtree/
tree.rs

1//! R*-tree public API and core structure.
2
3use nodedb_types::BoundingBox;
4
5use super::node::{Node, NodeKind, RTreeEntry};
6use super::search::NnResult;
7
8/// R*-tree spatial index.
9///
10/// Supports insert, delete, range search (bbox intersection), and
11/// incremental nearest-neighbor queries. Array-backed nodes for cache
12/// friendliness.
13///
14/// References:
15/// - Beckmann et al., "The R*-tree" (1990)
16/// - Hjaltason & Samet, "Distance Browsing in Spatial Databases" (1999)
17pub struct RTree {
18    pub(crate) nodes: Vec<Node>,
19    pub(crate) root: usize,
20    pub(crate) len: usize,
21}
22
23impl RTree {
24    pub fn new() -> Self {
25        Self {
26            nodes: vec![Node::new_leaf()],
27            root: 0,
28            len: 0,
29        }
30    }
31
32    pub fn len(&self) -> usize {
33        self.len
34    }
35
36    pub fn is_empty(&self) -> bool {
37        self.len == 0
38    }
39
40    /// Range search: return all entries whose bbox intersects the query bbox.
41    pub fn search(&self, query: &BoundingBox) -> Vec<&RTreeEntry> {
42        let mut results = Vec::new();
43        super::search::search_node(&self.nodes, self.root, query, &mut results);
44        results
45    }
46
47    /// Range search returning owned entries.
48    pub fn search_owned(&self, query: &BoundingBox) -> Vec<RTreeEntry> {
49        self.search(query).into_iter().cloned().collect()
50    }
51
52    /// Nearest-neighbor search using incremental distance ordering.
53    pub fn nearest(&self, query_lng: f64, query_lat: f64, k: usize) -> Vec<NnResult> {
54        super::search::nearest(
55            &self.nodes,
56            self.root,
57            query_lng,
58            query_lat,
59            k,
60            self.is_empty(),
61        )
62    }
63
64    /// Get all entries (for persistence serialization).
65    pub fn entries(&self) -> Vec<&RTreeEntry> {
66        let mut result = Vec::with_capacity(self.len);
67        collect_entries(&self.nodes, self.root, &mut result);
68        result
69    }
70
71    /// Find parent of a node by traversal from root.
72    pub(crate) fn find_parent(&self, current: usize, target: usize) -> Option<usize> {
73        if let NodeKind::Internal { children } = &self.nodes[current].kind {
74            for child in children {
75                if child.node_idx == target {
76                    return Some(current);
77                }
78                if let Some(p) = self.find_parent(child.node_idx, target) {
79                    return Some(p);
80                }
81            }
82        }
83        None
84    }
85
86    /// Condense root: if root is internal with 1 child, collapse.
87    pub(crate) fn condense_root(&mut self) {
88        if let NodeKind::Internal { children } = &self.nodes[self.root].kind
89            && children.len() == 1
90        {
91            self.root = children[0].node_idx;
92        }
93    }
94}
95
96impl Default for RTree {
97    fn default() -> Self {
98        Self::new()
99    }
100}
101
102fn collect_entries<'a>(nodes: &'a [Node], node_idx: usize, result: &mut Vec<&'a RTreeEntry>) {
103    match &nodes[node_idx].kind {
104        NodeKind::Leaf { entries } => result.extend(entries.iter()),
105        NodeKind::Internal { children } => {
106            for child in children {
107                collect_entries(nodes, child.node_idx, result);
108            }
109        }
110    }
111}
112
113pub(crate) fn collect_entries_owned(nodes: &[Node], node_idx: usize, result: &mut Vec<RTreeEntry>) {
114    match &nodes[node_idx].kind {
115        NodeKind::Leaf { entries } => result.extend(entries.iter().cloned()),
116        NodeKind::Internal { children } => {
117            for child in children {
118                collect_entries_owned(nodes, child.node_idx, result);
119            }
120        }
121    }
122}
123
124#[cfg(test)]
125mod tests {
126    use super::super::node::LEAF_CAPACITY;
127    use super::*;
128
129    fn make_entry(id: u64, lng: f64, lat: f64) -> RTreeEntry {
130        RTreeEntry {
131            id,
132            bbox: BoundingBox::from_point(lng, lat),
133        }
134    }
135
136    fn make_rect(id: u64, min_lng: f64, min_lat: f64, max_lng: f64, max_lat: f64) -> RTreeEntry {
137        RTreeEntry {
138            id,
139            bbox: BoundingBox::new(min_lng, min_lat, max_lng, max_lat),
140        }
141    }
142
143    #[test]
144    fn empty_tree() {
145        let tree = RTree::new();
146        assert!(tree.is_empty());
147        assert!(
148            tree.search(&BoundingBox::new(-180.0, -90.0, 180.0, 90.0))
149                .is_empty()
150        );
151    }
152
153    #[test]
154    fn insert_and_search_single() {
155        let mut tree = RTree::new();
156        tree.insert(make_entry(1, 10.0, 20.0));
157        assert_eq!(tree.len(), 1);
158        assert_eq!(
159            tree.search(&BoundingBox::new(5.0, 15.0, 15.0, 25.0)).len(),
160            1
161        );
162        assert!(
163            tree.search(&BoundingBox::new(50.0, 50.0, 60.0, 60.0))
164                .is_empty()
165        );
166    }
167
168    #[test]
169    fn insert_many_and_search() {
170        let mut tree = RTree::new();
171        for i in 0..200 {
172            tree.insert(make_entry(
173                i,
174                (i as f64) * 0.5 - 50.0,
175                (i as f64) * 0.3 - 30.0,
176            ));
177        }
178        assert_eq!(tree.len(), 200);
179        let all = tree.search(&BoundingBox::new(-180.0, -90.0, 180.0, 90.0));
180        assert_eq!(all.len(), 200);
181    }
182
183    #[test]
184    fn delete_entry() {
185        let mut tree = RTree::new();
186        for i in 0..50 {
187            tree.insert(make_entry(i, i as f64, i as f64));
188        }
189        assert!(tree.delete(25));
190        assert_eq!(tree.len(), 49);
191        let all = tree.search(&BoundingBox::new(-1.0, -1.0, 100.0, 100.0));
192        assert!(all.iter().all(|e| e.id != 25));
193        assert!(!tree.delete(999));
194    }
195
196    #[test]
197    fn nearest_neighbor() {
198        let mut tree = RTree::new();
199        tree.insert(make_entry(1, 0.0, 0.0));
200        tree.insert(make_entry(2, 10.0, 10.0));
201        tree.insert(make_entry(3, 5.0, 5.0));
202        let nn = tree.nearest(4.0, 4.0, 2);
203        assert_eq!(nn.len(), 2);
204        assert_eq!(nn[0].entry_id, 3);
205        assert_eq!(nn[1].entry_id, 1);
206    }
207
208    #[test]
209    fn nearest_empty() {
210        assert!(RTree::new().nearest(0.0, 0.0, 5).is_empty());
211    }
212
213    #[test]
214    fn rect_overlap_search() {
215        let mut tree = RTree::new();
216        tree.insert(make_rect(1, 0.0, 0.0, 10.0, 10.0));
217        tree.insert(make_rect(2, 5.0, 5.0, 15.0, 15.0));
218        tree.insert(make_rect(3, 20.0, 20.0, 30.0, 30.0));
219        let results = tree.search(&BoundingBox::new(3.0, 3.0, 7.0, 7.0));
220        assert_eq!(results.len(), 2);
221    }
222
223    #[test]
224    fn stress_insert_delete() {
225        let mut tree = RTree::new();
226        for i in 0..100_u64 {
227            tree.insert(make_entry(i, (i as f64) * 0.5, (i as f64) * 0.3));
228        }
229        for i in (0..100_u64).step_by(2) {
230            assert!(tree.delete(i));
231        }
232        assert_eq!(tree.len(), 50);
233        let all = tree.search(&BoundingBox::new(-180.0, -90.0, 180.0, 90.0));
234        assert!(all.iter().all(|e| e.id % 2 == 1));
235    }
236
237    #[test]
238    fn triggers_node_split() {
239        let mut tree = RTree::new();
240        let count = LEAF_CAPACITY * 3;
241        for i in 0..count as u64 {
242            tree.insert(make_entry(i, (i as f64) * 0.1, (i as f64) * 0.1));
243        }
244        assert_eq!(tree.len(), count);
245        let all = tree.search(&BoundingBox::new(-180.0, -90.0, 180.0, 90.0));
246        assert_eq!(all.len(), count);
247    }
248
249    #[test]
250    fn entries_enumeration() {
251        let mut tree = RTree::new();
252        for i in 0..10 {
253            tree.insert(make_entry(i, i as f64, i as f64));
254        }
255        assert_eq!(tree.entries().len(), 10);
256    }
257}
nodedb_spatial/rtree/tree.rs

nodedb_spatial/rtree/
tree.rs
