Auto merge of #33947 - xosmig:btree_split_off, r=gereeter

Implement split_off for BTreeMap and BTreeSet (RFC 509)

Fixes #19986 and refactors common with append methods.
It splits the tree with O(log n) operations and then calculates sizes by traversing the lower one.

CC @gereeter

4 files changed, 127 insertions, 0 deletions

diff --git a/src/libcollectionstest/btree/map.rs b/src/libcollectionstest/btree/map.rs
index e19090c7599..49fce68d15e 100644
--- a/src/libcollectionstest/btree/map.rs
+++ b/src/libcollectionstest/btree/map.rs
@@ -13,6 +13,9 @@ use std::collections::Bound::{self, Excluded, Included, Unbounded};
 use std::collections::btree_map::Entry::{Occupied, Vacant};
 use std::rc::Rc;
 
+use std::iter::FromIterator;
+use super::DeterministicRng;
+
 #[test]
 fn test_basic_large() {
     let mut map = BTreeMap::new();
@@ -528,6 +531,51 @@ create_append_test!(test_append_181, 181);
 create_append_test!(test_append_239, 239);
 create_append_test!(test_append_1700, 1700);
 
+fn rand_data(len: usize) -> Vec<(u32, u32)> {
+    let mut rng = DeterministicRng::new();
+    Vec::from_iter(
+        (0..len).map(|_| (rng.next(), rng.next()))
+    )
+}
+
+#[test]
+fn test_split_off_empty_right() {
+    let mut data = rand_data(173);
+
+    let mut map = BTreeMap::from_iter(data.clone());
+    let right = map.split_off(&(data.iter().max().unwrap().0 + 1));
+
+    data.sort();
+    assert!(map.into_iter().eq(data));
+    assert!(right.into_iter().eq(None));
+}
+
+#[test]
+fn test_split_off_empty_left() {
+    let mut data = rand_data(314);
+
+    let mut map = BTreeMap::from_iter(data.clone());
+    let right = map.split_off(&data.iter().min().unwrap().0);
+
+    data.sort();
+    assert!(map.into_iter().eq(None));
+    assert!(right.into_iter().eq(data));
+}
+
+#[test]
+fn test_split_off_large_random_sorted() {
+    let mut data = rand_data(1529);
+    // special case with maximum height.
+    data.sort();
+
+    let mut map = BTreeMap::from_iter(data.clone());
+    let key = data[data.len() / 2].0;
+    let right = map.split_off(&key);
+
+    assert!(map.into_iter().eq(data.clone().into_iter().filter(|x| x.0 < key)));
+    assert!(right.into_iter().eq(data.into_iter().filter(|x| x.0 >= key)));
+}
+
 mod bench {
     use std::collections::BTreeMap;
     use std::__rand::{Rng, thread_rng};
diff --git a/src/libcollectionstest/btree/mod.rs b/src/libcollectionstest/btree/mod.rs
index 0db48f3ce9e..ea43f423b7c 100644
--- a/src/libcollectionstest/btree/mod.rs
+++ b/src/libcollectionstest/btree/mod.rs
@@ -10,3 +10,33 @@
 
 mod map;
 mod set;
+
+/// XorShiftRng
+struct DeterministicRng {
+    x: u32,
+    y: u32,
+    z: u32,
+    w: u32,
+}
+
+impl DeterministicRng {
+    fn new() -> Self {
+        DeterministicRng {
+            x: 0x193a6754,
+            y: 0xa8a7d469,
+            z: 0x97830e05,
+            w: 0x113ba7bb
+        }
+    }
+
+    fn next(&mut self) -> u32 {
+        let x = self.x;
+        let t = x ^ (x << 11);
+        self.x = self.y;
+        self.y = self.z;
+        self.z = self.w;
+        let w_ = self.w;
+        self.w = w_ ^ (w_ >> 19) ^ (t ^ (t >> 8));
+        self.w
+    }
+}
diff --git a/src/libcollectionstest/btree/set.rs b/src/libcollectionstest/btree/set.rs
index 53ccfd5b4e2..f7b647d7772 100644
--- a/src/libcollectionstest/btree/set.rs
+++ b/src/libcollectionstest/btree/set.rs
@@ -10,6 +10,9 @@
 
 use std::collections::BTreeSet;
 
+use std::iter::FromIterator;
+use super::DeterministicRng;
+
 #[test]
 fn test_clone_eq() {
   let mut m = BTreeSet::new();
@@ -289,3 +292,48 @@ fn test_append() {
     assert_eq!(a.contains(&4), true);
     assert_eq!(a.contains(&5), true);
 }
+
+fn rand_data(len: usize) -> Vec<u32> {
+    let mut rng = DeterministicRng::new();
+    Vec::from_iter(
+        (0..len).map(|_| rng.next())
+    )
+}
+
+#[test]
+fn test_split_off_empty_right() {
+    let mut data = rand_data(173);
+
+    let mut set = BTreeSet::from_iter(data.clone());
+    let right = set.split_off(&(data.iter().max().unwrap() + 1));
+
+    data.sort();
+    assert!(set.into_iter().eq(data));
+    assert!(right.into_iter().eq(None));
+}
+
+#[test]
+fn test_split_off_empty_left() {
+    let mut data = rand_data(314);
+
+    let mut set = BTreeSet::from_iter(data.clone());
+    let right = set.split_off(data.iter().min().unwrap());
+
+    data.sort();
+    assert!(set.into_iter().eq(None));
+    assert!(right.into_iter().eq(data));
+}
+
+#[test]
+fn test_split_off_large_random_sorted() {
+    let mut data = rand_data(1529);
+    // special case with maximum height.
+    data.sort();
+
+    let mut set = BTreeSet::from_iter(data.clone());
+    let key = data[data.len() / 2];
+    let right = set.split_off(&key);
+
+    assert!(set.into_iter().eq(data.clone().into_iter().filter(|x| *x < key)));
+    assert!(right.into_iter().eq(data.into_iter().filter(|x| *x >= key)));
+}
diff --git a/src/libcollectionstest/lib.rs b/src/libcollectionstest/lib.rs
index e9754bc98aa..400d6140948 100644
--- a/src/libcollectionstest/lib.rs
+++ b/src/libcollectionstest/lib.rs
@@ -14,6 +14,7 @@
 #![feature(binary_heap_append)]
 #![feature(box_syntax)]
 #![feature(btree_append)]
+#![feature(btree_split_off)]
 #![feature(btree_range)]
 #![feature(collections)]
 #![feature(collections_bound)]