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::{DEFAULT_NODE_INLINE_CAPACITY, DEFAULT_VALUE_INLINE_CAPACITY, NearestMetrics};
232    use crate::{AsymmetricQuadratic, AsymmetricRStarSplit, Bounds, Rtree, SplitParameters};
233    use core::mem::size_of;
234
235    const FIELD: f64 = 50_000.0;
236    const CLUSTER_COUNT: usize = 16;
237    const CLUSTER_RADIUS: f64 = 100.0;
238    const N: usize = 50_000;
239    const Q: usize = 100;
240    const K: usize = 8;
241
242    struct Lcg {
243        state: u64,
244    }
245
246    impl Lcg {
247        fn new() -> Self {
248            Self {
249                state: 0x9E37_79B9_7F4A_7C15,
250            }
251        }
252
253        #[allow(
254            clippy::cast_precision_loss,
255            reason = "state >> 11 keeps 53 bits, exact in f64"
256        )]
257        fn next_f64(&mut self) -> f64 {
258            self.state = self
259                .state
260                .wrapping_mul(6_364_136_223_846_793_005)
261                .wrapping_add(1_442_695_040_888_963_407);
262            (self.state >> 11) as f64 / (1u64 << 53) as f64
263        }
264    }
265
266    fn uniform(n: usize) -> Vec<[f64; 2]> {
267        let mut lcg = Lcg::new();
268        (0..n)
269            .map(|_| [lcg.next_f64() * FIELD, lcg.next_f64() * FIELD])
270            .collect()
271    }
272
273    fn clustered(n: usize) -> Vec<[f64; 2]> {
274        let mut lcg = Lcg::new();
275        let centers: Vec<[f64; 2]> = (0..CLUSTER_COUNT)
276            .map(|_| [lcg.next_f64() * FIELD, lcg.next_f64() * FIELD])
277            .collect();
278        (0..n)
279            .map(|i| {
280                let center = centers[i % CLUSTER_COUNT];
281                [
282                    center[0] + lcg.next_f64() * 2.0 * CLUSTER_RADIUS - CLUSTER_RADIUS,
283                    center[1] + lcg.next_f64() * 2.0 * CLUSTER_RADIUS - CLUSTER_RADIUS,
284                ]
285            })
286            .collect()
287    }
288
289    fn queries(q: usize) -> Vec<[f64; 2]> {
290        let mut lcg = Lcg::new();
291        (0..q)
292            .map(|_| {
293                let x = lcg.next_f64() * FIELD;
294                lcg.next_f64();
295                let y = lcg.next_f64() * FIELD;
296                [x, y]
297            })
298            .collect()
299    }
300
301    fn measure<Params: SplitParameters>() -> (usize, usize, usize, usize, usize) {
302        let tree: Rtree<(Bounds, u32), Params> = uniform(N)
303            .into_iter()
304            .enumerate()
305            .map(|(i, point)| (Bounds::point(point), u32::try_from(i).expect("N fits u32")))
306            .collect();
307        let mut total_pushes = 0;
308        let mut total_pops = 0;
309        let mut max_high_water = 0;
310        let mut branch_expansions = 0;
311        let mut leaf_expansions = 0;
312        for query in queries(Q) {
313            let mut nearest = tree.nearest_iter(query);
314            assert_eq!(nearest.by_ref().take(K).count(), K);
315            let NearestMetrics {
316                frontier,
317                branch_expansions: branches,
318                leaf_expansions: leaves,
319                ..
320            } = nearest.metrics();
321            total_pushes += frontier.pushes;
322            total_pops += frontier.pops;
323            max_high_water = max_high_water.max(frontier.high_water);
324            branch_expansions += branches;
325            leaf_expansions += leaves;
326        }
327        (
328            total_pushes,
329            total_pops,
330            max_high_water,
331            branch_expansions,
332            leaf_expansions,
333        )
334    }
335
336    #[test]
337    fn asymmetric_fanout_reduces_knn_frontier_work() {
338        let baseline = measure::<crate::Quadratic<32, 9>>();
339        let default = measure::<AsymmetricRStarSplit<6, 2, 12, 4, 4, 4>>();
340        assert!(
341            default.0 < baseline.0 * 7 / 10,
342            "default asymmetric fanout must cut frontier pushes by at least 30%: baseline={baseline:?}, default={default:?}"
343        );
344        assert!(default.2 < baseline.2);
345        eprintln!(
346            "[rtree-asymmetric-fanout] config=branch32_leaf32 expected_high_water=267 observed={baseline:?}"
347        );
348        for (name, observed) in [
349            (
350                "branch6_leaf32",
351                measure::<AsymmetricQuadratic<6, 2, 32, 9>>(),
352            ),
353            ("default_insert6_leaf12_bulk4_4", default),
354            (
355                "branch12_leaf32",
356                measure::<AsymmetricQuadratic<12, 4, 32, 9>>(),
357            ),
358            (
359                "branch16_leaf32",
360                measure::<AsymmetricQuadratic<16, 4, 32, 9>>(),
361            ),
362        ] {
363            eprintln!(
364                "[rtree-asymmetric-fanout] config={name} expected_pushes_below={} observed={observed:?}",
365                baseline.0 * 7 / 10
366            );
367        }
368        for (name, observed) in [
369            (
370                "branch8_leaf6",
371                measure::<AsymmetricQuadratic<8, 3, 6, 2>>(),
372            ),
373            (
374                "branch8_leaf8",
375                measure::<AsymmetricQuadratic<8, 3, 8, 3>>(),
376            ),
377            (
378                "branch8_leaf12",
379                measure::<AsymmetricQuadratic<8, 3, 12, 4>>(),
380            ),
381            (
382                "branch8_leaf16",
383                measure::<AsymmetricQuadratic<8, 3, 16, 4>>(),
384            ),
385            (
386                "branch8_leaf24",
387                measure::<AsymmetricQuadratic<8, 3, 24, 7>>(),
388            ),
389            ("branch8_leaf32", default),
390        ] {
391            eprintln!("[rtree-leaf-fanout] config={name} observed={observed:?}");
392        }
393    }
394
395    #[test]
396    fn records_split_frontier_capacity_distribution() {
397        for (distribution, points) in [("uniform", uniform(N)), ("clustered", clustered(N))] {
398            let tree: Rtree<(Bounds, u32), AsymmetricRStarSplit<8, 3, 32, 9>> = points
399                .into_iter()
400                .enumerate()
401                .map(|(i, point)| (Bounds::point(point), u32::try_from(i).expect("N fits u32")))
402                .collect();
403            let mut node_high_waters = Vec::with_capacity(Q);
404            let mut value_high_waters = Vec::with_capacity(Q);
405            for query in queries(Q) {
406                let mut nearest = tree.nearest_iter(query);
407                assert_eq!(nearest.by_ref().take(K).count(), K);
408                let metrics = nearest.metrics();
409                node_high_waters.push(metrics.node_high_water);
410                value_high_waters.push(metrics.value_high_water);
411            }
412            node_high_waters.sort_unstable();
413            value_high_waters.sort_unstable();
414            let node_spills = Q - node_high_waters
415                .partition_point(|&water| water <= DEFAULT_NODE_INLINE_CAPACITY);
416            let value_spills = Q - value_high_waters
417                .partition_point(|&water| water <= DEFAULT_VALUE_INLINE_CAPACITY);
418            let spills_at = |high_waters: &[usize], capacity| {
419                Q - high_waters.partition_point(|&water| water <= capacity)
420            };
421            let release_entry_bytes = size_of::<f64>() + size_of::<&(Bounds, u32)>();
422            eprintln!(
423                "[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={}",
424                node_high_waters[Q / 2],
425                node_high_waters[Q * 95 / 100],
426                node_high_waters[Q - 1],
427                spills_at(&node_high_waters, 64),
428                spills_at(&node_high_waters, 96),
429                spills_at(&node_high_waters, 128),
430                value_high_waters[Q / 2],
431                value_high_waters[Q * 95 / 100],
432                value_high_waters[Q - 1],
433                spills_at(&value_high_waters, 128),
434                spills_at(&value_high_waters, 160),
435                spills_at(&value_high_waters, 192),
436                spills_at(&value_high_waters, 256),
437                release_entry_bytes,
438                (DEFAULT_NODE_INLINE_CAPACITY + DEFAULT_VALUE_INLINE_CAPACITY)
439                    * release_entry_bytes,
440            );
441        }
442    }
443
444    #[test]
445    fn caller_selected_inline_capacities_preserve_distance_order() {
446        let tree: Rtree<(Bounds, u32)> = uniform(2_000)
447            .into_iter()
448            .enumerate()
449            .map(|(i, point)| (Bounds::point(point), u32::try_from(i).expect("N fits u32")))
450            .collect();
451        let query = [12_345.0, 23_456.0];
452        let mut expected: Vec<f64> = tree
453            .nearest_iter(query)
454            .take(64)
455            .map(|(bounds, _)| bounds.comparable_min_distance_to(query))
456            .collect();
457        let actual: Vec<f64> = tree
458            .nearest_iter_with_inline_capacities::<0, 0>(query)
459            .take(64)
460            .map(|(bounds, _)| bounds.comparable_min_distance_to(query))
461            .collect();
462
463        expected.sort_by(f64::total_cmp);
464        assert_eq!(
465            actual, expected,
466            "the all-spill configuration must stay exact"
467        );
468    }
469}
geometry_rtree/nearest_iter.rs

geometry_rtree/
nearest_iter.rs
