geometry_rtree/
nearest_iter.rs

1//! Unbounded nearest-first streaming — the search behind
2//! [`Rtree::nearest_iter`](crate::rtree::Rtree::nearest_iter).
3//!
4//! A best-first stream over this crate's
5//! `SearchFrontier`: separate node and value frontiers keep each entry
6//! to a distance plus one reference. Each [`next`](Iterator::next)
7//! expands nodes until the closest pending value is no farther than the
8//! closest unexpanded node, then yields that value. No best-k
9//! pruning happens because no `k` exists; the bounded
10//! [`Rtree::nearest`](crate::rtree::Rtree::nearest) keeps the pruned
11//! path for the fixed-k case.
12
13use core::iter::FusedIterator;
14
15use crate::indexable::Indexable;
16use crate::node::Node;
17use crate::search_frontier::SearchFrontier;
18
19#[cfg(test)]
20use crate::search_frontier::FrontierMetrics;
21
22/// Default inline entry capacity of a nearest iterator's node frontier.
23pub const DEFAULT_NODE_INLINE_CAPACITY: usize = 32;
24
25/// Default inline entry capacity of a nearest iterator's value frontier.
26pub const DEFAULT_VALUE_INLINE_CAPACITY: usize = 64;
27
28/// A lazy iterator over ALL values in the tree in exact non-decreasing
29/// distance order from a query point — an unbounded ordered stream.
30///
31/// Created by [`Rtree::nearest_iter`](crate::rtree::Rtree::nearest_iter).
32/// The consumer supplies its own bound via
33/// [`take`](Iterator::take); consuming to exhaustion drains the whole
34/// tree in distance order. The default inline capacities can be
35/// overridden through
36/// [`Rtree::nearest_iter_with_inline_capacities`](crate::rtree::Rtree::nearest_iter_with_inline_capacities)
37/// when caller-specific measurements justify the stack/spill trade-off.
38pub struct NearestIter<
39    'a,
40    T,
41    const NODE_INLINE_CAPACITY: usize = DEFAULT_NODE_INLINE_CAPACITY,
42    const VALUE_INLINE_CAPACITY: usize = DEFAULT_VALUE_INLINE_CAPACITY,
43> {
44    query: [f64; 2],
45    nodes: SearchFrontier<DistanceEntry<'a, Node<T>>, NODE_INLINE_CAPACITY>,
46    values: SearchFrontier<DistanceEntry<'a, T>, VALUE_INLINE_CAPACITY>,
47    #[cfg(test)]
48    branch_expansions: usize,
49    #[cfg(test)]
50    leaf_expansions: usize,
51    #[cfg(test)]
52    pending_nodes: usize,
53    #[cfg(test)]
54    pending_values: usize,
55    #[cfg(test)]
56    node_high_water: usize,
57    #[cfg(test)]
58    value_high_water: usize,
59    #[cfg(test)]
60    combined_high_water: usize,
61}
62
63#[cfg(test)]
64#[derive(Clone, Copy, Debug)]
65pub(crate) struct NearestMetrics {
66    pub(crate) frontier: FrontierMetrics,
67    pub(crate) branch_expansions: usize,
68    pub(crate) leaf_expansions: usize,
69    pub(crate) node_high_water: usize,
70    pub(crate) value_high_water: usize,
71}
72
73impl<'a, T, const NODE_INLINE_CAPACITY: usize, const VALUE_INLINE_CAPACITY: usize>
74    NearestIter<'a, T, NODE_INLINE_CAPACITY, VALUE_INLINE_CAPACITY>
75{
76    pub(crate) fn new(root: &'a Node<T>, query: [f64; 2]) -> Self {
77        let mut nodes = SearchFrontier::new();
78        nodes.push(DistanceEntry {
79            dist: 0.0,
80            item: root,
81        });
82        Self {
83            query,
84            nodes,
85            values: SearchFrontier::new(),
86            #[cfg(test)]
87            branch_expansions: 0,
88            #[cfg(test)]
89            leaf_expansions: 0,
90            #[cfg(test)]
91            pending_nodes: 1,
92            #[cfg(test)]
93            pending_values: 0,
94            #[cfg(test)]
95            node_high_water: 1,
96            #[cfg(test)]
97            value_high_water: 0,
98            #[cfg(test)]
99            combined_high_water: 1,
100        }
101    }
102
103    #[cfg(test)]
104    pub(crate) fn metrics(&self) -> NearestMetrics {
105        let nodes = self.nodes.metrics();
106        let values = self.values.metrics();
107        NearestMetrics {
108            frontier: FrontierMetrics {
109                pushes: nodes.pushes + values.pushes,
110                pops: nodes.pops + values.pops,
111                high_water: self.combined_high_water,
112            },
113            branch_expansions: self.branch_expansions,
114            leaf_expansions: self.leaf_expansions,
115            node_high_water: self.node_high_water,
116            value_high_water: self.value_high_water,
117        }
118    }
119}
120
121impl<'a, T: Indexable, const NODE_INLINE_CAPACITY: usize, const VALUE_INLINE_CAPACITY: usize>
122    Iterator for NearestIter<'a, T, NODE_INLINE_CAPACITY, VALUE_INLINE_CAPACITY>
123{
124    type Item = &'a T;
125
126    /// Correctness of the yield order: a child's box is contained in
127    /// its parent's, so every entry pushed by an expansion is at least
128    /// as far as the node it came from — a value that pops is nearer
129    /// than everything still unexpanded.
130    fn next(&mut self) -> Option<&'a T> {
131        loop {
132            let node_dist = self.nodes.peek().map(DistanceEntry::dist);
133            let value_dist = self.values.peek().map(DistanceEntry::dist);
134            if value_dist
135                .is_some_and(|value| node_dist.is_none_or(|node| value.total_cmp(&node).is_le()))
136            {
137                let value = self
138                    .values
139                    .pop()
140                    .expect("a distance was just read from the value frontier");
141                #[cfg(test)]
142                {
143                    self.pending_values -= 1;
144                }
145                return Some(value.item);
146            }
147
148            let node = self.nodes.pop()?;
149            match node.item {
150                Node::Leaf(values) => {
151                    #[cfg(test)]
152                    {
153                        self.pending_nodes -= 1;
154                        self.leaf_expansions += 1;
155                        self.pending_values += values.len();
156                        self.value_high_water = self.value_high_water.max(self.pending_values);
157                        self.combined_high_water = self
158                            .combined_high_water
159                            .max(self.pending_nodes + self.pending_values);
160                    }
161                    let query = self.query;
162                    self.values.extend(values.iter().map(|value| DistanceEntry {
163                        dist: value.bounds().comparable_min_distance_to(query),
164                        item: value,
165                    }));
166                }
167                Node::Branch(children) => {
168                    #[cfg(test)]
169                    {
170                        self.pending_nodes -= 1;
171                        self.branch_expansions += 1;
172                        self.pending_nodes += children.len();
173                        self.node_high_water = self.node_high_water.max(self.pending_nodes);
174                        self.combined_high_water = self
175                            .combined_high_water
176                            .max(self.pending_nodes + self.pending_values);
177                    }
178                    let query = self.query;
179                    self.nodes
180                        .extend(children.iter().map(|(bounds, child)| DistanceEntry {
181                            dist: bounds.comparable_min_distance_to(query),
182                            item: child,
183                        }));
184                }
185            }
186        }
187    }
188}
189
190impl<T: Indexable, const NODE_INLINE_CAPACITY: usize, const VALUE_INLINE_CAPACITY: usize>
191    FusedIterator for NearestIter<'_, T, NODE_INLINE_CAPACITY, VALUE_INLINE_CAPACITY>
192{
193}
194
195/// A node or value keyed by its squared minimum distance to the query.
196struct DistanceEntry<'a, T> {
197    dist: f64,
198    item: &'a T,
199}
200
201impl<T> DistanceEntry<'_, T> {
202    fn dist(&self) -> f64 {
203        self.dist
204    }
205}
206
207impl<T> PartialEq for DistanceEntry<'_, T> {
208    fn eq(&self, other: &Self) -> bool {
209        self.dist().total_cmp(&other.dist()).is_eq()
210    }
211}
212
213impl<T> Eq for DistanceEntry<'_, T> {}
214
215impl<T> PartialOrd for DistanceEntry<'_, T> {
216    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
217        Some(self.cmp(other))
218    }
219}
220
221impl<T> Ord for DistanceEntry<'_, T> {
222    /// Reversed so the max-first [`SearchFrontier`] pops the SMALLEST
223    /// distance first.
224    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
225        other.dist().total_cmp(&self.dist())
226    }
227}
228
229#[cfg(test)]
230mod tests {
231    use super::{
232        DEFAULT_NODE_INLINE_CAPACITY, DEFAULT_VALUE_INLINE_CAPACITY, DistanceEntry, NearestMetrics,
233    };
234    use crate::{AsymmetricQuadratic, AsymmetricRStarSplit, Bounds, Rtree, SplitParameters};
235    use core::mem::size_of;
236
237    const FIELD: f64 = 50_000.0;
238    const CLUSTER_COUNT: usize = 16;
239    const CLUSTER_RADIUS: f64 = 100.0;
240    const N: usize = 50_000;
241    const Q: usize = 100;
242    const K: usize = 8;
243
244    #[test]
245    fn distance_entry_equality_is_distance_only() {
246        let first_value = 1_u8;
247        let second_value = 2_u8;
248        let first = DistanceEntry {
249            dist: 3.0,
250            item: &first_value,
251        };
252        let same_distance = DistanceEntry {
253            dist: 3.0,
254            item: &second_value,
255        };
256        let farther = DistanceEntry {
257            dist: 4.0,
258            item: &first_value,
259        };
260        assert!(first == same_distance);
261        assert!(first != farther);
262        assert_eq!(*first.item, first_value);
263    }
264
265    struct Lcg {
266        state: u64,
267    }
268
269    impl Lcg {
270        fn new() -> Self {
271            Self {
272                state: 0x9E37_79B9_7F4A_7C15,
273            }
274        }
275
276        #[allow(
277            clippy::cast_precision_loss,
278            reason = "state >> 11 keeps 53 bits, exact in f64"
279        )]
280        fn next_f64(&mut self) -> f64 {
281            self.state = self
282                .state
283                .wrapping_mul(6_364_136_223_846_793_005)
284                .wrapping_add(1_442_695_040_888_963_407);
285            (self.state >> 11) as f64 / (1u64 << 53) as f64
286        }
287    }
288
289    fn uniform(n: usize) -> Vec<[f64; 2]> {
290        let mut lcg = Lcg::new();
291        (0..n)
292            .map(|_| [lcg.next_f64() * FIELD, lcg.next_f64() * FIELD])
293            .collect()
294    }
295
296    fn clustered(n: usize) -> Vec<[f64; 2]> {
297        let mut lcg = Lcg::new();
298        let centers: Vec<[f64; 2]> = (0..CLUSTER_COUNT)
299            .map(|_| [lcg.next_f64() * FIELD, lcg.next_f64() * FIELD])
300            .collect();
301        (0..n)
302            .map(|i| {
303                let center = centers[i % CLUSTER_COUNT];
304                [
305                    center[0] + lcg.next_f64() * 2.0 * CLUSTER_RADIUS - CLUSTER_RADIUS,
306                    center[1] + lcg.next_f64() * 2.0 * CLUSTER_RADIUS - CLUSTER_RADIUS,
307                ]
308            })
309            .collect()
310    }
311
312    fn queries(q: usize) -> Vec<[f64; 2]> {
313        let mut lcg = Lcg::new();
314        (0..q)
315            .map(|_| {
316                let x = lcg.next_f64() * FIELD;
317                lcg.next_f64();
318                let y = lcg.next_f64() * FIELD;
319                [x, y]
320            })
321            .collect()
322    }
323
324    fn measure<Params: SplitParameters>() -> (usize, usize, usize, usize, usize) {
325        let tree: Rtree<(Bounds, u32), Params> = uniform(N)
326            .into_iter()
327            .enumerate()
328            .map(|(i, point)| (Bounds::point(point), u32::try_from(i).expect("N fits u32")))
329            .collect();
330        let mut total_pushes = 0;
331        let mut total_pops = 0;
332        let mut max_high_water = 0;
333        let mut branch_expansions = 0;
334        let mut leaf_expansions = 0;
335        for query in queries(Q) {
336            let mut nearest = tree.nearest_iter(query);
337            assert_eq!(nearest.by_ref().take(K).count(), K);
338            let NearestMetrics {
339                frontier,
340                branch_expansions: branches,
341                leaf_expansions: leaves,
342                ..
343            } = nearest.metrics();
344            total_pushes += frontier.pushes;
345            total_pops += frontier.pops;
346            max_high_water = max_high_water.max(frontier.high_water);
347            branch_expansions += branches;
348            leaf_expansions += leaves;
349        }
350        (
351            total_pushes,
352            total_pops,
353            max_high_water,
354            branch_expansions,
355            leaf_expansions,
356        )
357    }
358
359    #[test]
360    fn asymmetric_fanout_reduces_knn_frontier_work() {
361        let baseline = measure::<crate::Quadratic<32, 9>>();
362        let default = measure::<AsymmetricRStarSplit<6, 2, 12, 4, 4, 4>>();
363        assert!(
364            default.0 < baseline.0 * 7 / 10,
365            "default asymmetric fanout must cut frontier pushes by at least 30%: baseline={baseline:?}, default={default:?}"
366        );
367        assert!(default.2 < baseline.2);
368        eprintln!(
369            "[rtree-asymmetric-fanout] config=branch32_leaf32 expected_high_water=267 observed={baseline:?}"
370        );
371        for (name, observed) in [
372            (
373                "branch6_leaf32",
374                measure::<AsymmetricQuadratic<6, 2, 32, 9>>(),
375            ),
376            ("default_insert6_leaf12_bulk4_4", default),
377            (
378                "branch12_leaf32",
379                measure::<AsymmetricQuadratic<12, 4, 32, 9>>(),
380            ),
381            (
382                "branch16_leaf32",
383                measure::<AsymmetricQuadratic<16, 4, 32, 9>>(),
384            ),
385        ] {
386            eprintln!(
387                "[rtree-asymmetric-fanout] config={name} expected_pushes_below={} observed={observed:?}",
388                baseline.0 * 7 / 10
389            );
390        }
391        for (name, observed) in [
392            (
393                "branch8_leaf6",
394                measure::<AsymmetricQuadratic<8, 3, 6, 2>>(),
395            ),
396            (
397                "branch8_leaf8",
398                measure::<AsymmetricQuadratic<8, 3, 8, 3>>(),
399            ),
400            (
401                "branch8_leaf12",
402                measure::<AsymmetricQuadratic<8, 3, 12, 4>>(),
403            ),
404            (
405                "branch8_leaf16",
406                measure::<AsymmetricQuadratic<8, 3, 16, 4>>(),
407            ),
408            (
409                "branch8_leaf24",
410                measure::<AsymmetricQuadratic<8, 3, 24, 7>>(),
411            ),
412            ("branch8_leaf32", default),
413        ] {
414            eprintln!("[rtree-leaf-fanout] config={name} observed={observed:?}");
415        }
416    }
417
418    #[test]
419    fn records_split_frontier_capacity_distribution() {
420        for (distribution, points) in [("uniform", uniform(N)), ("clustered", clustered(N))] {
421            let tree: Rtree<(Bounds, u32), AsymmetricRStarSplit<8, 3, 32, 9>> = points
422                .into_iter()
423                .enumerate()
424                .map(|(i, point)| (Bounds::point(point), u32::try_from(i).expect("N fits u32")))
425                .collect();
426            let mut node_high_waters = Vec::with_capacity(Q);
427            let mut value_high_waters = Vec::with_capacity(Q);
428            for query in queries(Q) {
429                let mut nearest = tree.nearest_iter(query);
430                assert_eq!(nearest.by_ref().take(K).count(), K);
431                let metrics = nearest.metrics();
432                node_high_waters.push(metrics.node_high_water);
433                value_high_waters.push(metrics.value_high_water);
434            }
435            node_high_waters.sort_unstable();
436            value_high_waters.sort_unstable();
437            let node_spills = Q - node_high_waters
438                .partition_point(|&water| water <= DEFAULT_NODE_INLINE_CAPACITY);
439            let value_spills = Q - value_high_waters
440                .partition_point(|&water| water <= DEFAULT_VALUE_INLINE_CAPACITY);
441            let spills_at = |high_waters: &[usize], capacity| {
442                Q - high_waters.partition_point(|&water| water <= capacity)
443            };
444            let release_entry_bytes = size_of::<f64>() + size_of::<&(Bounds, u32)>();
445            eprintln!(
446                "[rtree-split-frontier] distribution={distribution} expected_node_capacity={DEFAULT_NODE_INLINE_CAPACITY} observed_node_spills={node_spills}/{Q} observed_node_p50={} observed_node_p95={} observed_node_max={} observed_node_spills_64_96_128={}/{}/{} expected_value_capacity={DEFAULT_VALUE_INLINE_CAPACITY} observed_value_spills={value_spills}/{Q} observed_value_p50={} observed_value_p95={} observed_value_max={} observed_value_spills_128_160_192_256={}/{}/{}/{} observed_entry_bytes={} observed_total_inline_bytes={}",
447                node_high_waters[Q / 2],
448                node_high_waters[Q * 95 / 100],
449                node_high_waters[Q - 1],
450                spills_at(&node_high_waters, 64),
451                spills_at(&node_high_waters, 96),
452                spills_at(&node_high_waters, 128),
453                value_high_waters[Q / 2],
454                value_high_waters[Q * 95 / 100],
455                value_high_waters[Q - 1],
456                spills_at(&value_high_waters, 128),
457                spills_at(&value_high_waters, 160),
458                spills_at(&value_high_waters, 192),
459                spills_at(&value_high_waters, 256),
460                release_entry_bytes,
461                (DEFAULT_NODE_INLINE_CAPACITY + DEFAULT_VALUE_INLINE_CAPACITY)
462                    * release_entry_bytes,
463            );
464        }
465    }
466
467    #[test]
468    fn caller_selected_inline_capacities_preserve_distance_order() {
469        let tree: Rtree<(Bounds, u32)> = uniform(2_000)
470            .into_iter()
471            .enumerate()
472            .map(|(i, point)| (Bounds::point(point), u32::try_from(i).expect("N fits u32")))
473            .collect();
474        let query = [12_345.0, 23_456.0];
475        let mut expected: Vec<f64> = tree
476            .nearest_iter(query)
477            .take(64)
478            .map(|(bounds, _)| bounds.comparable_min_distance_to(query))
479            .collect();
480        let actual: Vec<f64> = tree
481            .nearest_iter_with_inline_capacities::<0, 0>(query)
482            .take(64)
483            .map(|(bounds, _)| bounds.comparable_min_distance_to(query))
484            .collect();
485
486        expected.sort_by(f64::total_cmp);
487        assert_eq!(
488            actual, expected,
489            "the all-spill configuration must stay exact"
490        );
491    }
492}
geometry_rtree/nearest_iter.rs

geometry_rtree/
nearest_iter.rs
