nodedb_spatial/rtree/
tree.rs

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

nodedb_spatial/rtree/
tree.rs
