nodedb_spatial/rtree/
search.rs

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

nodedb_spatial/rtree/
search.rs
