1use super::node::Removed;
4use super::{Comparator, Forest, MAX_PATH, Node, NodeData, NodePool, slice_insert, slice_shift};
5use core::borrow::Borrow;
6use core::marker::PhantomData;
7use wasmtime_core::error::OutOfMemory;
8
9#[cfg(test)]
10use core::fmt;
11
12pub(super) struct Path<F: Forest> {
13 size: usize,
15
16 node: [Node; MAX_PATH],
18
19 entry: [u8; MAX_PATH],
21
22 unused: PhantomData<F>,
23}
24
25impl<F: Forest> Default for Path<F> {
26 fn default() -> Self {
27 Self {
28 size: 0,
29 node: [Node(0); MAX_PATH],
30 entry: [0; MAX_PATH],
31 unused: PhantomData,
32 }
33 }
34}
35
36impl<F: Forest> Path<F> {
37 pub fn find(
49 &mut self,
50 key: F::Key,
51 root: Node,
52 pool: &NodePool<F>,
53 comp: &dyn Comparator<F::Key>,
54 ) -> Option<F::Value> {
55 let mut node = root;
56 for level in 0.. {
57 self.size = level + 1;
58 self.node[level] = node;
59 match pool[node] {
60 NodeData::Inner { size, keys, tree } => {
61 let i = match comp.search(key, &keys[0..size.into()]) {
64 Ok(i) => i + 1,
66 Err(i) => i,
68 };
69 self.entry[level] = i as u8;
70 node = tree[i];
71 }
72 NodeData::Leaf { size, keys, vals } => {
73 return match comp.search(key, &keys.borrow()[0..size.into()]) {
75 Ok(i) => {
76 self.entry[level] = i as u8;
77 Some(vals.borrow()[i])
78 }
79 Err(i) => {
80 self.entry[level] = i as u8;
81 None
82 }
83 };
84 }
85 NodeData::Free { .. } => panic!("Free {} reached from {}", node, root),
86 }
87 }
88 unreachable!();
89 }
90
91 pub fn first(&mut self, root: Node, pool: &NodePool<F>) -> (F::Key, F::Value) {
93 let mut node = root;
94 for level in 0.. {
95 self.size = level + 1;
96 self.node[level] = node;
97 self.entry[level] = 0;
98 match pool[node] {
99 NodeData::Inner { tree, .. } => node = tree[0],
100 NodeData::Leaf { keys, vals, .. } => return (keys.borrow()[0], vals.borrow()[0]),
101 NodeData::Free { .. } => panic!("Free {} reached from {}", node, root),
102 }
103 }
104 unreachable!();
105 }
106
107 pub fn next(&mut self, pool: &NodePool<F>) -> Option<(F::Key, F::Value)> {
109 match self.leaf_pos() {
110 None => return None,
111 Some((node, entry)) => {
112 let (keys, vals) = pool[node].unwrap_leaf();
113 if entry + 1 < keys.len() {
114 self.entry[self.size - 1] += 1;
115 return Some((keys[entry + 1], vals[entry + 1]));
116 }
117 }
118 }
119
120 let leaf_level = self.size - 1;
122 self.next_node(leaf_level, pool).map(|node| {
123 let (keys, vals) = pool[node].unwrap_leaf();
124 (keys[0], vals[0])
125 })
126 }
127
128 pub fn prev(&mut self, root: Node, pool: &NodePool<F>) -> Option<(F::Key, F::Value)> {
134 if self.size == 0 {
136 self.goto_subtree_last(0, root, pool);
137 let (node, entry) = self.leaf_pos().unwrap();
138 let (keys, vals) = pool[node].unwrap_leaf();
139 return Some((keys[entry], vals[entry]));
140 }
141
142 match self.leaf_pos() {
143 None => return None,
144 Some((node, entry)) => {
145 if entry > 0 {
146 self.entry[self.size - 1] -= 1;
147 let (keys, vals) = pool[node].unwrap_leaf();
148 return Some((keys[entry - 1], vals[entry - 1]));
149 }
150 }
151 }
152
153 self.prev_leaf(pool).map(|node| {
155 let (keys, vals) = pool[node].unwrap_leaf();
156 let e = self.leaf_entry();
157 (keys[e], vals[e])
158 })
159 }
160
161 fn next_node(&mut self, level: usize, pool: &NodePool<F>) -> Option<Node> {
167 match self.right_sibling_branch_level(level, pool) {
168 None => {
169 self.size = 0;
170 None
171 }
172 Some(bl) => {
173 let (_, bnodes) = pool[self.node[bl]].unwrap_inner();
174 self.entry[bl] += 1;
175 let mut node = bnodes[usize::from(self.entry[bl])];
176
177 for l in bl + 1..level {
178 self.node[l] = node;
179 self.entry[l] = 0;
180 node = pool[node].unwrap_inner().1[0];
181 }
182
183 self.node[level] = node;
184 self.entry[level] = 0;
185 Some(node)
186 }
187 }
188 }
189
190 fn prev_leaf(&mut self, pool: &NodePool<F>) -> Option<Node> {
196 self.left_sibling_branch_level(self.size - 1).map(|bl| {
197 let entry = self.entry[bl] - 1;
198 self.entry[bl] = entry;
199 let (_, bnodes) = pool[self.node[bl]].unwrap_inner();
200 self.goto_subtree_last(bl + 1, bnodes[usize::from(entry)], pool)
201 })
202 }
203
204 fn goto_subtree_last(&mut self, level: usize, root: Node, pool: &NodePool<F>) -> Node {
206 let mut node = root;
207 for l in level.. {
208 self.node[l] = node;
209 match pool[node] {
210 NodeData::Inner { size, ref tree, .. } => {
211 self.entry[l] = size;
212 node = tree[usize::from(size)];
213 }
214 NodeData::Leaf { size, .. } => {
215 self.entry[l] = size - 1;
216 self.size = l + 1;
217 break;
218 }
219 NodeData::Free { .. } => panic!("Free {} reached from {}", node, root),
220 }
221 }
222 node
223 }
224
225 pub fn set_root_node(&mut self, root: Node) {
227 self.size = 1;
228 self.node[0] = root;
229 self.entry[0] = 0;
230 }
231
232 pub fn leaf_pos(&self) -> Option<(Node, usize)> {
234 let i = self.size.wrapping_sub(1);
235 self.node.get(i).map(|&n| (n, self.entry[i].into()))
236 }
237
238 fn leaf_node(&self) -> Node {
240 self.node[self.size - 1]
241 }
242
243 fn leaf_entry(&self) -> usize {
245 self.entry[self.size - 1].into()
246 }
247
248 fn at_first_entry(&self) -> bool {
251 self.entry[0..self.size].iter().all(|&i| i == 0)
252 }
253
254 pub fn value_mut<'a>(&self, pool: &'a mut NodePool<F>) -> &'a mut F::Value {
257 &mut pool[self.leaf_node()].unwrap_leaf_mut().1[self.leaf_entry()]
258 }
259
260 pub fn insert(
265 &mut self,
266 key: F::Key,
267 value: F::Value,
268 pool: &mut NodePool<F>,
269 ) -> Result<Node, OutOfMemory> {
270 if !self.try_leaf_insert(key, value, pool) {
271 self.split_and_insert(key, value, pool)?;
272 }
273 Ok(self.node[0])
274 }
275
276 fn try_leaf_insert(&self, key: F::Key, value: F::Value, pool: &mut NodePool<F>) -> bool {
279 let index = self.leaf_entry();
280
281 debug_assert!(index > 0 || self.at_first_entry());
285
286 pool[self.leaf_node()].try_leaf_insert(index, key, value)
287 }
288
289 fn split_and_insert(
292 &mut self,
293 mut key: F::Key,
294 value: F::Value,
295 pool: &mut NodePool<F>,
296 ) -> Result<(), OutOfMemory> {
297 let orig_root = self.node[0];
298
299 let mut ins_node = None;
302 let mut split;
303 for level in (0..self.size).rev() {
304 let mut node = self.node[level];
306 let mut entry = self.entry[level].into();
307 split = pool[node].split(entry);
308 let rhs_node = pool.alloc_node(split.rhs_data)?;
309
310 if entry > split.lhs_entries
318 || (entry == split.lhs_entries
319 && (split.lhs_entries > split.rhs_entries || ins_node.is_some()))
320 {
321 node = rhs_node;
322 entry -= split.lhs_entries;
323 self.node[level] = node;
324 self.entry[level] = entry as u8;
325 }
326
327 match ins_node {
329 None => {
330 let inserted = pool[node].try_leaf_insert(entry, key, value);
331 debug_assert!(inserted);
332 if entry == 0 && node == rhs_node {
335 split.crit_key = key;
336 }
337 }
338 Some(n) => {
339 let inserted = pool[node].try_inner_insert(entry, key, n);
340 debug_assert!(inserted);
341 if n == self.node[level + 1] {
344 self.entry[level] += 1;
345 }
346 }
347 }
348
349 key = split.crit_key;
352 ins_node = Some(rhs_node);
353 if level > 0 {
354 let pnode = &mut pool[self.node[level - 1]];
355 let pentry = self.entry[level - 1].into();
356 if pnode.try_inner_insert(pentry, key, rhs_node) {
357 if node == rhs_node {
359 self.entry[level - 1] += 1;
360 }
361 return Ok(());
362 }
363 }
364 }
365
366 let rhs_node = ins_node.expect("empty path");
368 let root = pool.alloc_node(NodeData::inner(orig_root, key, rhs_node))?;
369 let entry = if self.node[0] == rhs_node { 1 } else { 0 };
370 self.size += 1;
371 slice_insert(&mut self.node[0..self.size], 0, root);
372 slice_insert(&mut self.entry[0..self.size], 0, entry);
373 Ok(())
374 }
375
376 pub fn remove(&mut self, pool: &mut NodePool<F>) -> Option<Node> {
381 let e = self.leaf_entry();
382 match pool[self.leaf_node()].leaf_remove(e) {
383 Removed::Healthy => {
384 if e == 0 {
385 self.update_crit_key(pool)
386 }
387 Some(self.node[0])
388 }
389 status => self.balance_nodes(status, pool),
390 }
391 }
392
393 fn current_crit_key(&self, level: usize, pool: &NodePool<F>) -> Option<F::Key> {
400 self.left_sibling_branch_level(level).map(|bl| {
402 let (keys, _) = pool[self.node[bl]].unwrap_inner();
403 keys[usize::from(self.entry[bl]) - 1]
404 })
405 }
406
407 fn update_crit_key(&mut self, pool: &mut NodePool<F>) {
409 let crit_level = match self.left_sibling_branch_level(self.size - 1) {
411 None => return,
412 Some(l) => l,
413 };
414 let crit_kidx = self.entry[crit_level] - 1;
415
416 let crit_key = pool[self.leaf_node()].leaf_crit_key();
418 let crit_node = self.node[crit_level];
419
420 match pool[crit_node] {
421 NodeData::Inner {
422 size, ref mut keys, ..
423 } => {
424 debug_assert!(crit_kidx < size);
425 keys[usize::from(crit_kidx)] = crit_key;
426 }
427 _ => panic!("Expected inner node"),
428 }
429 }
430
431 fn balance_nodes(&mut self, status: Removed, pool: &mut NodePool<F>) -> Option<Node> {
436 if status != Removed::Empty && self.leaf_entry() == 0 {
441 self.update_crit_key(pool);
442 }
443
444 let leaf_level = self.size - 1;
445 if self.heal_level(status, leaf_level, pool) {
446 self.size = 0;
448 return None;
449 }
450
451 let mut ns = 0;
453 while let NodeData::Inner {
454 size: 0, ref tree, ..
455 } = pool[self.node[ns]]
456 {
457 ns += 1;
458 self.node[ns] = tree[0];
459 }
460
461 if ns > 0 {
462 for l in 0..ns {
463 pool.free_node(self.node[l]);
464 }
465
466 slice_shift(&mut self.node, ns);
469 slice_shift(&mut self.entry, ns);
470
471 if self.size > 0 {
472 self.size -= ns;
473 }
474 }
475
476 Some(self.node[0])
479 }
480
481 fn heal_level(&mut self, status: Removed, level: usize, pool: &mut NodePool<F>) -> bool {
488 match status {
489 Removed::Healthy => {}
490 Removed::Rightmost => {
491 debug_assert_eq!(
494 usize::from(self.entry[level]),
495 pool[self.node[level]].entries()
496 );
497 self.next_node(level, pool);
498 }
499 Removed::Underflow => self.underflowed_node(level, pool),
500 Removed::Empty => return self.empty_node(level, pool),
501 }
502 false
503 }
504
505 fn underflowed_node(&mut self, level: usize, pool: &mut NodePool<F>) {
512 if let Some((crit_key, rhs_node)) = self.right_sibling(level, pool) {
515 let new_ck: Option<F::Key>;
517 let empty;
518 let mut rhs = pool[rhs_node];
520 match pool[self.node[level]].balance(crit_key, &mut rhs) {
521 None => {
522 new_ck = self.current_crit_key(level, pool);
524 empty = true;
525 }
526 Some(key) => {
527 new_ck = Some(key);
529 empty = false;
530 }
531 }
532 pool[rhs_node] = rhs;
534 if let Some(ck) = new_ck {
537 self.update_right_crit_key(level, ck, pool);
538 }
539 if empty {
540 let empty_tree = self.empty_node(level, pool);
541 debug_assert!(!empty_tree);
542 }
543
544 debug_assert!(usize::from(self.entry[level]) < pool[self.node[level]].entries());
548 } else if usize::from(self.entry[level]) >= pool[self.node[level]].entries() {
549 self.size = 0;
552 }
553 }
554
555 fn empty_node(&mut self, level: usize, pool: &mut NodePool<F>) -> bool {
564 pool.free_node(self.node[level]);
565 if level == 0 {
566 return true;
568 }
569
570 let rhs_node = self.right_sibling(level, pool).map(|(_, n)| n);
572
573 let pl = level - 1;
575 let pe = self.entry[pl].into();
576 let status = pool[self.node[pl]].inner_remove(pe);
577 self.heal_level(status, pl, pool);
578
579 match rhs_node {
581 Some(rhs) => self.node[level] = rhs,
584 None => self.size = 0,
587 }
588 false
589 }
590
591 fn right_sibling_branch_level(&self, level: usize, pool: &NodePool<F>) -> Option<usize> {
598 (0..level).rposition(|l| match pool[self.node[l]] {
599 NodeData::Inner { size, .. } => self.entry[l] < size,
600 _ => panic!("Expected inner node"),
601 })
602 }
603
604 fn left_sibling_branch_level(&self, level: usize) -> Option<usize> {
606 self.entry[0..level].iter().rposition(|&e| e != 0)
607 }
608
609 fn right_sibling(&self, level: usize, pool: &NodePool<F>) -> Option<(F::Key, Node)> {
612 self.right_sibling_branch_level(level, pool).map(|bl| {
615 let be = usize::from(self.entry[bl]);
617 let crit_key;
618 let mut node;
619 {
620 let (keys, tree) = pool[self.node[bl]].unwrap_inner();
621 crit_key = keys[be];
622 node = tree[be + 1];
623 }
624
625 for _ in bl + 1..level {
627 node = pool[node].unwrap_inner().1[0];
628 }
629
630 (crit_key, node)
631 })
632 }
633
634 fn update_right_crit_key(&self, level: usize, crit_key: F::Key, pool: &mut NodePool<F>) {
636 let bl = self
637 .right_sibling_branch_level(level, pool)
638 .expect("No right sibling exists");
639 match pool[self.node[bl]] {
640 NodeData::Inner { ref mut keys, .. } => {
641 keys[usize::from(self.entry[bl])] = crit_key;
642 }
643 _ => panic!("Expected inner node"),
644 }
645 }
646
647 pub fn normalize(&mut self, pool: &NodePool<F>) {
650 if let Some((leaf, entry)) = self.leaf_pos() {
651 if entry >= pool[leaf].entries() {
652 let leaf_level = self.size - 1;
653 self.next_node(leaf_level, pool);
654 }
655 }
656 }
657}
658
659#[cfg(test)]
660impl<F: Forest> Path<F> {
661 pub fn verify(&self, pool: &NodePool<F>) {
663 for level in 0..self.size {
664 match pool[self.node[level]] {
665 NodeData::Inner { size, tree, .. } => {
666 assert!(level < self.size - 1, "Expected leaf node at level {level}");
667 assert!(
668 self.entry[level] <= size,
669 "OOB inner entry {}/{} at level {}",
670 self.entry[level],
671 size,
672 level
673 );
674 assert_eq!(
675 self.node[level + 1],
676 tree[usize::from(self.entry[level])],
677 "Node mismatch at level {level}"
678 );
679 }
680 NodeData::Leaf { size, .. } => {
681 assert_eq!(level, self.size - 1, "Expected inner node");
682 assert!(
683 self.entry[level] <= size,
684 "OOB leaf entry {}/{}",
685 self.entry[level],
686 size,
687 );
688 }
689 NodeData::Free { .. } => {
690 panic!("Free {} in path", self.node[level]);
691 }
692 }
693 }
694 }
695}
696
697#[cfg(test)]
698impl<F: Forest> fmt::Display for Path<F> {
699 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
700 if self.size == 0 {
701 write!(f, "<empty path>")
702 } else {
703 write!(f, "{}[{}]", self.node[0], self.entry[0])?;
704 for i in 1..self.size {
705 write!(f, "--{}[{}]", self.node[i], self.entry[i])?;
706 }
707 Ok(())
708 }
709 }
710}
711
712#[cfg(test)]
713mod tests {
714 use wasmtime_core::alloc::PanicOnOom;
715
716 use super::*;
717 use core::cmp::Ordering;
718
719 struct TC();
720
721 impl Comparator<i32> for TC {
722 fn cmp(&self, a: i32, b: i32) -> Ordering {
723 a.cmp(&b)
724 }
725 }
726
727 struct TF();
728
729 impl Forest for TF {
730 type Key = i32;
731 type Value = char;
732 type LeafKeys = [i32; 7];
733 type LeafValues = [char; 7];
734
735 fn splat_key(key: Self::Key) -> Self::LeafKeys {
736 [key; 7]
737 }
738
739 fn splat_value(value: Self::Value) -> Self::LeafValues {
740 [value; 7]
741 }
742 }
743
744 #[test]
745 fn search_single_leaf() {
746 let mut pool = NodePool::<TF>::new();
748 let root = pool.alloc_node(NodeData::leaf(10, 'a')).panic_on_oom();
749 let mut p = Path::default();
750 let comp = TC();
751
752 assert_eq!(p.find(5, root, &pool, &comp), None);
754 assert_eq!(p.size, 1);
755 assert_eq!(p.node[0], root);
756 assert_eq!(p.entry[0], 0);
757
758 assert_eq!(p.find(10, root, &pool, &comp), Some('a'));
760 assert_eq!(p.size, 1);
761 assert_eq!(p.node[0], root);
762 assert_eq!(p.entry[0], 0);
763
764 assert_eq!(p.find(15, root, &pool, &comp), None);
766 assert_eq!(p.size, 1);
767 assert_eq!(p.node[0], root);
768 assert_eq!(p.entry[0], 1);
769
770 match pool[root] {
772 NodeData::Leaf {
773 ref mut size,
774 ref mut keys,
775 ref mut vals,
776 } => {
777 *size = 2;
778 keys[1] = 20;
779 vals[1] = 'b';
780 }
781 _ => unreachable!(),
782 }
783
784 assert_eq!(p.find(15, root, &pool, &comp), None);
786 assert_eq!(p.size, 1);
787 assert_eq!(p.node[0], root);
788 assert_eq!(p.entry[0], 1);
789
790 assert_eq!(p.find(25, root, &pool, &comp), None);
792 assert_eq!(p.size, 1);
793 assert_eq!(p.node[0], root);
794 assert_eq!(p.entry[0], 2);
795 }
796
797 #[test]
798 fn search_single_inner() {
799 let mut pool = NodePool::<TF>::new();
801 let leaf1 = pool.alloc_node(NodeData::leaf(10, 'a')).panic_on_oom();
802 let leaf2 = pool.alloc_node(NodeData::leaf(20, 'b')).panic_on_oom();
803 let root = pool
804 .alloc_node(NodeData::inner(leaf1, 20, leaf2))
805 .panic_on_oom();
806 let mut p = Path::default();
807 let comp = TC();
808
809 assert_eq!(p.find(5, root, &pool, &comp), None);
811 assert_eq!(p.size, 2);
812 assert_eq!(p.node[0], root);
813 assert_eq!(p.entry[0], 0);
814 assert_eq!(p.node[1], leaf1);
815 assert_eq!(p.entry[1], 0);
816
817 assert_eq!(p.find(10, root, &pool, &comp), Some('a'));
818 assert_eq!(p.size, 2);
819 assert_eq!(p.node[0], root);
820 assert_eq!(p.entry[0], 0);
821 assert_eq!(p.node[1], leaf1);
822 assert_eq!(p.entry[1], 0);
823
824 assert_eq!(p.find(15, root, &pool, &comp), None);
826 assert_eq!(p.size, 2);
827 assert_eq!(p.node[0], root);
828 assert_eq!(p.entry[0], 0);
829 assert_eq!(p.node[1], leaf1);
830 assert_eq!(p.entry[1], 1);
831
832 assert_eq!(p.find(20, root, &pool, &comp), Some('b'));
833 assert_eq!(p.size, 2);
834 assert_eq!(p.node[0], root);
835 assert_eq!(p.entry[0], 1);
836 assert_eq!(p.node[1], leaf2);
837 assert_eq!(p.entry[1], 0);
838
839 assert_eq!(p.find(25, root, &pool, &comp), None);
840 assert_eq!(p.size, 2);
841 assert_eq!(p.node[0], root);
842 assert_eq!(p.entry[0], 1);
843 assert_eq!(p.node[1], leaf2);
844 assert_eq!(p.entry[1], 1);
845 }
846}