osom_lib_btree/btree/operations/
iterate.rs

1#![allow(clippy::needless_return)]
2
3use osom_lib_arrays::fixed_array::InlineFixedArray;
4use osom_lib_arrays::traits::MutableArray;
5use osom_lib_primitives::kvp::KVP;
6
7use crate::btree::{BTree, BTreeConfig, node_ptr::BTreeNodePtr};
8
9struct Frame<TKey, TValue, TConfig>
10where
11    TKey: Ord,
12    TConfig: BTreeConfig,
13{
14    node: BTreeNodePtr<TKey, TValue, TConfig>,
15    index: usize,
16}
17
18const MAX_TREE_DEPTH: usize = 70;
19
20struct BTreeIterator<TKey, TValue, TConfig>
21where
22    TKey: Ord,
23    TConfig: BTreeConfig,
24{
25    stack: InlineFixedArray<MAX_TREE_DEPTH, Frame<TKey, TValue, TConfig>>,
26}
27
28impl<TKey, TValue, TConfig> BTreeIterator<TKey, TValue, TConfig>
29where
30    TKey: Ord,
31    TConfig: BTreeConfig,
32{
33    pub fn new(tree: &BTree<TKey, TValue, TConfig>) -> Self {
34        const {
35            assert!(
36                TConfig::CHILDREN_COUNT >= 4,
37                "BTreeConfig::CHILDREN_COUNT must be greater or equal to 4"
38            );
39        }
40
41        // With CHILDREN_COUNT at least 4, there is no way that
42        // the tree ever reaches the max depth MAX_TREE_DEPTH.
43        // This means that the iterator takes around ~1kb of
44        // stack space, but avoids recursion entirely.
45
46        let mut iter = Self {
47            stack: InlineFixedArray::new(),
48        };
49        if !tree.data.is_null() {
50            iter.push_left(tree.data);
51        }
52        iter
53    }
54
55    pub fn push_left(&mut self, mut node: BTreeNodePtr<TKey, TValue, TConfig>) {
56        unsafe {
57            loop {
58                self.stack
59                    .try_push(Frame { node, index: 0 })
60                    .expect("[BTreeIterator] failed to push frame to stack");
61
62                let child = node.children_ptr().read();
63
64                if child.is_null() {
65                    break;
66                }
67
68                node = child;
69            }
70        }
71    }
72}
73
74impl<TKey, TValue, TConfig> Iterator for BTreeIterator<TKey, TValue, TConfig>
75where
76    TKey: Ord,
77    TConfig: BTreeConfig,
78{
79    type Item = KVP<*mut TKey, *mut TValue>;
80
81    fn next(&mut self) -> Option<Self::Item> {
82        unsafe {
83            loop {
84                let stack_slice = self.stack.as_mut();
85                let top_frame = stack_slice.last_mut()?;
86
87                let node_len = top_frame.node.len();
88                if top_frame.index < node_len {
89                    let i = top_frame.index;
90                    top_frame.index += 1;
91
92                    let key = top_frame.node.keys_ptr().add(i);
93                    let value = top_frame.node.values_ptr().add(i);
94                    let item = KVP { key, value };
95
96                    if !top_frame.node.is_leaf() {
97                        let next_child = top_frame.node.children_ptr().add(i + 1).read();
98                        self.push_left(next_child);
99                    }
100
101                    return Some(item);
102                }
103
104                self.stack
105                    .try_pop()
106                    .expect("[BTreeIterator] failed to pop frame from stack");
107            }
108        }
109    }
110}
111
112impl<TKey, TValue, TConfig> BTree<TKey, TValue, TConfig>
113where
114    TKey: Ord,
115    TConfig: BTreeConfig,
116{
117    /// Iterates over the tree in ascending order.
118    #[inline]
119    pub fn iter(&self) -> impl Iterator<Item = KVP<&TKey, &TValue>> {
120        BTreeIterator::new(self).map(|kvp| unsafe {
121            let (key, value) = kvp.unpack();
122            KVP {
123                key: key.as_ref_unchecked(),
124                value: value.as_ref_unchecked(),
125            }
126        })
127    }
128
129    /// Iterates over the tree in ascending order.
130    #[inline]
131    pub fn iter_mut(&mut self) -> impl Iterator<Item = KVP<&TKey, &mut TValue>> {
132        BTreeIterator::new(self).map(|kvp| unsafe {
133            let (key, value) = kvp.unpack();
134            KVP {
135                key: key.as_ref_unchecked(),
136                value: value.as_mut_unchecked(),
137            }
138        })
139    }
140}
osom_lib_btree/btree/operations/iterate.rs

osom_lib_btree/btree/operations/
iterate.rs
