nodedb_spatial/rtree/
search.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! R-tree range search and nearest-neighbor queries.
4
5use nodedb_types::BoundingBox;
6use std::collections::BinaryHeap;
7
8use super::node::{EntryId, Node, NodeKind, RTreeEntry};
9
10/// Result from nearest-neighbor search.
11#[derive(Debug, Clone)]
12pub struct NnResult {
13    pub entry_id: EntryId,
14    pub bbox: BoundingBox,
15    /// Minimum distance in degrees (approximate).
16    pub distance: f64,
17}
18
19/// Recursive range search.
20pub(crate) fn search_node<'a>(
21    nodes: &'a [Node],
22    node_idx: usize,
23    query: &BoundingBox,
24    results: &mut Vec<&'a RTreeEntry>,
25) {
26    let node = &nodes[node_idx];
27    if !node.bbox.intersects(query) {
28        return;
29    }
30    match &node.kind {
31        NodeKind::Leaf { entries } => {
32            for entry in entries {
33                if entry.bbox.intersects(query) {
34                    results.push(entry);
35                }
36            }
37        }
38        NodeKind::Internal { children } => {
39            for child in children {
40                if child.bbox.intersects(query) {
41                    search_node(nodes, child.node_idx, query, results);
42                }
43            }
44        }
45    }
46}
47
48/// Nearest-neighbor search via priority queue (min-heap).
49pub(crate) fn nearest(
50    nodes: &[Node],
51    root: usize,
52    query_lng: f64,
53    query_lat: f64,
54    k: usize,
55    is_empty: bool,
56) -> Vec<NnResult> {
57    if k == 0 || is_empty {
58        return Vec::new();
59    }
60
61    let mut heap: BinaryHeap<HeapItem> = BinaryHeap::new();
62    // no-governor: hot-path kNN search results; bounded by k (query limit, small)
63    let mut results: Vec<NnResult> = Vec::with_capacity(k);
64
65    heap.push(HeapItem {
66        dist: min_dist_point_bbox(query_lng, query_lat, &nodes[root].bbox),
67        node_idx: root,
68    });
69
70    while let Some(item) = heap.pop() {
71        if results.len() >= k && item.dist > results[k - 1].distance {
72            continue;
73        }
74        let node = &nodes[item.node_idx];
75        match &node.kind {
76            NodeKind::Internal { children } => {
77                for child in children {
78                    let d = min_dist_point_bbox(query_lng, query_lat, &child.bbox);
79                    if results.len() < k || d <= results[results.len() - 1].distance {
80                        heap.push(HeapItem {
81                            dist: d,
82                            node_idx: child.node_idx,
83                        });
84                    }
85                }
86            }
87            NodeKind::Leaf { entries } => {
88                for entry in entries {
89                    let d = min_dist_point_bbox(query_lng, query_lat, &entry.bbox);
90                    if results.len() < k || d < results[results.len() - 1].distance {
91                        let nn = NnResult {
92                            entry_id: entry.id,
93                            bbox: entry.bbox,
94                            distance: d,
95                        };
96                        insert_sorted(&mut results, nn, k);
97                    }
98                }
99            }
100        }
101    }
102
103    results
104}
105
106/// Min-heap item for NN traversal.
107#[derive(Debug)]
108struct HeapItem {
109    dist: f64,
110    node_idx: usize,
111}
112
113impl PartialEq for HeapItem {
114    fn eq(&self, other: &Self) -> bool {
115        self.dist == other.dist
116    }
117}
118impl Eq for HeapItem {}
119
120impl PartialOrd for HeapItem {
121    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
122        Some(self.cmp(other))
123    }
124}
125
126impl Ord for HeapItem {
127    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
128        // Reverse for min-heap (BinaryHeap is max-heap).
129        other
130            .dist
131            .partial_cmp(&self.dist)
132            .unwrap_or(std::cmp::Ordering::Equal)
133    }
134}
135
136/// Minimum distance from point to bbox (in degrees, approximate).
137fn min_dist_point_bbox(lng: f64, lat: f64, bbox: &BoundingBox) -> f64 {
138    let dlat = if lat < bbox.min_lat {
139        bbox.min_lat - lat
140    } else if lat > bbox.max_lat {
141        lat - bbox.max_lat
142    } else {
143        0.0
144    };
145
146    let dlng = if bbox.crosses_antimeridian() {
147        if lng >= bbox.min_lng || lng <= bbox.max_lng {
148            0.0
149        } else {
150            (bbox.min_lng - lng).min(lng - bbox.max_lng).max(0.0)
151        }
152    } else if lng < bbox.min_lng {
153        bbox.min_lng - lng
154    } else if lng > bbox.max_lng {
155        lng - bbox.max_lng
156    } else {
157        0.0
158    };
159
160    (dlat * dlat + dlng * dlng).sqrt()
161}
162
163fn insert_sorted(results: &mut Vec<NnResult>, item: NnResult, k: usize) {
164    let pos = results
165        .binary_search_by(|r| {
166            r.distance
167                .partial_cmp(&item.distance)
168                .unwrap_or(std::cmp::Ordering::Equal)
169        })
170        .unwrap_or_else(|pos| pos);
171    results.insert(pos, item);
172    if results.len() > k {
173        results.truncate(k);
174    }
175}
nodedb_spatial/rtree/search.rs

nodedb_spatial/rtree/
search.rs
