1use super::{slice_insert, slice_shift, Forest, Node, SetValue, INNER_SIZE};
4use core::borrow::{Borrow, BorrowMut};
5use core::fmt;
6
7#[allow(dead_code)] pub(super) enum NodeData<F: Forest> {
18 Inner {
19 size: u8,
22
23 keys: [F::Key; INNER_SIZE - 1],
28
29 tree: [Node; INNER_SIZE],
31 },
32 Leaf {
33 size: u8,
35
36 keys: F::LeafKeys,
38
39 vals: F::LeafValues,
41 },
42 Free { next: Option<Node> },
44}
45
46impl<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 pub fn is_free(&self) -> bool {
58 match *self {
59 Self::Free { .. } => true,
60 _ => false,
61 }
62 }
63
64 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 pub fn inner(left: Node, key: F::Key, right: Node) -> Self {
78 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 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 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 (&keys[0..size], &tree[0..=size])
110 }
111 _ => panic!("Expected inner node"),
112 }
113 }
114
115 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 (&keys[0..size], &vals[0..size])
129 }
130 _ => panic!("Expected leaf node"),
131 }
132 }
133
134 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 (&mut keys[0..size], &mut vals[0..size])
148 }
149 _ => panic!("Expected leaf node"),
150 }
151 }
152
153 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 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 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 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 let l_ents = split_pos(tree.len(), insert_index + 1);
236 let r_ents = tree.len() - l_ents;
237
238 *size = (l_ents - 1) as u8;
247
248 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 *size = l_size as u8;
280
281 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 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 *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 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 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 *l_size = 0;
386 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 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 *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 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
467fn split_pos(len: usize, ins: usize) -> usize {
475 if ins <= len / 2 {
477 len / 2
478 } else {
479 (len + 1) / 2
480 }
481}
482
483pub(super) struct SplitOff<F: Forest> {
485 pub lhs_entries: usize,
489
490 pub rhs_entries: usize,
492
493 pub crit_key: F::Key,
497
498 pub rhs_data: NodeData<F>,
501}
502
503#[derive(Clone, Copy, Debug, PartialEq, Eq)]
505pub(super) enum Removed {
506 Healthy,
508
509 Rightmost,
511
512 Underflow,
514
515 Empty,
518}
519
520impl Removed {
521 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
537pub(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 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 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 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 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 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 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 assert_eq!(split_pos(8, 0), 4);
721 assert_eq!(split_pos(8, 8), 4);
722
723 assert_eq!(split_pos(7, 0), 3);
725 assert_eq!(split_pos(7, 7), 4);
726
727 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 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 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 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 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}