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

index_db/
iter.rs
