alloc/collections/btree/
split.rs

1use core::alloc::Allocator;
2use core::borrow::Borrow;
3
4use super::node::ForceResult::*;
5use super::node::Root;
6use super::search::SearchResult::*;
7
8impl<K, V> Root<K, V> {
9    /// Calculates the length of both trees that result from splitting up
10    /// a given number of distinct key-value pairs.
11    pub(super) fn calc_split_length(
12        total_num: usize,
13        root_a: &Root<K, V>,
14        root_b: &Root<K, V>,
15    ) -> (usize, usize) {
16        let (length_a, length_b);
17        if root_a.height() < root_b.height() {
18            length_a = root_a.reborrow().calc_length();
19            length_b = total_num - length_a;
20            debug_assert_eq!(length_b, root_b.reborrow().calc_length());
21        } else {
22            length_b = root_b.reborrow().calc_length();
23            length_a = total_num - length_b;
24            debug_assert_eq!(length_a, root_a.reborrow().calc_length());
25        }
26        (length_a, length_b)
27    }
28
29    /// Split off a tree with key-value pairs at and after the given key.
30    /// The result is meaningful only if the tree is ordered by key,
31    /// and if the ordering of `Q` corresponds to that of `K`.
32    /// If `self` respects all `BTreeMap` tree invariants, then both
33    /// `self` and the returned tree will respect those invariants.
34    pub(super) fn split_off<Q: ?Sized + Ord, A: Allocator + Clone>(
35        &mut self,
36        key: &Q,
37        alloc: A,
38    ) -> Self
39    where
40        K: Borrow<Q>,
41    {
42        let left_root = self;
43        let mut right_root = Root::new_pillar(left_root.height(), alloc.clone());
44        let mut left_node = left_root.borrow_mut();
45        let mut right_node = right_root.borrow_mut();
46
47        loop {
48            let mut split_edge = match left_node.search_node(key) {
49                // key is going to the right tree
50                Found(kv) => kv.left_edge(),
51                GoDown(edge) => edge,
52            };
53
54            split_edge.move_suffix(&mut right_node);
55
56            match (split_edge.force(), right_node.force()) {
57                (Internal(edge), Internal(node)) => {
58                    left_node = edge.descend();
59                    right_node = node.first_edge().descend();
60                }
61                (Leaf(_), Leaf(_)) => break,
62                _ => unreachable!(),
63            }
64        }
65
66        left_root.fix_right_border(alloc.clone());
67        right_root.fix_left_border(alloc);
68        right_root
69    }
70
71    /// Creates a tree consisting of empty nodes.
72    fn new_pillar<A: Allocator + Clone>(height: usize, alloc: A) -> Self {
73        let mut root = Root::new(alloc.clone());
74        for _ in 0..height {
75            root.push_internal_level(alloc.clone());
76        }
77        root
78    }
79}