subms_treap/range.rs
1//! Sorted-range iteration over a treap. Part of the default path: an ordered
2//! index that cannot answer "everything between these two keys" is a hash map
3//! with extra steps.
4//!
5//! `range(from, to)` yields every `(&K, &V)` whose key falls between
6//! the bounds, in ascending key order. Bounds are either inclusive,
7//! exclusive, or unbounded - mix freely.
8//!
9//! Iteration is stack-based so a deep treap
10//! still walks at the expected `O(log N)` peak stack depth. Each
11//! `next()` is amortised `O(1)`; the full traversal of an N-key range
12//! is `O(N)` plus `O(log T)` to locate the left boundary in a treap
13//! of T total entries.
14//!
15//! The iterator borrows the treap immutably for its lifetime - it is
16//! a "stable iteration over a snapshot" in the sense that the
17//! compile-time borrow ensures no concurrent writer can mutate the
18//! source while the iter is alive. Compose with `TreapSnapshot` from
19//! the `concurrent-reads` feature when readers and writers cross
20//! thread boundaries.
21
22use crate::{NIL, Treap};
23
24/// One end of a range query.
25pub enum RangeBound<'a, K> {
26 Unbounded,
27 Inclusive(&'a K),
28 Exclusive(&'a K),
29}
30
31impl<K: Ord, V> Treap<K, V> {
32 /// Iterate every `(&K, &V)` with `from <= key <= to` (or whichever
33 /// inclusion shape the bounds declare), in ascending key order.
34 pub fn range<'a>(
35 &'a self,
36 from: RangeBound<'a, K>,
37 to: RangeBound<'a, K>,
38 ) -> RangeIter<'a, K, V> {
39 let mut iter = RangeIter {
40 treap: self,
41 stack: Vec::new(),
42 to,
43 };
44 iter.descend_to_lower_bound(self.root, &from);
45 iter
46 }
47}
48
49pub struct RangeIter<'a, K, V> {
50 treap: &'a Treap<K, V>,
51 /// In-order stack: ancestors with the left subtree consumed but
52 /// the node itself not yet emitted.
53 stack: Vec<u32>,
54 to: RangeBound<'a, K>,
55}
56
57impl<'a, K: Ord, V> RangeIter<'a, K, V> {
58 fn descend_to_lower_bound(&mut self, mut idx: u32, from: &RangeBound<'a, K>) {
59 while idx != NIL {
60 let node = &self.treap.nodes[idx as usize];
61 let take_left = match from {
62 RangeBound::Unbounded => true,
63 RangeBound::Inclusive(k) => &*node.key >= *k,
64 RangeBound::Exclusive(k) => &*node.key > *k,
65 };
66 if take_left {
67 self.stack.push(idx);
68 idx = node.left;
69 } else {
70 idx = node.right;
71 }
72 }
73 }
74
75 fn in_upper_bound(&self, key: &K) -> bool {
76 match &self.to {
77 RangeBound::Unbounded => true,
78 RangeBound::Inclusive(k) => key <= k,
79 RangeBound::Exclusive(k) => key < k,
80 }
81 }
82}
83
84impl<'a, K: Ord, V> Iterator for RangeIter<'a, K, V> {
85 type Item = (&'a K, &'a V);
86
87 fn next(&mut self) -> Option<Self::Item> {
88 let idx = self.stack.pop()?;
89 let node = &self.treap.nodes[idx as usize];
90 if !self.in_upper_bound(&node.key) {
91 self.stack.clear();
92 return None;
93 }
94 let mut right = node.right;
95 // Standard in-order: descend left from the right child, push the spine.
96 while right != NIL {
97 self.stack.push(right);
98 right = self.treap.nodes[right as usize].left;
99 }
100 Some((&node.key, &node.value))
101 }
102}
103
104#[cfg(test)]
105#[path = "range_tests.rs"]
106mod tests;