Skip to main content

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;