indigo_kernel/
btree_map.rs

1use arrayvec::ArrayVec;
2use std::{mem, sync::Arc};
3
4/// Ordered map based on an in-memory copy-on-write B+ tree
5#[derive(Clone, Debug)]
6pub struct BTreeMap<K, V> {
7    root: Arc<Node<K, V>>,
8    length: usize,
9}
10
11impl<K, V> BTreeMap<K, V> {
12    /// Create a new, empty map
13    #[must_use]
14    pub fn new() -> Self {
15        Self {
16            root: Arc::new(Node::Leaf(NodeLeaf {
17                keys: ArrayVec::new(),
18                values: ArrayVec::new(),
19            })),
20            length: 0,
21        }
22    }
23
24    /// Insert an entry, returning the existing value if present
25    pub fn insert(&mut self, key: K, value: V) -> Option<V>
26    where
27        K: Clone + Ord,
28        V: Clone,
29    {
30        let root = Arc::make_mut(&mut self.root);
31        match root.insert(key, value) {
32            NodeInsertResult::Replaced(previous_value) => Some(previous_value),
33            NodeInsertResult::Inserted => {
34                self.length += 1;
35                None
36            }
37            // Branch and leaf nodes implement insert-with-split by returning their sibling and the
38            // key they'd like inserted in their parent. Only the root node is capable of creating
39            // itself a new parent.
40            NodeInsertResult::Split(parent_key, child_right) => {
41                // Root node becomes the parent to the incoming `child_right` and its former self
42                // (now called `child_left`).
43                let child_left = {
44                    let parent = Node::Branch(NodeBranch {
45                        keys: ArrayVec::new(),
46                        children: ArrayVec::new(),
47                    });
48                    mem::replace(root, parent)
49                };
50
51                let child_left_max = &child_left.keys()[child_left.keys().len() - 1];
52                let child_right_min = &child_right.keys()[0];
53                assert!(child_left_max < child_right_min, "Left child < right child");
54                assert!(*child_left_max < parent_key, "Left child max < parent key");
55                assert!(
56                    child_right.is_branch() || *child_right_min == parent_key,
57                    "Right child min == parent key"
58                );
59
60                let Node::Branch(parent) = root else {
61                    unreachable!()
62                };
63                parent.keys.push(parent_key);
64                parent.children.push(Arc::new(child_left));
65                parent.children.push(Arc::new(child_right));
66
67                self.length += 1;
68                None
69            }
70        }
71    }
72
73    /// Remove an entry, returning the existing value if present
74    ///
75    /// We follow the advice of the [Deletion without rebalancing in multiway search trees][1]
76    /// paper, which suggests "rebalancing on deletion not only is unnecessary but may be harmful."
77    ///
78    /// [1]: https://doi.org/10.1145/2540068
79    pub fn remove(&mut self, key: &K) -> Option<V>
80    where
81        K: Clone + Ord,
82        V: Clone,
83    {
84        let root = Arc::make_mut(&mut self.root);
85        if let Some(value) = root.remove(key) {
86            self.length -= 1;
87            Some(value)
88        } else {
89            None
90        }
91    }
92
93    /// Get a shared reference to the value associated with the given key
94    pub fn get(&self, key: &K) -> Option<&V>
95    where
96        K: Ord,
97    {
98        self.root.get(key)
99    }
100
101    /// Get an exclusive reference to the value associated with the given key
102    pub fn get_mut(&mut self, key: &K) -> Option<&mut V>
103    where
104        K: Clone + Ord,
105        V: Clone,
106    {
107        let root = Arc::make_mut(&mut self.root);
108        root.get_mut(key)
109    }
110
111    /// Check whether the key is present in the map
112    pub fn contains_key(&self, key: &K) -> bool
113    where
114        K: Ord,
115    {
116        self.root.contains_key(key)
117    }
118
119    /// Return the number of entries in the map
120    #[must_use]
121    pub fn len(&self) -> usize {
122        self.length
123    }
124
125    /// Check whether the map has any entries
126    #[must_use]
127    pub fn is_empty(&self) -> bool {
128        self.len() == 0
129    }
130
131    #[expect(dead_code)]
132    fn assert_invariants(&self)
133    where
134        K: Ord,
135    {
136        self.root.assert_invariants("root", 0);
137    }
138}
139
140impl<K, V> Default for BTreeMap<K, V> {
141    fn default() -> Self {
142        Self::new()
143    }
144}
145
146impl<K, V> Extend<(K, V)> for BTreeMap<K, V>
147where
148    K: Clone + Ord,
149    V: Clone,
150{
151    fn extend<I>(&mut self, iter: I)
152    where
153        I: IntoIterator<Item = (K, V)>,
154    {
155        // TODO: Use more efficient bulk-loading algorithm:
156        // https://cs186berkeley.net/notes/note4/#bulk-loading
157        for (key, value) in iter {
158            self.insert(key, value);
159        }
160    }
161}
162
163#[cfg(not(test))]
164const D: usize = 42; // TODO: Choose a real value
165#[cfg(test)]
166const D: usize = 2;
167
168/// Maximum possible height of the tree
169///
170/// I'm making an educated guess that the height will will never exceed 38 based on these sources:
171///
172/// - [Wikipedia > B-tree > Best case and worst case heights][1]
173/// - [WolframAlpha > "limit of (floor(log base 2 of ((n + 1) / 2))) as n approaches 1 trillion"][2]
174///
175/// [1]: https://en.wikipedia.org/wiki/B-tree#Best_case_and_worst_case_heights
176/// [2]: https://www.wolframalpha.com/input?i=limit+of+%28floor%28log+base+2+of+%28%28n+%2B+1%29+%2F+2%29%29%29+as+n+approaches+1+trillion
177const H_MAX: usize = 38;
178const _: () = assert!(D >= 2, "`H_MAX` assumes `D` >= 2");
179
180#[cfg_attr(not(test), expect(clippy::large_enum_variant))]
181#[derive(Clone, Debug)]
182enum Node<K, V> {
183    Branch(NodeBranch<K, V>),
184    Leaf(NodeLeaf<K, V>),
185}
186
187impl<K, V> Node<K, V> {
188    fn insert(&mut self, key: K, value: V) -> NodeInsertResult<K, V>
189    where
190        K: Clone + Ord,
191        V: Clone,
192    {
193        match self {
194            Self::Branch(branch) => branch.insert(key, value),
195            Self::Leaf(leaf) => leaf.insert(key, value),
196        }
197    }
198
199    fn remove(&mut self, key: &K) -> Option<V>
200    where
201        K: Clone + Ord,
202        V: Clone,
203    {
204        match self {
205            Self::Branch(branch) => branch.remove(key),
206            Self::Leaf(leaf) => leaf.remove(key),
207        }
208    }
209
210    fn get(&self, key: &K) -> Option<&V>
211    where
212        K: Ord,
213    {
214        match self {
215            Self::Branch(branch) => branch.get(key),
216            Self::Leaf(leaf) => leaf.get(key),
217        }
218    }
219
220    fn get_mut(&mut self, key: &K) -> Option<&mut V>
221    where
222        K: Clone + Ord,
223        V: Clone,
224    {
225        match self {
226            Self::Branch(branch) => branch.get_mut(key),
227            Self::Leaf(leaf) => leaf.get_mut(key),
228        }
229    }
230
231    fn contains_key(&self, key: &K) -> bool
232    where
233        K: Ord,
234    {
235        match self {
236            Self::Branch(branch) => branch.contains_key(key),
237            Self::Leaf(leaf) => leaf.contains_key(key),
238        }
239    }
240
241    fn keys(&self) -> &[K] {
242        match self {
243            Self::Branch(branch) => branch.keys.as_slice(),
244            Self::Leaf(leaf) => leaf.keys.as_slice(),
245        }
246    }
247
248    fn is_branch(&self) -> bool {
249        matches!(self, Self::Branch(_))
250    }
251
252    #[cfg_attr(not(test), expect(dead_code))]
253    fn is_leaf(&self) -> bool {
254        matches!(self, Self::Leaf(_))
255    }
256
257    fn assert_invariants(&self, path: &str, depth: u8)
258    where
259        K: Ord,
260    {
261        match self {
262            Self::Branch(branch) => branch.assert_invariants(path, depth),
263            Self::Leaf(leaf) => leaf.assert_invariants(path, depth),
264        }
265    }
266}
267
268#[expect(dead_code)]
269struct NodeSearchResult {
270    branches: ArrayVec<usize, H_MAX>,
271    leaf: LeafSearchResult,
272}
273
274#[cfg_attr(not(test), expect(clippy::large_enum_variant))]
275enum NodeInsertResult<K, V> {
276    Replaced(V),
277    Inserted,
278    Split(K, Node<K, V>),
279}
280
281#[derive(Clone, Debug)]
282struct NodeBranch<K, V> {
283    keys: ArrayVec<K, { (2 * D) + 1 }>,
284    children: ArrayVec<Arc<Node<K, V>>, { (2 * D) + 2 }>,
285}
286
287impl<K, V> NodeBranch<K, V> {
288    fn search(&self, key: &K) -> usize
289    where
290        K: Ord,
291    {
292        // TODO: Go back to binary search; it's better once the nodes get larger (e.g. >=4 KiB).
293        self.keys
294            .iter()
295            .position(|k| k > key)
296            .unwrap_or(self.children.len() - 1)
297    }
298
299    fn split(&mut self) -> (K, Self)
300    where
301        K: Clone,
302    {
303        assert!(self.is_overflowing(), "Only overflowing branches are split");
304        let mut sibling = Self {
305            keys: self.keys.drain(D..).collect(),
306            children: self.children.drain(D + 1..).collect(),
307        };
308        let parent_key = sibling.keys.remove(0);
309        assert_eq!(self.keys.len(), D);
310        assert_eq!(self.children.len(), D + 1);
311        assert_eq!(sibling.keys.len(), D);
312        assert_eq!(sibling.children.len(), D + 1);
313        (parent_key, sibling)
314    }
315
316    fn insert(&mut self, key: K, value: V) -> NodeInsertResult<K, V>
317    where
318        K: Clone + Ord,
319        V: Clone,
320    {
321        let index = self.search(&key);
322        let child = Arc::make_mut(&mut self.children[index]);
323        match child.insert(key, value) {
324            NodeInsertResult::Split(key, child) => {
325                self.keys.insert(index, key);
326                self.children.insert(index + 1, Arc::new(child));
327                if self.is_overflowing() {
328                    let (parent_key, sibling) = self.split();
329                    NodeInsertResult::Split(parent_key, Node::Branch(sibling))
330                } else {
331                    NodeInsertResult::Inserted
332                }
333            }
334            no_split => no_split,
335        }
336    }
337
338    fn remove(&mut self, key: &K) -> Option<V>
339    where
340        K: Clone + Ord,
341        V: Clone,
342    {
343        let index = self.search(key);
344        let child = Arc::make_mut(&mut self.children[index]);
345        child.remove(key)
346    }
347
348    fn get(&self, key: &K) -> Option<&V>
349    where
350        K: Ord,
351    {
352        let index = self.search(key);
353        let child = &self.children[index];
354        child.get(key)
355    }
356
357    fn get_mut(&mut self, key: &K) -> Option<&mut V>
358    where
359        K: Clone + Ord,
360        V: Clone,
361    {
362        let index = self.search(key);
363        let child = Arc::make_mut(&mut self.children[index]);
364        child.get_mut(key)
365    }
366
367    fn contains_key(&self, key: &K) -> bool
368    where
369        K: Ord,
370    {
371        let index = self.search(key);
372        let child = &self.children[index];
373        child.contains_key(key)
374    }
375
376    #[expect(dead_code)]
377    fn is_full(&self) -> bool {
378        self.keys.len() + 1 == self.keys.capacity()
379    }
380
381    fn is_overflowing(&self) -> bool {
382        self.keys.is_full()
383    }
384
385    fn assert_invariants(&self, path: &str, depth: u8)
386    where
387        K: Ord,
388    {
389        assert_eq!(
390            self.keys.len() + 1,
391            self.children.len(),
392            "{path}: keys.len() + 1 != children.len()"
393        );
394        assert!(!self.is_overflowing(), "{path}: overflowing");
395        // TODO: Assert all leaves are at the same depth.
396        // TODO: Assert ordering is correct (e.g. keys to the left are less than). Could pass child
397        // bounds and have it check itself.
398        for (i, child) in self.children.iter().enumerate() {
399            child.assert_invariants(
400                &format!(
401                    "{path}->{}{i}",
402                    if child.is_branch() { "branch" } else { "leaf" }
403                ),
404                depth + 1,
405            );
406        }
407    }
408}
409
410#[derive(Clone, Debug)]
411struct NodeLeaf<K, V> {
412    keys: ArrayVec<K, { (2 * D) + 1 }>,
413    values: ArrayVec<V, { (2 * D) + 1 }>,
414}
415
416impl<K, V> NodeLeaf<K, V> {
417    fn search(&self, key: &K) -> LeafSearchResult
418    where
419        K: Ord,
420    {
421        match self.keys.binary_search(key) {
422            Ok(index) => LeafSearchResult::Found(index),
423            Err(index) => LeafSearchResult::Missing(index),
424        }
425    }
426
427    fn split(&mut self) -> (K, Self)
428    where
429        K: Clone,
430    {
431        assert!(self.is_overflowing(), "Only overflowing leaves are split");
432        let self_len = D;
433        let sibling_len = D + 1;
434        let sibling = Self {
435            keys: self.keys.drain(self_len..).collect(),
436            values: self.values.drain(self_len..).collect(),
437        };
438        let parent_key = sibling.keys[0].clone();
439        assert_eq!(self.keys.len(), self_len);
440        assert_eq!(sibling.keys.len(), sibling_len);
441        (parent_key, sibling)
442    }
443
444    fn insert(&mut self, key: K, value: V) -> NodeInsertResult<K, V>
445    where
446        K: Clone + Ord,
447    {
448        match self.search(&key) {
449            LeafSearchResult::Found(index) => {
450                let previous_value = mem::replace(&mut self.values[index], value);
451                NodeInsertResult::Replaced(previous_value)
452            }
453            LeafSearchResult::Missing(index) => {
454                self.keys.insert(index, key);
455                self.values.insert(index, value);
456                if self.is_overflowing() {
457                    let (parent_key, sibling) = self.split();
458                    NodeInsertResult::Split(parent_key, Node::Leaf(sibling))
459                } else {
460                    NodeInsertResult::Inserted
461                }
462            }
463        }
464    }
465
466    fn remove(&mut self, key: &K) -> Option<V>
467    where
468        K: Ord,
469    {
470        match self.search(key) {
471            LeafSearchResult::Found(index) => {
472                self.keys.remove(index);
473                let value = self.values.remove(index);
474                Some(value)
475            }
476            LeafSearchResult::Missing(_) => None,
477        }
478    }
479
480    fn get(&self, key: &K) -> Option<&V>
481    where
482        K: Ord,
483    {
484        match self.search(key) {
485            LeafSearchResult::Found(index) => Some(&self.values[index]),
486            LeafSearchResult::Missing(_) => None,
487        }
488    }
489
490    fn get_mut(&mut self, key: &K) -> Option<&mut V>
491    where
492        K: Ord,
493    {
494        match self.search(key) {
495            LeafSearchResult::Found(index) => Some(&mut self.values[index]),
496            LeafSearchResult::Missing(_) => None,
497        }
498    }
499
500    fn contains_key(&self, key: &K) -> bool
501    where
502        K: Ord,
503    {
504        matches!(self.search(key), LeafSearchResult::Found(_))
505    }
506
507    #[expect(dead_code)]
508    fn is_full(&self) -> bool {
509        self.keys.len() + 1 == self.keys.capacity()
510    }
511
512    fn is_overflowing(&self) -> bool {
513        self.keys.is_full()
514    }
515
516    fn assert_invariants(&self, path: &str, _depth: u8)
517    where
518        K: Ord,
519    {
520        assert!(
521            self.keys.len() == self.values.len(),
522            "{path}: keys.len() != values.len()"
523        );
524        assert!(!self.is_overflowing(), "{path}: overflowing");
525        assert!(self.keys.is_sorted(), "{path}: keys are not sorted");
526    }
527}
528
529enum LeafSearchResult {
530    Found(usize),
531    Missing(usize),
532}
533
534#[cfg(test)]
535mod tests {
536    use super::*;
537
538    fn xorshift(mut x: u32) -> u32 {
539        assert!(x != 0);
540        x ^= x << 13;
541        x ^= x >> 17;
542        x ^= x << 5;
543        x
544    }
545
546    #[test]
547    fn big_insert() {
548        let mut map = BTreeMap::new();
549        let count: u32 = 1000;
550        for x in (1..=count).map(xorshift) {
551            map.insert(x, x);
552        }
553        assert_eq!(map.len(), 1000);
554        for x in (1..=count).map(xorshift) {
555            assert_eq!(map.remove(&x), Some(x));
556        }
557        assert!(map.is_empty());
558    }
559
560    #[test]
561    fn leaf_split() {
562        let mut leaf1 = NodeLeaf {
563            keys: ArrayVec::new(),
564            values: ArrayVec::new(),
565        };
566
567        assert!(matches!(leaf1.insert(1, 11), NodeInsertResult::Inserted));
568        assert!(matches!(leaf1.insert(2, 22), NodeInsertResult::Inserted));
569        assert!(matches!(leaf1.insert(3, 33), NodeInsertResult::Inserted));
570        assert!(matches!(leaf1.insert(4, 44), NodeInsertResult::Inserted));
571        let NodeInsertResult::Split(parent_key, Node::Leaf(leaf2)) = leaf1.insert(5, 55) else {
572            panic!();
573        };
574
575        assert_eq!(leaf1.keys.as_slice(), &[1, 2]);
576        assert_eq!(leaf1.values.as_slice(), &[11, 22]);
577
578        assert_eq!(parent_key, 3);
579
580        assert_eq!(leaf2.keys.as_slice(), &[3, 4, 5]);
581        assert_eq!(leaf2.values.as_slice(), &[33, 44, 55]);
582    }
583
584    #[test]
585    fn leaf_to_branch() {
586        let mut map = BTreeMap::new();
587        assert!(map.insert(1, 11).is_none());
588        assert!(map.insert(2, 22).is_none());
589        assert!(map.insert(3, 33).is_none());
590        assert!(map.insert(4, 44).is_none());
591        assert!(map.root.is_leaf());
592        assert!(map.insert(5, 55).is_none());
593        assert!(map.root.is_branch());
594        assert_eq!(map.len(), 5);
595    }
596}
indigo_kernel/btree_map.rs

indigo_kernel/
btree_map.rs
