library/alloc/src/collections/btree/split.rs

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