index_db/
iter.rs

1//! In-order iteration over a [`BPlusTree`](crate::BPlusTree).
2//!
3//! A [`Cursor`] is a position in the tree — a path of internal nodes from the
4//! root down to a leaf, plus an index within that leaf. It resolves child ids
5//! through the store as it descends, but holds the resolved node references, so
6//! reading the current entry is a direct array access. Stepping a cursor walks
7//! the entries in key order: within a leaf it moves the index, and at a leaf
8//! boundary it ascends the path to the next subtree and descends again. The
9//! public [`Iter`] holds two cursors, one advancing from the front and one
10//! retreating from the back, so it serves both forward and reverse scans; they
11//! meet in the middle, at which point iteration is done.
12
13use alloc::vec::Vec;
14use core::ops::Bound;
15
16use crate::node::{Internal, Leaf, Node};
17use crate::store::{InMemoryStore, NodeId, NodeStore};
18
19/// The path from the root to a leaf: each internal node on the way down, paired
20/// with the index of the child that was descended into.
21type Path<'a, K> = Vec<(&'a Internal<K>, usize)>;
22
23/// A position in the tree: the store it reads through, the path of internal
24/// nodes descended from the root, and the current leaf and entry index.
25struct Cursor<'a, K, V> {
26    store: &'a InMemoryStore<K, V>,
27    path: Path<'a, K>,
28    leaf: &'a Leaf<K, V>,
29    idx: usize,
30}
31
32impl<'a, K, V> Cursor<'a, K, V> {
33    /// The key/value at the current position.
34    #[inline]
35    fn current(&self) -> (&'a K, &'a V) {
36        (&self.leaf.keys[self.idx], &self.leaf.vals[self.idx])
37    }
38
39    /// The key at the current position.
40    #[inline]
41    fn key(&self) -> &'a K {
42        &self.leaf.keys[self.idx]
43    }
44
45    /// Whether two cursors point at the same entry. Leaves are compared by
46    /// address, so this is exact and needs no ordering on `K`.
47    #[inline]
48    fn same_position(&self, other: &Self) -> bool {
49        core::ptr::eq(self.leaf, other.leaf) && self.idx == other.idx
50    }
51
52    /// A cursor at the first (smallest) entry, or `None` if the tree is empty.
53    fn first(store: &'a InMemoryStore<K, V>, root: NodeId) -> Option<Self> {
54        let mut path = Vec::new();
55        let mut node = store.get(root);
56        loop {
57            match node {
58                Node::Internal(internal) => {
59                    path.push((internal, 0));
60                    node = store.get(internal.children[0]);
61                }
62                Node::Leaf(leaf) => {
63                    return if leaf.keys.is_empty() {
64                        None
65                    } else {
66                        Some(Cursor {
67                            store,
68                            path,
69                            leaf,
70                            idx: 0,
71                        })
72                    };
73                }
74            }
75        }
76    }
77
78    /// A cursor at the last (largest) entry, or `None` if the tree is empty.
79    fn last(store: &'a InMemoryStore<K, V>, root: NodeId) -> Option<Self> {
80        let mut path = Vec::new();
81        let mut node = store.get(root);
82        loop {
83            match node {
84                Node::Internal(internal) => {
85                    let ci = internal.children.len() - 1;
86                    path.push((internal, ci));
87                    node = store.get(internal.children[ci]);
88                }
89                Node::Leaf(leaf) => {
90                    return if leaf.keys.is_empty() {
91                        None
92                    } else {
93                        Some(Cursor {
94                            store,
95                            path,
96                            leaf,
97                            idx: leaf.keys.len() - 1,
98                        })
99                    };
100                }
101            }
102        }
103    }
104
105    /// Advance to the next entry in key order. Returns `false` at the end.
106    fn step_forward(&mut self) -> bool {
107        if self.idx + 1 < self.leaf.keys.len() {
108            self.idx += 1;
109            return true;
110        }
111        loop {
112            let (parent, ci) = match self.path.last() {
113                Some(&frame) => frame,
114                None => return false,
115            };
116            if ci + 1 < parent.children.len() {
117                if let Some(frame) = self.path.last_mut() {
118                    frame.1 = ci + 1;
119                }
120                self.descend_leftmost(self.store.get(parent.children[ci + 1]));
121                return true;
122            }
123            let _popped = self.path.pop();
124        }
125    }
126
127    /// Retreat to the previous entry in key order. Returns `false` at the start.
128    fn step_backward(&mut self) -> bool {
129        if self.idx > 0 {
130            self.idx -= 1;
131            return true;
132        }
133        loop {
134            let (parent, ci) = match self.path.last() {
135                Some(&frame) => frame,
136                None => return false,
137            };
138            if ci > 0 {
139                if let Some(frame) = self.path.last_mut() {
140                    frame.1 = ci - 1;
141                }
142                self.descend_rightmost(self.store.get(parent.children[ci - 1]));
143                return true;
144            }
145            let _popped = self.path.pop();
146        }
147    }
148
149    /// Descend `node`'s leftmost path, pushing frames, and land on its first
150    /// entry. `node` must belong to the tree this cursor walks.
151    fn descend_leftmost(&mut self, node: &'a Node<K, V>) {
152        let mut node = node;
153        loop {
154            match node {
155                Node::Internal(internal) => {
156                    self.path.push((internal, 0));
157                    node = self.store.get(internal.children[0]);
158                }
159                Node::Leaf(leaf) => {
160                    self.leaf = leaf;
161                    self.idx = 0;
162                    return;
163                }
164            }
165        }
166    }
167
168    /// Descend `node`'s rightmost path, pushing frames, and land on its last
169    /// entry. `node` must belong to the tree this cursor walks.
170    fn descend_rightmost(&mut self, node: &'a Node<K, V>) {
171        let mut node = node;
172        loop {
173            match node {
174                Node::Internal(internal) => {
175                    let ci = internal.children.len() - 1;
176                    self.path.push((internal, ci));
177                    node = self.store.get(internal.children[ci]);
178                }
179                Node::Leaf(leaf) => {
180                    self.leaf = leaf;
181                    self.idx = leaf.keys.len() - 1;
182                    return;
183                }
184            }
185        }
186    }
187}
188
189impl<'a, K: Ord, V> Cursor<'a, K, V> {
190    /// A cursor at the first entry satisfying the lower `bound`, or `None`.
191    fn lower_bound(store: &'a InMemoryStore<K, V>, root: NodeId, bound: Bound<&K>) -> Option<Self> {
192        let probe = match bound {
193            Bound::Unbounded => return Self::first(store, root),
194            Bound::Included(b) | Bound::Excluded(b) => b,
195        };
196        let (path, leaf) = descend_to(store, root, probe);
197        if leaf.keys.is_empty() {
198            return None;
199        }
200        let idx = match leaf.keys.binary_search(probe) {
201            Ok(i) => match bound {
202                Bound::Excluded(_) => i + 1,
203                _ => i,
204            },
205            Err(i) => i,
206        };
207        let mut cursor = Cursor {
208            store,
209            path,
210            leaf,
211            idx: leaf.keys.len() - 1,
212        };
213        if idx < leaf.keys.len() {
214            cursor.idx = idx;
215            Some(cursor)
216        } else if cursor.step_forward() {
217            Some(cursor)
218        } else {
219            None
220        }
221    }
222
223    /// A cursor at the last entry satisfying the upper `bound`, or `None`.
224    fn upper_bound(store: &'a InMemoryStore<K, V>, root: NodeId, bound: Bound<&K>) -> Option<Self> {
225        let probe = match bound {
226            Bound::Unbounded => return Self::last(store, root),
227            Bound::Included(b) | Bound::Excluded(b) => b,
228        };
229        let (path, leaf) = descend_to(store, root, probe);
230        if leaf.keys.is_empty() {
231            return None;
232        }
233        let idx = match leaf.keys.binary_search(probe) {
234            Ok(i) => match bound {
235                Bound::Excluded(_) => i.checked_sub(1),
236                _ => Some(i),
237            },
238            Err(i) => i.checked_sub(1),
239        };
240        let mut cursor = Cursor {
241            store,
242            path,
243            leaf,
244            idx: 0,
245        };
246        match idx {
247            Some(idx) => {
248                cursor.idx = idx;
249                Some(cursor)
250            }
251            None if cursor.step_backward() => Some(cursor),
252            None => None,
253        }
254    }
255}
256
257/// Descend from `root` to the leaf whose key range covers `probe`, returning the
258/// path taken and that leaf.
259fn descend_to<'a, K: Ord, V>(
260    store: &'a InMemoryStore<K, V>,
261    root: NodeId,
262    probe: &K,
263) -> (Path<'a, K>, &'a Leaf<K, V>) {
264    let mut path = Vec::new();
265    let mut node = store.get(root);
266    loop {
267        match node {
268            Node::Internal(internal) => {
269                let ci = internal.child_index(probe);
270                path.push((internal, ci));
271                node = store.get(internal.children[ci]);
272            }
273            Node::Leaf(leaf) => return (path, leaf),
274        }
275    }
276}
277
278/// A forward-and-reverse iterator over a [`BPlusTree`](crate::BPlusTree)'s
279/// entries in ascending key order, yielding `(&K, &V)`.
280///
281/// Returned by [`BPlusTree::iter`](crate::BPlusTree::iter) and
282/// [`BPlusTree::range`](crate::BPlusTree::range). It is a
283/// [`DoubleEndedIterator`], so a reverse scan is `iter.rev()` and a range can be
284/// walked from either end.
285///
286/// # Examples
287///
288/// ```
289/// use index_db::BPlusTree;
290///
291/// let mut index = BPlusTree::new();
292/// for k in 0..5_u32 {
293///     index.insert(k, k * k);
294/// }
295///
296/// let forward: Vec<_> = index.iter().map(|(&k, &v)| (k, v)).collect();
297/// assert_eq!(forward, vec![(0, 0), (1, 1), (2, 4), (3, 9), (4, 16)]);
298///
299/// let reverse: Vec<_> = index.iter().rev().map(|(&k, _)| k).collect();
300/// assert_eq!(reverse, vec![4, 3, 2, 1, 0]);
301/// ```
302pub struct Iter<'a, K, V> {
303    span: Option<Span<'a, K, V>>,
304}
305
306/// The not-yet-yielded slice of an iteration: cursors at its first and last
307/// remaining entries. Once they coincide, one element is left.
308struct Span<'a, K, V> {
309    front: Cursor<'a, K, V>,
310    back: Cursor<'a, K, V>,
311}
312
313impl<'a, K, V> Iter<'a, K, V> {
314    /// An iterator over every entry in the tree.
315    pub(crate) fn full(store: &'a InMemoryStore<K, V>, root: NodeId) -> Self {
316        match (Cursor::first(store, root), Cursor::last(store, root)) {
317            (Some(front), Some(back)) => Iter {
318                span: Some(Span { front, back }),
319            },
320            _ => Iter { span: None },
321        }
322    }
323}
324
325impl<'a, K: Ord, V> Iter<'a, K, V> {
326    /// An iterator over the entries whose keys fall within the two bounds.
327    pub(crate) fn range(
328        store: &'a InMemoryStore<K, V>,
329        root: NodeId,
330        lower: Bound<&K>,
331        upper: Bound<&K>,
332    ) -> Self {
333        match (
334            Cursor::lower_bound(store, root, lower),
335            Cursor::upper_bound(store, root, upper),
336        ) {
337            (Some(front), Some(back)) if front.key() <= back.key() => Iter {
338                span: Some(Span { front, back }),
339            },
340            _ => Iter { span: None },
341        }
342    }
343}
344
345impl<'a, K, V> Iterator for Iter<'a, K, V> {
346    type Item = (&'a K, &'a V);
347
348    fn next(&mut self) -> Option<Self::Item> {
349        let span = self.span.as_mut()?;
350        let out = span.front.current();
351        if span.front.same_position(&span.back) {
352            self.span = None;
353        } else {
354            let _advanced = span.front.step_forward();
355        }
356        Some(out)
357    }
358}
359
360impl<K, V> DoubleEndedIterator for Iter<'_, K, V> {
361    fn next_back(&mut self) -> Option<Self::Item> {
362        let span = self.span.as_mut()?;
363        let out = span.back.current();
364        if span.front.same_position(&span.back) {
365            self.span = None;
366        } else {
367            let _retreated = span.back.step_backward();
368        }
369        Some(out)
370    }
371}
index_db/iter.rs

index_db/
iter.rs
