cranelift_bforest/
node.rs

1//! B+-tree nodes.
2
3use super::{slice_insert, slice_shift, Forest, Node, SetValue, INNER_SIZE};
4use core::borrow::{Borrow, BorrowMut};
5use core::fmt;
6
7/// B+-tree node.
8///
9/// A B+-tree has different node types for inner nodes and leaf nodes. Inner nodes contain M node
10/// references and M-1 keys while leaf nodes contain N keys and values. Values for M and N are
11/// chosen such that a node is exactly 64 bytes (a cache line) when keys and values are 32 bits
12/// each.
13///
14/// An inner node contains at least M/2 node references unless it is the right-most node at its
15/// level. A leaf node contains at least N/2 keys unless it is the right-most leaf.
16#[allow(dead_code)] // workaround for https://github.com/rust-lang/rust/issues/64362
17pub(super) enum NodeData<F: Forest> {
18    Inner {
19        /// The number of keys in this node.
20        /// The number of node references is always one more.
21        size: u8,
22
23        /// Keys discriminating sub-trees.
24        ///
25        /// The key in `keys[i]` is greater than all keys in `tree[i]` and less than or equal to
26        /// all keys in `tree[i+1]`.
27        keys: [F::Key; INNER_SIZE - 1],
28
29        /// Sub-trees.
30        tree: [Node; INNER_SIZE],
31    },
32    Leaf {
33        /// Number of key-value pairs in this node.
34        size: u8,
35
36        // Key array.
37        keys: F::LeafKeys,
38
39        // Value array.
40        vals: F::LeafValues,
41    },
42    /// An unused node on the free list.
43    Free { next: Option<Node> },
44}
45
46// Implement `Clone` and `Copy` manually, because deriving them would also require `Forest` to
47// implement `Clone`.
48impl<F: Forest> Copy for NodeData<F> {}
49impl<F: Forest> Clone for NodeData<F> {
50    fn clone(&self) -> Self {
51        *self
52    }
53}
54
55impl<F: Forest> NodeData<F> {
56    /// Is this a free/unused node?
57    pub fn is_free(&self) -> bool {
58        match *self {
59            Self::Free { .. } => true,
60            _ => false,
61        }
62    }
63
64    /// Get the number of entries in this node.
65    ///
66    /// This is the number of outgoing edges in an inner node, or the number of key-value pairs in
67    /// a leaf node.
68    pub fn entries(&self) -> usize {
69        match *self {
70            Self::Inner { size, .. } => usize::from(size) + 1,
71            Self::Leaf { size, .. } => usize::from(size),
72            Self::Free { .. } => panic!("freed node"),
73        }
74    }
75
76    /// Create an inner node with a single key and two sub-trees.
77    pub fn inner(left: Node, key: F::Key, right: Node) -> Self {
78        // Splat the key and right node to the whole array.
79        // Saves us from inventing a default/reserved value.
80        let mut tree = [right; INNER_SIZE];
81        tree[0] = left;
82        Self::Inner {
83            size: 1,
84            keys: [key; INNER_SIZE - 1],
85            tree,
86        }
87    }
88
89    /// Create a leaf node with a single key-value pair.
90    pub fn leaf(key: F::Key, value: F::Value) -> Self {
91        Self::Leaf {
92            size: 1,
93            keys: F::splat_key(key),
94            vals: F::splat_value(value),
95        }
96    }
97
98    /// Unwrap an inner node into two slices (keys, trees).
99    pub fn unwrap_inner(&self) -> (&[F::Key], &[Node]) {
100        match *self {
101            Self::Inner {
102                size,
103                ref keys,
104                ref tree,
105            } => {
106                let size = usize::from(size);
107                // TODO: We could probably use `get_unchecked()` here since `size` is always in
108                // range.
109                (&keys[0..size], &tree[0..=size])
110            }
111            _ => panic!("Expected inner node"),
112        }
113    }
114
115    /// Unwrap a leaf node into two slices (keys, values) of the same length.
116    pub fn unwrap_leaf(&self) -> (&[F::Key], &[F::Value]) {
117        match *self {
118            Self::Leaf {
119                size,
120                ref keys,
121                ref vals,
122            } => {
123                let size = usize::from(size);
124                let keys = keys.borrow();
125                let vals = vals.borrow();
126                // TODO: We could probably use `get_unchecked()` here since `size` is always in
127                // range.
128                (&keys[0..size], &vals[0..size])
129            }
130            _ => panic!("Expected leaf node"),
131        }
132    }
133
134    /// Unwrap a mutable leaf node into two slices (keys, values) of the same length.
135    pub fn unwrap_leaf_mut(&mut self) -> (&mut [F::Key], &mut [F::Value]) {
136        match *self {
137            Self::Leaf {
138                size,
139                ref mut keys,
140                ref mut vals,
141            } => {
142                let size = usize::from(size);
143                let keys = keys.borrow_mut();
144                let vals = vals.borrow_mut();
145                // TODO: We could probably use `get_unchecked_mut()` here since `size` is always in
146                // range.
147                (&mut keys[0..size], &mut vals[0..size])
148            }
149            _ => panic!("Expected leaf node"),
150        }
151    }
152
153    /// Get the critical key for a leaf node.
154    /// This is simply the first key.
155    pub fn leaf_crit_key(&self) -> F::Key {
156        match *self {
157            Self::Leaf { size, ref keys, .. } => {
158                debug_assert!(size > 0, "Empty leaf node");
159                keys.borrow()[0]
160            }
161            _ => panic!("Expected leaf node"),
162        }
163    }
164
165    /// Try to insert `(key, node)` at key-position `index` in an inner node.
166    /// This means that `key` is inserted at `keys[i]` and `node` is inserted at `tree[i + 1]`.
167    /// If the node is full, this leaves the node unchanged and returns false.
168    pub fn try_inner_insert(&mut self, index: usize, key: F::Key, node: Node) -> bool {
169        match *self {
170            Self::Inner {
171                ref mut size,
172                ref mut keys,
173                ref mut tree,
174            } => {
175                let sz = usize::from(*size);
176                debug_assert!(sz <= keys.len());
177                debug_assert!(index <= sz, "Can't insert at {} with {} keys", index, sz);
178
179                if let Some(ks) = keys.get_mut(0..=sz) {
180                    *size = (sz + 1) as u8;
181                    slice_insert(ks, index, key);
182                    slice_insert(&mut tree[1..=sz + 1], index, node);
183                    true
184                } else {
185                    false
186                }
187            }
188            _ => panic!("Expected inner node"),
189        }
190    }
191
192    /// Try to insert `key, value` at `index` in a leaf node, but fail and return false if the node
193    /// is full.
194    pub fn try_leaf_insert(&mut self, index: usize, key: F::Key, value: F::Value) -> bool {
195        match *self {
196            Self::Leaf {
197                ref mut size,
198                ref mut keys,
199                ref mut vals,
200            } => {
201                let sz = usize::from(*size);
202                let keys = keys.borrow_mut();
203                let vals = vals.borrow_mut();
204                debug_assert!(sz <= keys.len());
205                debug_assert!(index <= sz);
206
207                if let Some(ks) = keys.get_mut(0..=sz) {
208                    *size = (sz + 1) as u8;
209                    slice_insert(ks, index, key);
210                    slice_insert(&mut vals[0..=sz], index, value);
211                    true
212                } else {
213                    false
214                }
215            }
216            _ => panic!("Expected leaf node"),
217        }
218    }
219
220    /// Split off the second half of this node.
221    /// It is assumed that this a completely full inner or leaf node.
222    ///
223    /// The `insert_index` parameter is the position where an insertion was tried and failed. The
224    /// node will be split in half with a bias towards an even split after the insertion is retried.
225    pub fn split(&mut self, insert_index: usize) -> SplitOff<F> {
226        match *self {
227            Self::Inner {
228                ref mut size,
229                ref keys,
230                ref tree,
231            } => {
232                debug_assert_eq!(usize::from(*size), keys.len(), "Node not full");
233
234                // Number of tree entries in the lhs node.
235                let l_ents = split_pos(tree.len(), insert_index + 1);
236                let r_ents = tree.len() - l_ents;
237
238                // With INNER_SIZE=8, we get l_ents=4 and:
239                //
240                // self: [ n0 k0 n1 k1 n2 k2 n3 k3 n4 k4 n5 k5 n6 k6 n7 ]
241                // lhs:  [ n0 k0 n1 k1 n2 k2 n3 ]
242                // crit_key = k3 (not present in either node)
243                // rhs:  [ n4 k4 n5 k5 n6 k6 n7 ]
244
245                // 1. Truncate the LHS.
246                *size = (l_ents - 1) as u8;
247
248                // 2. Copy second half to `rhs_data`.
249                let mut r_keys = *keys;
250                r_keys[0..r_ents - 1].copy_from_slice(&keys[l_ents..]);
251
252                let mut r_tree = *tree;
253                r_tree[0..r_ents].copy_from_slice(&tree[l_ents..]);
254
255                SplitOff {
256                    lhs_entries: l_ents,
257                    rhs_entries: r_ents,
258                    crit_key: keys[l_ents - 1],
259                    rhs_data: Self::Inner {
260                        size: (r_ents - 1) as u8,
261                        keys: r_keys,
262                        tree: r_tree,
263                    },
264                }
265            }
266            Self::Leaf {
267                ref mut size,
268                ref keys,
269                ref vals,
270            } => {
271                let o_keys = keys.borrow();
272                let o_vals = vals.borrow();
273                debug_assert_eq!(usize::from(*size), o_keys.len(), "Node not full");
274
275                let l_size = split_pos(o_keys.len(), insert_index);
276                let r_size = o_keys.len() - l_size;
277
278                // 1. Truncate the LHS node at `l_size`.
279                *size = l_size as u8;
280
281                // 2. Copy second half to `rhs_data`.
282                let mut r_keys = *keys;
283                r_keys.borrow_mut()[0..r_size].copy_from_slice(&o_keys[l_size..]);
284
285                let mut r_vals = *vals;
286                r_vals.borrow_mut()[0..r_size].copy_from_slice(&o_vals[l_size..]);
287
288                SplitOff {
289                    lhs_entries: l_size,
290                    rhs_entries: r_size,
291                    crit_key: o_keys[l_size],
292                    rhs_data: Self::Leaf {
293                        size: r_size as u8,
294                        keys: r_keys,
295                        vals: r_vals,
296                    },
297                }
298            }
299            _ => panic!("Expected leaf node"),
300        }
301    }
302
303    /// Remove the sub-tree at `index` from this inner node.
304    ///
305    /// Note that `index` refers to a sub-tree entry and not a key entry as it does for
306    /// `try_inner_insert()`. It is possible to remove the first sub-tree (which can't be inserted
307    /// by `try_inner_insert()`).
308    ///
309    /// Return an indication of the node's health (i.e. below half capacity).
310    pub fn inner_remove(&mut self, index: usize) -> Removed {
311        match *self {
312            Self::Inner {
313                ref mut size,
314                ref mut keys,
315                ref mut tree,
316            } => {
317                let ents = usize::from(*size) + 1;
318                debug_assert!(ents <= tree.len());
319                debug_assert!(index < ents);
320                // Leave an invalid 0xff size when node becomes empty.
321                *size = ents.wrapping_sub(2) as u8;
322                if ents > 1 {
323                    slice_shift(&mut keys[index.saturating_sub(1)..ents - 1], 1);
324                }
325                slice_shift(&mut tree[index..ents], 1);
326                Removed::new(index, ents - 1, tree.len())
327            }
328            _ => panic!("Expected inner node"),
329        }
330    }
331
332    /// Remove the key-value pair at `index` from this leaf node.
333    ///
334    /// Return an indication of the node's health (i.e. below half capacity).
335    pub fn leaf_remove(&mut self, index: usize) -> Removed {
336        match *self {
337            Self::Leaf {
338                ref mut size,
339                ref mut keys,
340                ref mut vals,
341            } => {
342                let sz = usize::from(*size);
343                let keys = keys.borrow_mut();
344                let vals = vals.borrow_mut();
345                *size -= 1;
346                slice_shift(&mut keys[index..sz], 1);
347                slice_shift(&mut vals[index..sz], 1);
348                Removed::new(index, sz - 1, keys.len())
349            }
350            _ => panic!("Expected leaf node"),
351        }
352    }
353
354    /// Balance this node with its right sibling.
355    ///
356    /// It is assumed that the current node has underflowed. Look at the right sibling node and do
357    /// one of two things:
358    ///
359    /// 1. Move all entries to the right node, leaving this node empty, or
360    /// 2. Distribute entries evenly between the two nodes.
361    ///
362    /// In the first case, `None` is returned. In the second case, the new critical key for the
363    /// right sibling node is returned.
364    pub fn balance(&mut self, crit_key: F::Key, rhs: &mut Self) -> Option<F::Key> {
365        match (self, rhs) {
366            (
367                &mut Self::Inner {
368                    size: ref mut l_size,
369                    keys: ref mut l_keys,
370                    tree: ref mut l_tree,
371                },
372                &mut Self::Inner {
373                    size: ref mut r_size,
374                    keys: ref mut r_keys,
375                    tree: ref mut r_tree,
376                },
377            ) => {
378                let l_ents = usize::from(*l_size) + 1;
379                let r_ents = usize::from(*r_size) + 1;
380                let ents = l_ents + r_ents;
381
382                if ents <= r_tree.len() {
383                    // All entries will fit in the RHS node.
384                    // We'll leave the LHS node empty, but first use it as a scratch space.
385                    *l_size = 0;
386                    // Insert `crit_key` between the two nodes.
387                    l_keys[l_ents - 1] = crit_key;
388                    l_keys[l_ents..ents - 1].copy_from_slice(&r_keys[0..r_ents - 1]);
389                    r_keys[0..ents - 1].copy_from_slice(&l_keys[0..ents - 1]);
390                    l_tree[l_ents..ents].copy_from_slice(&r_tree[0..r_ents]);
391                    r_tree[0..ents].copy_from_slice(&l_tree[0..ents]);
392                    *r_size = (ents - 1) as u8;
393                    None
394                } else {
395                    // The entries don't all fit in one node. Distribute some from RHS -> LHS.
396                    // Split evenly with a bias to putting one entry in LHS.
397                    let r_goal = ents / 2;
398                    let l_goal = ents - r_goal;
399                    debug_assert!(l_goal > l_ents, "Node must be underflowed");
400
401                    l_keys[l_ents - 1] = crit_key;
402                    l_keys[l_ents..l_goal - 1].copy_from_slice(&r_keys[0..l_goal - 1 - l_ents]);
403                    l_tree[l_ents..l_goal].copy_from_slice(&r_tree[0..l_goal - l_ents]);
404                    *l_size = (l_goal - 1) as u8;
405
406                    let new_crit = r_keys[r_ents - r_goal - 1];
407                    slice_shift(&mut r_keys[0..r_ents - 1], r_ents - r_goal);
408                    slice_shift(&mut r_tree[0..r_ents], r_ents - r_goal);
409                    *r_size = (r_goal - 1) as u8;
410
411                    Some(new_crit)
412                }
413            }
414            (
415                &mut Self::Leaf {
416                    size: ref mut l_size,
417                    keys: ref mut l_keys,
418                    vals: ref mut l_vals,
419                },
420                &mut Self::Leaf {
421                    size: ref mut r_size,
422                    keys: ref mut r_keys,
423                    vals: ref mut r_vals,
424                },
425            ) => {
426                let l_ents = usize::from(*l_size);
427                let l_keys = l_keys.borrow_mut();
428                let l_vals = l_vals.borrow_mut();
429                let r_ents = usize::from(*r_size);
430                let r_keys = r_keys.borrow_mut();
431                let r_vals = r_vals.borrow_mut();
432                let ents = l_ents + r_ents;
433
434                if ents <= r_vals.len() {
435                    // We can fit all entries in the RHS node.
436                    // We'll leave the LHS node empty, but first use it as a scratch space.
437                    *l_size = 0;
438                    l_keys[l_ents..ents].copy_from_slice(&r_keys[0..r_ents]);
439                    r_keys[0..ents].copy_from_slice(&l_keys[0..ents]);
440                    l_vals[l_ents..ents].copy_from_slice(&r_vals[0..r_ents]);
441                    r_vals[0..ents].copy_from_slice(&l_vals[0..ents]);
442                    *r_size = ents as u8;
443                    None
444                } else {
445                    // The entries don't all fit in one node. Distribute some from RHS -> LHS.
446                    // Split evenly with a bias to putting one entry in LHS.
447                    let r_goal = ents / 2;
448                    let l_goal = ents - r_goal;
449                    debug_assert!(l_goal > l_ents, "Node must be underflowed");
450
451                    l_keys[l_ents..l_goal].copy_from_slice(&r_keys[0..l_goal - l_ents]);
452                    l_vals[l_ents..l_goal].copy_from_slice(&r_vals[0..l_goal - l_ents]);
453                    *l_size = l_goal as u8;
454
455                    slice_shift(&mut r_keys[0..r_ents], r_ents - r_goal);
456                    slice_shift(&mut r_vals[0..r_ents], r_ents - r_goal);
457                    *r_size = r_goal as u8;
458
459                    Some(r_keys[0])
460                }
461            }
462            _ => panic!("Mismatched nodes"),
463        }
464    }
465}
466
467/// Find the right split position for halving a full node with `len` entries to recover from a
468/// failed insertion at `ins`.
469///
470/// If `len` is even, we should split straight down the middle regardless of `len`.
471///
472/// If `len` is odd, we should split the node such that the two halves are the same size after the
473/// insertion is retried.
474fn split_pos(len: usize, ins: usize) -> usize {
475    // Anticipate `len` being a compile time constant, so this all folds away when `len` is even.
476    if ins <= len / 2 {
477        len / 2
478    } else {
479        (len + 1) / 2
480    }
481}
482
483/// The result of splitting off the second half of a node.
484pub(super) struct SplitOff<F: Forest> {
485    /// The number of entries left in the original node which becomes the left-hand-side of the
486    /// pair. This is the number of outgoing node edges for an inner node, and the number of
487    /// key-value pairs for a leaf node.
488    pub lhs_entries: usize,
489
490    /// The number of entries in the new RHS node.
491    pub rhs_entries: usize,
492
493    /// The critical key separating the LHS and RHS nodes. All keys in the LHS sub-tree are less
494    /// than the critical key, and all entries in the RHS sub-tree are greater or equal to the
495    /// critical key.
496    pub crit_key: F::Key,
497
498    /// The RHS node data containing the elements that were removed from the original node (now the
499    /// LHS).
500    pub rhs_data: NodeData<F>,
501}
502
503/// The result of removing an entry from a node.
504#[derive(Clone, Copy, Debug, PartialEq, Eq)]
505pub(super) enum Removed {
506    /// An entry was removed, and the node is still in good shape.
507    Healthy,
508
509    /// The node is in good shape after removing the rightmost element.
510    Rightmost,
511
512    /// The node has too few entries now, and it should be balanced with a sibling node.
513    Underflow,
514
515    /// The last entry was removed. For an inner node, this means that the `keys` array is empty
516    /// and there is just a single sub-tree left.
517    Empty,
518}
519
520impl Removed {
521    /// Create a `Removed` status from a size and capacity.
522    fn new(removed: usize, new_size: usize, capacity: usize) -> Self {
523        if 2 * new_size >= capacity {
524            if removed == new_size {
525                Self::Rightmost
526            } else {
527                Self::Healthy
528            }
529        } else if new_size > 0 {
530            Self::Underflow
531        } else {
532            Self::Empty
533        }
534    }
535}
536
537// Display ": value" or nothing at all for `()`.
538pub(super) trait ValDisp {
539    fn valfmt(&self, f: &mut fmt::Formatter) -> fmt::Result;
540}
541
542impl ValDisp for SetValue {
543    fn valfmt(&self, _: &mut fmt::Formatter) -> fmt::Result {
544        Ok(())
545    }
546}
547
548impl<T: fmt::Display> ValDisp for T {
549    fn valfmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
550        write!(f, ":{}", self)
551    }
552}
553
554impl<F> fmt::Display for NodeData<F>
555where
556    F: Forest,
557    F::Key: fmt::Display,
558    F::Value: ValDisp,
559{
560    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
561        match *self {
562            Self::Inner { size, keys, tree } => {
563                write!(f, "[ {}", tree[0])?;
564                for i in 0..usize::from(size) {
565                    write!(f, " {} {}", keys[i], tree[i + 1])?;
566                }
567                write!(f, " ]")
568            }
569            Self::Leaf { size, keys, vals } => {
570                let keys = keys.borrow();
571                let vals = vals.borrow();
572                write!(f, "[")?;
573                for i in 0..usize::from(size) {
574                    write!(f, " {}", keys[i])?;
575                    vals[i].valfmt(f)?;
576                }
577                write!(f, " ]")
578            }
579            Self::Free { next: Some(n) } => write!(f, "[ free -> {} ]", n),
580            Self::Free { next: None } => write!(f, "[ free ]"),
581        }
582    }
583}
584
585#[cfg(test)]
586mod tests {
587    use super::*;
588    use alloc::string::ToString;
589    use core::mem;
590
591    // Forest impl for a set implementation.
592    struct TF();
593
594    impl Forest for TF {
595        type Key = char;
596        type Value = SetValue;
597        type LeafKeys = [char; 15];
598        type LeafValues = [SetValue; 15];
599
600        fn splat_key(key: Self::Key) -> Self::LeafKeys {
601            [key; 15]
602        }
603
604        fn splat_value(value: Self::Value) -> Self::LeafValues {
605            [value; 15]
606        }
607    }
608
609    #[test]
610    fn inner() {
611        let n1 = Node(1);
612        let n2 = Node(2);
613        let n3 = Node(3);
614        let n4 = Node(4);
615        let mut inner = NodeData::<TF>::inner(n1, 'c', n4);
616        assert_eq!(mem::size_of_val(&inner), 64);
617        assert_eq!(inner.to_string(), "[ node1 c node4 ]");
618
619        assert!(inner.try_inner_insert(0, 'a', n2));
620        assert_eq!(inner.to_string(), "[ node1 a node2 c node4 ]");
621
622        assert!(inner.try_inner_insert(1, 'b', n3));
623        assert_eq!(inner.to_string(), "[ node1 a node2 b node3 c node4 ]");
624
625        for i in 3..7 {
626            assert!(inner.try_inner_insert(
627                usize::from(i),
628                ('a' as u8 + i) as char,
629                Node(i as u32 + 2),
630            ));
631        }
632        assert_eq!(
633            inner.to_string(),
634            "[ node1 a node2 b node3 c node4 d node5 e node6 f node7 g node8 ]"
635        );
636
637        // Now the node is full and insertion should fail anywhere.
638        assert!(!inner.try_inner_insert(0, 'x', n3));
639        assert!(!inner.try_inner_insert(4, 'x', n3));
640        assert!(!inner.try_inner_insert(7, 'x', n3));
641
642        // Splitting should be independent of the hint because we have an even number of node
643        // references.
644        let saved = inner.clone();
645        let sp = inner.split(1);
646        assert_eq!(sp.lhs_entries, 4);
647        assert_eq!(sp.rhs_entries, 4);
648        assert_eq!(sp.crit_key, 'd');
649        // The critical key is not present in either of the resulting nodes.
650        assert_eq!(inner.to_string(), "[ node1 a node2 b node3 c node4 ]");
651        assert_eq!(sp.rhs_data.to_string(), "[ node5 e node6 f node7 g node8 ]");
652
653        assert_eq!(inner.inner_remove(0), Removed::Underflow);
654        assert_eq!(inner.to_string(), "[ node2 b node3 c node4 ]");
655
656        assert_eq!(inner.inner_remove(1), Removed::Underflow);
657        assert_eq!(inner.to_string(), "[ node2 c node4 ]");
658
659        assert_eq!(inner.inner_remove(1), Removed::Underflow);
660        assert_eq!(inner.to_string(), "[ node2 ]");
661
662        assert_eq!(inner.inner_remove(0), Removed::Empty);
663
664        inner = saved;
665        let sp = inner.split(6);
666        assert_eq!(sp.lhs_entries, 4);
667        assert_eq!(sp.rhs_entries, 4);
668        assert_eq!(sp.crit_key, 'd');
669        assert_eq!(inner.to_string(), "[ node1 a node2 b node3 c node4 ]");
670        assert_eq!(sp.rhs_data.to_string(), "[ node5 e node6 f node7 g node8 ]");
671    }
672
673    #[test]
674    fn leaf() {
675        let mut leaf = NodeData::<TF>::leaf('d', SetValue());
676        assert_eq!(leaf.to_string(), "[ d ]");
677
678        assert!(leaf.try_leaf_insert(0, 'a', SetValue()));
679        assert_eq!(leaf.to_string(), "[ a d ]");
680        assert!(leaf.try_leaf_insert(1, 'b', SetValue()));
681        assert!(leaf.try_leaf_insert(2, 'c', SetValue()));
682        assert_eq!(leaf.to_string(), "[ a b c d ]");
683        for i in 4..15 {
684            assert!(leaf.try_leaf_insert(usize::from(i), ('a' as u8 + i) as char, SetValue()));
685        }
686        assert_eq!(leaf.to_string(), "[ a b c d e f g h i j k l m n o ]");
687
688        // Now the node is full and insertion should fail anywhere.
689        assert!(!leaf.try_leaf_insert(0, 'x', SetValue()));
690        assert!(!leaf.try_leaf_insert(8, 'x', SetValue()));
691        assert!(!leaf.try_leaf_insert(15, 'x', SetValue()));
692
693        // The index given to `split` is not the split position, it's a hint for balancing the node.
694        let saved = leaf.clone();
695        let sp = leaf.split(12);
696        assert_eq!(sp.lhs_entries, 8);
697        assert_eq!(sp.rhs_entries, 7);
698        assert_eq!(sp.crit_key, 'i');
699        assert_eq!(leaf.to_string(), "[ a b c d e f g h ]");
700        assert_eq!(sp.rhs_data.to_string(), "[ i j k l m n o ]");
701
702        assert!(leaf.try_leaf_insert(8, 'i', SetValue()));
703        assert_eq!(leaf.leaf_remove(2), Removed::Healthy);
704        assert_eq!(leaf.to_string(), "[ a b d e f g h i ]");
705        assert_eq!(leaf.leaf_remove(7), Removed::Underflow);
706        assert_eq!(leaf.to_string(), "[ a b d e f g h ]");
707
708        leaf = saved;
709        let sp = leaf.split(7);
710        assert_eq!(sp.lhs_entries, 7);
711        assert_eq!(sp.rhs_entries, 8);
712        assert_eq!(sp.crit_key, 'h');
713        assert_eq!(leaf.to_string(), "[ a b c d e f g ]");
714        assert_eq!(sp.rhs_data.to_string(), "[ h i j k l m n o ]");
715    }
716
717    #[test]
718    fn optimal_split_pos() {
719        // An even split is easy.
720        assert_eq!(split_pos(8, 0), 4);
721        assert_eq!(split_pos(8, 8), 4);
722
723        // Easy cases for odd splits.
724        assert_eq!(split_pos(7, 0), 3);
725        assert_eq!(split_pos(7, 7), 4);
726
727        // If the insertion point is the same as the split position, we
728        // will append to the lhs node.
729        assert_eq!(split_pos(7, 3), 3);
730        assert_eq!(split_pos(7, 4), 4);
731    }
732
733    #[test]
734    fn inner_balance() {
735        let n1 = Node(1);
736        let n2 = Node(2);
737        let n3 = Node(3);
738        let mut lhs = NodeData::<TF>::inner(n1, 'a', n2);
739        assert!(lhs.try_inner_insert(1, 'b', n3));
740        assert_eq!(lhs.to_string(), "[ node1 a node2 b node3 ]");
741
742        let n11 = Node(11);
743        let n12 = Node(12);
744        let mut rhs = NodeData::<TF>::inner(n11, 'p', n12);
745
746        for i in 1..4 {
747            assert!(rhs.try_inner_insert(
748                usize::from(i),
749                ('p' as u8 + i) as char,
750                Node(i as u32 + 12),
751            ));
752        }
753        assert_eq!(
754            rhs.to_string(),
755            "[ node11 p node12 q node13 r node14 s node15 ]"
756        );
757
758        // 3+5 elements fit in RHS.
759        assert_eq!(lhs.balance('o', &mut rhs), None);
760        assert_eq!(
761            rhs.to_string(),
762            "[ node1 a node2 b node3 o node11 p node12 q node13 r node14 s node15 ]"
763        );
764
765        // 2+8 elements are redistributed.
766        lhs = NodeData::<TF>::inner(Node(20), 'x', Node(21));
767        assert_eq!(lhs.balance('y', &mut rhs), Some('o'));
768        assert_eq!(
769            lhs.to_string(),
770            "[ node20 x node21 y node1 a node2 b node3 ]"
771        );
772        assert_eq!(
773            rhs.to_string(),
774            "[ node11 p node12 q node13 r node14 s node15 ]"
775        );
776    }
777
778    #[test]
779    fn leaf_balance() {
780        let mut lhs = NodeData::<TF>::leaf('a', SetValue());
781        for i in 1..6 {
782            assert!(lhs.try_leaf_insert(usize::from(i), ('a' as u8 + i) as char, SetValue()));
783        }
784        assert_eq!(lhs.to_string(), "[ a b c d e f ]");
785
786        let mut rhs = NodeData::<TF>::leaf('0', SetValue());
787        for i in 1..8 {
788            assert!(rhs.try_leaf_insert(usize::from(i), ('0' as u8 + i) as char, SetValue()));
789        }
790        assert_eq!(rhs.to_string(), "[ 0 1 2 3 4 5 6 7 ]");
791
792        // 6+8 elements all fits in rhs.
793        assert_eq!(lhs.balance('0', &mut rhs), None);
794        assert_eq!(rhs.to_string(), "[ a b c d e f 0 1 2 3 4 5 6 7 ]");
795
796        assert!(lhs.try_leaf_insert(0, 'x', SetValue()));
797        assert!(lhs.try_leaf_insert(1, 'y', SetValue()));
798        assert!(lhs.try_leaf_insert(2, 'z', SetValue()));
799        assert_eq!(lhs.to_string(), "[ x y z ]");
800
801        // 3+14 elements need redistribution.
802        assert_eq!(lhs.balance('a', &mut rhs), Some('0'));
803        assert_eq!(lhs.to_string(), "[ x y z a b c d e f ]");
804        assert_eq!(rhs.to_string(), "[ 0 1 2 3 4 5 6 7 ]");
805    }
806}