Implement append and split_off for BTreeMap and BTreeSet

src/libcollections/btree/map.rs

@@ -148,6 +148,13 @@ pub struct OccupiedEntry<'a, K:'a, V:'a> {
     stack: stack::SearchStack<'a, K, V, node::handle::KV, node::handle::LeafOrInternal>,
 }
 
+struct MergeIter<K, V, I: Iterator<Item=(K, V)>> {
+    left: I,
+    right: I,
+    left_cur: Option<(K, V)>,
+    right_cur: Option<(K, V)>,
+}
+
 impl<K: Ord, V> BTreeMap<K, V> {
     /// Makes a new empty BTreeMap with a reasonable choice for B.
     #[stable(feature = "rust1", since = "1.0.0")]
@@ -348,7 +355,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
 
         loop {
             let result = stack.with(move |pusher, node| {
-                // Same basic logic as found in `find`, but with PartialSearchStack mediating the
+                // Same basic logic as found in `get`, but with PartialSearchStack mediating the
                 // actual nodes for us
                 match Node::search(node, &key) {
                     Found(mut handle) => {
@@ -438,14 +445,14 @@ impl<K: Ord, V> BTreeMap<K, V> {
     /// ```
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn remove<Q: ?Sized>(&mut self, key: &Q) -> Option<V> where K: Borrow<Q>, Q: Ord {
-        // See `swap` for a more thorough description of the stuff going on in here
+        // See `insert` for a more thorough description of the stuff going on in here
         let mut stack = stack::PartialSearchStack::new(self);
         loop {
             let result = stack.with(move |pusher, node| {
                 match Node::search(node, key) {
                     Found(handle) => {
                         // Perfect match. Terminate the stack here, and remove the entry
-                        Finished(Some(pusher.seal(handle).remove()))
+                        Finished(Some(pusher.seal(handle).remove().1))
                     },
                     GoDown(handle) => {
                         // We need to keep searching, try to go down the next edge
@@ -463,6 +470,157 @@ impl<K: Ord, V> BTreeMap<K, V> {
             }
         }
     }
+
+    /// Moves all elements from `other` into `Self`, leaving `other` empty.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #![feature(btree_append_split_off)]
+    /// use std::collections::BTreeMap;
+    ///
+    /// let mut a = BTreeMap::new();
+    /// a.insert(1, "a");
+    /// a.insert(2, "b");
+    /// a.insert(3, "c");
+    ///
+    /// let mut b = BTreeMap::new();
+    /// b.insert(3, "d");
+    /// b.insert(4, "e");
+    /// b.insert(5, "f");
+    ///
+    /// a.append(&mut b);
+    ///
+    /// assert_eq!(a.len(), 5);
+    /// assert_eq!(b.len(), 0);
+    ///
+    /// assert_eq!(a[&1], "a");
+    /// assert_eq!(a[&2], "b");
+    /// assert_eq!(a[&3], "d");
+    /// assert_eq!(a[&4], "e");
+    /// assert_eq!(a[&5], "f");
+    /// ```
+    #[unstable(feature = "append",
+               reason = "recently added as part of collections reform 2")]
+    pub fn append(&mut self, other: &mut Self) {
+        // Do we have to append anything at all?
+        if other.len() == 0 {
+            return;
+        }
+
+        // If the values of `b` of `self` and `other` are equal, we can just swap them if `self` is
+        // empty.
+        if self.len() == 0 && self.b == other.b {
+            mem::swap(self, other);
+        }
+
+        // First, we merge `self` and `other` into a sorted sequence in linear time.
+        let self_b = self.b;
+        let other_b = other.b;
+        let mut self_iter = mem::replace(self, BTreeMap::with_b(self_b)).into_iter();
+        let mut other_iter = mem::replace(other, BTreeMap::with_b(other_b)).into_iter();
+        let self_cur = self_iter.next();
+        let other_cur = other_iter.next();
+
+        // Second, we build a tree from the sorted sequence in linear time.
+        let (length, depth, root) = Node::from_sorted_iter(MergeIter {
+            left: self_iter,
+            right: other_iter,
+            left_cur: self_cur,
+            right_cur: other_cur,
+        }, self_b);
+
+        self.length = length;
+        self.depth = depth;
+        self.root = root.unwrap(); // `unwrap` won't panic because length can't be zero.
+    }
+
+    /// Splits the map into two at the given key,
+    /// retaining the first half in-place and returning the second one.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #![feature(btree_append_split_off)]
+    /// use std::collections::BTreeMap;
+    ///
+    /// a.insert(1, "a");
+    /// a.insert(2, "b");
+    /// a.insert(3, "c");
+    /// a.insert(4, "d");
+    /// a.insert(5, "e");
+    ///
+    /// let b = a.split_off(3);
+    ///
+    /// assert_eq!(a.len(), 2);
+    /// assert_eq!(b.len(), 3);
+    ///
+    /// assert_eq!(a[&1], "a");
+    /// assert_eq!(a[&2], "b");
+    /// assert_eq!(b[&3], "c");
+    /// assert_eq!(b[&4], "d");
+    /// assert_eq!(b[&5], "e");
+    /// ```
+    #[unstable(feature = "split_off",
+               reason = "recently added as part of collections reform 2")]
+    pub fn split_off<Q: ?Sized>(&mut self, at: &Q) -> Self where K: Borrow<Q>, Q: Ord {
+        let mut other = BTreeMap::new();
+
+        if self.len() == 0 {
+            return other;
+        }
+
+        // FIXME(RFC #811) We can't check for `at` pointing before the
+        // first element and then swap `self` and `other`, because
+        // `self` will still be borrowed immutably.
+        // `unwrap` won't panic because `self.len()` > 0.
+        let should_swap = at <= self.keys().next().unwrap().borrow();
+
+        // Does `at` point before the first element?
+        if should_swap {
+            mem::swap(self, &mut other);
+            return other;
+        }
+        // Does `at` point behind the last element?
+        // `unwrap` won't panic because `self.len()` > 0.
+        else if at > self.keys().rev().next().unwrap().borrow() {
+            return other;
+        }
+
+        let mut remove_greater_or_equal = || {
+            let mut stack = stack::PartialSearchStack::new(self);
+            loop {
+                let result = stack.with(move |pusher, node| {
+                    match Node::greater_or_equal(node, at) {
+                        Found(handle) => {
+                            // Found a matching key. Terminate the stack here, and remove the entry
+                            Finished(Some(pusher.seal(handle).remove()))
+                        },
+                        GoDown(handle) => {
+                            // We need to keep searching, try to go down the next edge
+                            match handle.force() {
+                                // We're at a leaf; no matching key found
+                                Leaf(_) => Finished(None),
+                                Internal(internal_handle) => Continue(pusher.push(internal_handle))
+                            }
+                        }
+                    }
+                });
+                match result {
+                    Finished(ret) => return ret,
+                    Continue(new_stack) => stack = new_stack
+                }
+            }
+        };
+
+        // Remove and move all elements greater than or equal to at
+        loop {
+            match remove_greater_or_equal() {
+                Some((key, value)) => other.insert(key, value),
+                None => return other,
+            };
+        }
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
@@ -519,6 +677,56 @@ impl<'a, K, V> IntoIterator for &'a mut BTreeMap<K, V> {
     }
 }
 
+// Helper enum for MergeIter
+enum MergeResult {
+    Left,
+    Right,
+    Both,
+    None,
+}
+
+impl<K: Ord, V, I: Iterator<Item=(K, V)>> Iterator for MergeIter<K, V, I> {
+    type Item = (K, V);
+
+    fn next(&mut self) -> Option<(K, V)> {
+        let res = match (&self.left_cur, &self.right_cur) {
+            (&Some((ref left_key, _)), &Some((ref right_key, _))) => {
+                match left_key.cmp(right_key) {
+                    Ordering::Less => MergeResult::Left,
+                    Ordering::Equal => MergeResult::Both,
+                    Ordering::Greater => MergeResult::Right,
+                }
+            },
+            (&Some(_), &None) => MergeResult::Left,
+            (&None, &Some(_)) => MergeResult::Right,
+            (&None, &None) => MergeResult::None,
+        };
+
+        // Check which elements comes first and only advance the corresponding iterator.
+        // If two keys are equal, take the value from `right`.
+        match res {
+            MergeResult::Left => {
+                let ret = self.left_cur.take();
+                self.left_cur = self.left.next();
+                ret
+            },
+            MergeResult::Right => {
+                let ret = self.right_cur.take();
+                self.right_cur = self.right.next();
+                ret
+            },
+            MergeResult::Both => {
+                let ret = self.right_cur.take();
+                self.left_cur = self.left.next();
+                self.right_cur = self.right.next();
+                ret
+            },
+            MergeResult::None => None,
+        }
+    }
+}
+
+
 /// A helper enum useful for deciding whether to continue a loop since we can't
 /// return from a closure
 enum Continuation<A, B> {
@@ -693,16 +901,16 @@ mod stack {
     impl<'a, K, V> SearchStack<'a, K, V, handle::KV, handle::Leaf> {
         /// Removes the key and value in the top element of the stack, then handles underflows as
         /// described in BTree's pop function.
-        fn remove_leaf(mut self) -> V {
+        fn remove_leaf(mut self) -> (K, V) {
             self.map.length -= 1;
 
             // Remove the key-value pair from the leaf that this search stack points to.
             // Then, note if the leaf is underfull, and promptly forget the leaf and its ptr
             // to avoid ownership issues.
-            let (value, mut underflow) = unsafe {
-                let (_, value) = self.top.from_raw_mut().remove_as_leaf();
+            let (key, value, mut underflow) = unsafe {
+                let (key, value) = self.top.from_raw_mut().remove_as_leaf();
                 let underflow = self.top.from_raw().node().is_underfull();
-                (value, underflow)
+                (key, value, underflow)
             };
 
             loop {
@@ -717,7 +925,7 @@ mod stack {
                             self.map.depth -= 1;
                             self.map.root.hoist_lone_child();
                         }
-                        return value;
+                        return (key, value);
                     }
                     Some(mut handle) => {
                         if underflow {
@@ -728,7 +936,7 @@ mod stack {
                             }
                         } else {
                             // All done!
-                            return value;
+                            return (key, value);
                         }
                     }
                 }
@@ -739,7 +947,7 @@ mod stack {
     impl<'a, K, V> SearchStack<'a, K, V, handle::KV, handle::LeafOrInternal> {
         /// Removes the key and value in the top element of the stack, then handles underflows as
         /// described in BTree's pop function.
-        pub fn remove(self) -> V {
+        pub fn remove(self) -> (K, V) {
             // Ensure that the search stack goes to a leaf. This is necessary to perform deletion
             // in a BTree. Note that this may put the tree in an inconsistent state (further
             // described in into_leaf's comments), but this is immediately fixed by the
@@ -1220,7 +1428,7 @@ impl<'a, K: Ord, V> OccupiedEntry<'a, K, V> {
     /// Takes the value of the entry out of the map, and returns it.
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn remove(self) -> V {
-        self.stack.remove()
+        self.stack.remove().1
     }
 }