1use super::node::Removed;
4use super::{slice_insert, slice_shift, Comparator, Forest, Node, NodeData, NodePool, MAX_PATH};
5use core::borrow::Borrow;
6use core::marker::PhantomData;
7
8#[cfg(test)]
9use core::fmt;
10
11pub(super) struct Path<F: Forest> {
12 size: usize,
14
15 node: [Node; MAX_PATH],
17
18 entry: [u8; MAX_PATH],
20
21 unused: PhantomData<F>,
22}
23
24impl<F: Forest> Default for Path<F> {
25 fn default() -> Self {
26 Self {
27 size: 0,
28 node: [Node(0); MAX_PATH],
29 entry: [0; MAX_PATH],
30 unused: PhantomData,
31 }
32 }
33}
34
35impl<F: Forest> Path<F> {
36 pub fn find(
48 &mut self,
49 key: F::Key,
50 root: Node,
51 pool: &NodePool<F>,
52 comp: &dyn Comparator<F::Key>,
53 ) -> Option<F::Value> {
54 let mut node = root;
55 for level in 0.. {
56 self.size = level + 1;
57 self.node[level] = node;
58 match pool[node] {
59 NodeData::Inner { size, keys, tree } => {
60 let i = match comp.search(key, &keys[0..size.into()]) {
63 Ok(i) => i + 1,
65 Err(i) => i,
67 };
68 self.entry[level] = i as u8;
69 node = tree[i];
70 }
71 NodeData::Leaf { size, keys, vals } => {
72 return match comp.search(key, &keys.borrow()[0..size.into()]) {
74 Ok(i) => {
75 self.entry[level] = i as u8;
76 Some(vals.borrow()[i])
77 }
78 Err(i) => {
79 self.entry[level] = i as u8;
80 None
81 }
82 };
83 }
84 NodeData::Free { .. } => panic!("Free {} reached from {}", node, root),
85 }
86 }
87 unreachable!();
88 }
89
90 pub fn first(&mut self, root: Node, pool: &NodePool<F>) -> (F::Key, F::Value) {
92 let mut node = root;
93 for level in 0.. {
94 self.size = level + 1;
95 self.node[level] = node;
96 self.entry[level] = 0;
97 match pool[node] {
98 NodeData::Inner { tree, .. } => node = tree[0],
99 NodeData::Leaf { keys, vals, .. } => return (keys.borrow()[0], vals.borrow()[0]),
100 NodeData::Free { .. } => panic!("Free {} reached from {}", node, root),
101 }
102 }
103 unreachable!();
104 }
105
106 pub fn next(&mut self, pool: &NodePool<F>) -> Option<(F::Key, F::Value)> {
108 match self.leaf_pos() {
109 None => return None,
110 Some((node, entry)) => {
111 let (keys, vals) = pool[node].unwrap_leaf();
112 if entry + 1 < keys.len() {
113 self.entry[self.size - 1] += 1;
114 return Some((keys[entry + 1], vals[entry + 1]));
115 }
116 }
117 }
118
119 let leaf_level = self.size - 1;
121 self.next_node(leaf_level, pool).map(|node| {
122 let (keys, vals) = pool[node].unwrap_leaf();
123 (keys[0], vals[0])
124 })
125 }
126
127 pub fn prev(&mut self, root: Node, pool: &NodePool<F>) -> Option<(F::Key, F::Value)> {
133 if self.size == 0 {
135 self.goto_subtree_last(0, root, pool);
136 let (node, entry) = self.leaf_pos().unwrap();
137 let (keys, vals) = pool[node].unwrap_leaf();
138 return Some((keys[entry], vals[entry]));
139 }
140
141 match self.leaf_pos() {
142 None => return None,
143 Some((node, entry)) => {
144 if entry > 0 {
145 self.entry[self.size - 1] -= 1;
146 let (keys, vals) = pool[node].unwrap_leaf();
147 return Some((keys[entry - 1], vals[entry - 1]));
148 }
149 }
150 }
151
152 self.prev_leaf(pool).map(|node| {
154 let (keys, vals) = pool[node].unwrap_leaf();
155 let e = self.leaf_entry();
156 (keys[e], vals[e])
157 })
158 }
159
160 fn next_node(&mut self, level: usize, pool: &NodePool<F>) -> Option<Node> {
166 match self.right_sibling_branch_level(level, pool) {
167 None => {
168 self.size = 0;
169 None
170 }
171 Some(bl) => {
172 let (_, bnodes) = pool[self.node[bl]].unwrap_inner();
173 self.entry[bl] += 1;
174 let mut node = bnodes[usize::from(self.entry[bl])];
175
176 for l in bl + 1..level {
177 self.node[l] = node;
178 self.entry[l] = 0;
179 node = pool[node].unwrap_inner().1[0];
180 }
181
182 self.node[level] = node;
183 self.entry[level] = 0;
184 Some(node)
185 }
186 }
187 }
188
189 fn prev_leaf(&mut self, pool: &NodePool<F>) -> Option<Node> {
195 self.left_sibling_branch_level(self.size - 1).map(|bl| {
196 let entry = self.entry[bl] - 1;
197 self.entry[bl] = entry;
198 let (_, bnodes) = pool[self.node[bl]].unwrap_inner();
199 self.goto_subtree_last(bl + 1, bnodes[usize::from(entry)], pool)
200 })
201 }
202
203 fn goto_subtree_last(&mut self, level: usize, root: Node, pool: &NodePool<F>) -> Node {
205 let mut node = root;
206 for l in level.. {
207 self.node[l] = node;
208 match pool[node] {
209 NodeData::Inner { size, ref tree, .. } => {
210 self.entry[l] = size;
211 node = tree[usize::from(size)];
212 }
213 NodeData::Leaf { size, .. } => {
214 self.entry[l] = size - 1;
215 self.size = l + 1;
216 break;
217 }
218 NodeData::Free { .. } => panic!("Free {} reached from {}", node, root),
219 }
220 }
221 node
222 }
223
224 pub fn set_root_node(&mut self, root: Node) {
226 self.size = 1;
227 self.node[0] = root;
228 self.entry[0] = 0;
229 }
230
231 pub fn leaf_pos(&self) -> Option<(Node, usize)> {
233 let i = self.size.wrapping_sub(1);
234 self.node.get(i).map(|&n| (n, self.entry[i].into()))
235 }
236
237 fn leaf_node(&self) -> Node {
239 self.node[self.size - 1]
240 }
241
242 fn leaf_entry(&self) -> usize {
244 self.entry[self.size - 1].into()
245 }
246
247 fn at_first_entry(&self) -> bool {
250 self.entry[0..self.size].iter().all(|&i| i == 0)
251 }
252
253 pub fn value_mut<'a>(&self, pool: &'a mut NodePool<F>) -> &'a mut F::Value {
256 &mut pool[self.leaf_node()].unwrap_leaf_mut().1[self.leaf_entry()]
257 }
258
259 pub fn insert(&mut self, key: F::Key, value: F::Value, pool: &mut NodePool<F>) -> Node {
264 if !self.try_leaf_insert(key, value, pool) {
265 self.split_and_insert(key, value, pool);
266 }
267 self.node[0]
268 }
269
270 fn try_leaf_insert(&self, key: F::Key, value: F::Value, pool: &mut NodePool<F>) -> bool {
273 let index = self.leaf_entry();
274
275 debug_assert!(index > 0 || self.at_first_entry());
279
280 pool[self.leaf_node()].try_leaf_insert(index, key, value)
281 }
282
283 fn split_and_insert(&mut self, mut key: F::Key, value: F::Value, pool: &mut NodePool<F>) {
286 let orig_root = self.node[0];
287
288 let mut ins_node = None;
291 let mut split;
292 for level in (0..self.size).rev() {
293 let mut node = self.node[level];
295 let mut entry = self.entry[level].into();
296 split = pool[node].split(entry);
297 let rhs_node = pool.alloc_node(split.rhs_data);
298
299 if entry > split.lhs_entries
307 || (entry == split.lhs_entries
308 && (split.lhs_entries > split.rhs_entries || ins_node.is_some()))
309 {
310 node = rhs_node;
311 entry -= split.lhs_entries;
312 self.node[level] = node;
313 self.entry[level] = entry as u8;
314 }
315
316 match ins_node {
318 None => {
319 let inserted = pool[node].try_leaf_insert(entry, key, value);
320 debug_assert!(inserted);
321 if entry == 0 && node == rhs_node {
324 split.crit_key = key;
325 }
326 }
327 Some(n) => {
328 let inserted = pool[node].try_inner_insert(entry, key, n);
329 debug_assert!(inserted);
330 if n == self.node[level + 1] {
333 self.entry[level] += 1;
334 }
335 }
336 }
337
338 key = split.crit_key;
341 ins_node = Some(rhs_node);
342 if level > 0 {
343 let pnode = &mut pool[self.node[level - 1]];
344 let pentry = self.entry[level - 1].into();
345 if pnode.try_inner_insert(pentry, key, rhs_node) {
346 if node == rhs_node {
348 self.entry[level - 1] += 1;
349 }
350 return;
351 }
352 }
353 }
354
355 let rhs_node = ins_node.expect("empty path");
357 let root = pool.alloc_node(NodeData::inner(orig_root, key, rhs_node));
358 let entry = if self.node[0] == rhs_node { 1 } else { 0 };
359 self.size += 1;
360 slice_insert(&mut self.node[0..self.size], 0, root);
361 slice_insert(&mut self.entry[0..self.size], 0, entry);
362 }
363
364 pub fn remove(&mut self, pool: &mut NodePool<F>) -> Option<Node> {
369 let e = self.leaf_entry();
370 match pool[self.leaf_node()].leaf_remove(e) {
371 Removed::Healthy => {
372 if e == 0 {
373 self.update_crit_key(pool)
374 }
375 Some(self.node[0])
376 }
377 status => self.balance_nodes(status, pool),
378 }
379 }
380
381 fn current_crit_key(&self, level: usize, pool: &NodePool<F>) -> Option<F::Key> {
388 self.left_sibling_branch_level(level).map(|bl| {
390 let (keys, _) = pool[self.node[bl]].unwrap_inner();
391 keys[usize::from(self.entry[bl]) - 1]
392 })
393 }
394
395 fn update_crit_key(&mut self, pool: &mut NodePool<F>) {
397 let crit_level = match self.left_sibling_branch_level(self.size - 1) {
399 None => return,
400 Some(l) => l,
401 };
402 let crit_kidx = self.entry[crit_level] - 1;
403
404 let crit_key = pool[self.leaf_node()].leaf_crit_key();
406 let crit_node = self.node[crit_level];
407
408 match pool[crit_node] {
409 NodeData::Inner {
410 size, ref mut keys, ..
411 } => {
412 debug_assert!(crit_kidx < size);
413 keys[usize::from(crit_kidx)] = crit_key;
414 }
415 _ => panic!("Expected inner node"),
416 }
417 }
418
419 fn balance_nodes(&mut self, status: Removed, pool: &mut NodePool<F>) -> Option<Node> {
424 if status != Removed::Empty && self.leaf_entry() == 0 {
429 self.update_crit_key(pool);
430 }
431
432 let leaf_level = self.size - 1;
433 if self.heal_level(status, leaf_level, pool) {
434 self.size = 0;
436 return None;
437 }
438
439 let mut ns = 0;
441 while let NodeData::Inner {
442 size: 0, ref tree, ..
443 } = pool[self.node[ns]]
444 {
445 ns += 1;
446 self.node[ns] = tree[0];
447 }
448
449 if ns > 0 {
450 for l in 0..ns {
451 pool.free_node(self.node[l]);
452 }
453
454 slice_shift(&mut self.node, ns);
457 slice_shift(&mut self.entry, ns);
458
459 if self.size > 0 {
460 self.size -= ns;
461 }
462 }
463
464 Some(self.node[0])
467 }
468
469 fn heal_level(&mut self, status: Removed, level: usize, pool: &mut NodePool<F>) -> bool {
476 match status {
477 Removed::Healthy => {}
478 Removed::Rightmost => {
479 debug_assert_eq!(
482 usize::from(self.entry[level]),
483 pool[self.node[level]].entries()
484 );
485 self.next_node(level, pool);
486 }
487 Removed::Underflow => self.underflowed_node(level, pool),
488 Removed::Empty => return self.empty_node(level, pool),
489 }
490 false
491 }
492
493 fn underflowed_node(&mut self, level: usize, pool: &mut NodePool<F>) {
500 if let Some((crit_key, rhs_node)) = self.right_sibling(level, pool) {
503 let new_ck: Option<F::Key>;
505 let empty;
506 let mut rhs = pool[rhs_node];
508 match pool[self.node[level]].balance(crit_key, &mut rhs) {
509 None => {
510 new_ck = self.current_crit_key(level, pool);
512 empty = true;
513 }
514 Some(key) => {
515 new_ck = Some(key);
517 empty = false;
518 }
519 }
520 pool[rhs_node] = rhs;
522 if let Some(ck) = new_ck {
525 self.update_right_crit_key(level, ck, pool);
526 }
527 if empty {
528 let empty_tree = self.empty_node(level, pool);
529 debug_assert!(!empty_tree);
530 }
531
532 debug_assert!(usize::from(self.entry[level]) < pool[self.node[level]].entries());
536 } else if usize::from(self.entry[level]) >= pool[self.node[level]].entries() {
537 self.size = 0;
540 }
541 }
542
543 fn empty_node(&mut self, level: usize, pool: &mut NodePool<F>) -> bool {
552 pool.free_node(self.node[level]);
553 if level == 0 {
554 return true;
556 }
557
558 let rhs_node = self.right_sibling(level, pool).map(|(_, n)| n);
560
561 let pl = level - 1;
563 let pe = self.entry[pl].into();
564 let status = pool[self.node[pl]].inner_remove(pe);
565 self.heal_level(status, pl, pool);
566
567 match rhs_node {
569 Some(rhs) => self.node[level] = rhs,
572 None => self.size = 0,
575 }
576 false
577 }
578
579 fn right_sibling_branch_level(&self, level: usize, pool: &NodePool<F>) -> Option<usize> {
586 (0..level).rposition(|l| match pool[self.node[l]] {
587 NodeData::Inner { size, .. } => self.entry[l] < size,
588 _ => panic!("Expected inner node"),
589 })
590 }
591
592 fn left_sibling_branch_level(&self, level: usize) -> Option<usize> {
594 self.entry[0..level].iter().rposition(|&e| e != 0)
595 }
596
597 fn right_sibling(&self, level: usize, pool: &NodePool<F>) -> Option<(F::Key, Node)> {
600 self.right_sibling_branch_level(level, pool).map(|bl| {
603 let be = usize::from(self.entry[bl]);
605 let crit_key;
606 let mut node;
607 {
608 let (keys, tree) = pool[self.node[bl]].unwrap_inner();
609 crit_key = keys[be];
610 node = tree[be + 1];
611 }
612
613 for _ in bl + 1..level {
615 node = pool[node].unwrap_inner().1[0];
616 }
617
618 (crit_key, node)
619 })
620 }
621
622 fn update_right_crit_key(&self, level: usize, crit_key: F::Key, pool: &mut NodePool<F>) {
624 let bl = self
625 .right_sibling_branch_level(level, pool)
626 .expect("No right sibling exists");
627 match pool[self.node[bl]] {
628 NodeData::Inner { ref mut keys, .. } => {
629 keys[usize::from(self.entry[bl])] = crit_key;
630 }
631 _ => panic!("Expected inner node"),
632 }
633 }
634
635 pub fn normalize(&mut self, pool: &mut NodePool<F>) {
638 if let Some((leaf, entry)) = self.leaf_pos() {
639 if entry >= pool[leaf].entries() {
640 let leaf_level = self.size - 1;
641 self.next_node(leaf_level, pool);
642 }
643 }
644 }
645}
646
647#[cfg(test)]
648impl<F: Forest> Path<F> {
649 pub fn verify(&self, pool: &NodePool<F>) {
651 for level in 0..self.size {
652 match pool[self.node[level]] {
653 NodeData::Inner { size, tree, .. } => {
654 assert!(level < self.size - 1, "Expected leaf node at level {level}");
655 assert!(
656 self.entry[level] <= size,
657 "OOB inner entry {}/{} at level {}",
658 self.entry[level],
659 size,
660 level
661 );
662 assert_eq!(
663 self.node[level + 1],
664 tree[usize::from(self.entry[level])],
665 "Node mismatch at level {level}"
666 );
667 }
668 NodeData::Leaf { size, .. } => {
669 assert_eq!(level, self.size - 1, "Expected inner node");
670 assert!(
671 self.entry[level] <= size,
672 "OOB leaf entry {}/{}",
673 self.entry[level],
674 size,
675 );
676 }
677 NodeData::Free { .. } => {
678 panic!("Free {} in path", self.node[level]);
679 }
680 }
681 }
682 }
683}
684
685#[cfg(test)]
686impl<F: Forest> fmt::Display for Path<F> {
687 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
688 if self.size == 0 {
689 write!(f, "<empty path>")
690 } else {
691 write!(f, "{}[{}]", self.node[0], self.entry[0])?;
692 for i in 1..self.size {
693 write!(f, "--{}[{}]", self.node[i], self.entry[i])?;
694 }
695 Ok(())
696 }
697 }
698}
699
700#[cfg(test)]
701mod tests {
702 use super::*;
703 use core::cmp::Ordering;
704
705 struct TC();
706
707 impl Comparator<i32> for TC {
708 fn cmp(&self, a: i32, b: i32) -> Ordering {
709 a.cmp(&b)
710 }
711 }
712
713 struct TF();
714
715 impl Forest for TF {
716 type Key = i32;
717 type Value = char;
718 type LeafKeys = [i32; 7];
719 type LeafValues = [char; 7];
720
721 fn splat_key(key: Self::Key) -> Self::LeafKeys {
722 [key; 7]
723 }
724
725 fn splat_value(value: Self::Value) -> Self::LeafValues {
726 [value; 7]
727 }
728 }
729
730 #[test]
731 fn search_single_leaf() {
732 let mut pool = NodePool::<TF>::new();
734 let root = pool.alloc_node(NodeData::leaf(10, 'a'));
735 let mut p = Path::default();
736 let comp = TC();
737
738 assert_eq!(p.find(5, root, &pool, &comp), None);
740 assert_eq!(p.size, 1);
741 assert_eq!(p.node[0], root);
742 assert_eq!(p.entry[0], 0);
743
744 assert_eq!(p.find(10, root, &pool, &comp), Some('a'));
746 assert_eq!(p.size, 1);
747 assert_eq!(p.node[0], root);
748 assert_eq!(p.entry[0], 0);
749
750 assert_eq!(p.find(15, root, &pool, &comp), None);
752 assert_eq!(p.size, 1);
753 assert_eq!(p.node[0], root);
754 assert_eq!(p.entry[0], 1);
755
756 match pool[root] {
758 NodeData::Leaf {
759 ref mut size,
760 ref mut keys,
761 ref mut vals,
762 } => {
763 *size = 2;
764 keys[1] = 20;
765 vals[1] = 'b';
766 }
767 _ => unreachable!(),
768 }
769
770 assert_eq!(p.find(15, root, &pool, &comp), None);
772 assert_eq!(p.size, 1);
773 assert_eq!(p.node[0], root);
774 assert_eq!(p.entry[0], 1);
775
776 assert_eq!(p.find(25, root, &pool, &comp), None);
778 assert_eq!(p.size, 1);
779 assert_eq!(p.node[0], root);
780 assert_eq!(p.entry[0], 2);
781 }
782
783 #[test]
784 fn search_single_inner() {
785 let mut pool = NodePool::<TF>::new();
787 let leaf1 = pool.alloc_node(NodeData::leaf(10, 'a'));
788 let leaf2 = pool.alloc_node(NodeData::leaf(20, 'b'));
789 let root = pool.alloc_node(NodeData::inner(leaf1, 20, leaf2));
790 let mut p = Path::default();
791 let comp = TC();
792
793 assert_eq!(p.find(5, root, &pool, &comp), None);
795 assert_eq!(p.size, 2);
796 assert_eq!(p.node[0], root);
797 assert_eq!(p.entry[0], 0);
798 assert_eq!(p.node[1], leaf1);
799 assert_eq!(p.entry[1], 0);
800
801 assert_eq!(p.find(10, root, &pool, &comp), Some('a'));
802 assert_eq!(p.size, 2);
803 assert_eq!(p.node[0], root);
804 assert_eq!(p.entry[0], 0);
805 assert_eq!(p.node[1], leaf1);
806 assert_eq!(p.entry[1], 0);
807
808 assert_eq!(p.find(15, root, &pool, &comp), None);
810 assert_eq!(p.size, 2);
811 assert_eq!(p.node[0], root);
812 assert_eq!(p.entry[0], 0);
813 assert_eq!(p.node[1], leaf1);
814 assert_eq!(p.entry[1], 1);
815
816 assert_eq!(p.find(20, root, &pool, &comp), Some('b'));
817 assert_eq!(p.size, 2);
818 assert_eq!(p.node[0], root);
819 assert_eq!(p.entry[0], 1);
820 assert_eq!(p.node[1], leaf2);
821 assert_eq!(p.entry[1], 0);
822
823 assert_eq!(p.find(25, root, &pool, &comp), None);
824 assert_eq!(p.size, 2);
825 assert_eq!(p.node[0], root);
826 assert_eq!(p.entry[0], 1);
827 assert_eq!(p.node[1], leaf2);
828 assert_eq!(p.entry[1], 1);
829 }
830}