Implement append and split_off for BTreeMap and BTreeSet

Changes the internal SearchStack API to return the key on removal as well.

Author: jooert

src/libcollections/btree/map.rs

@@ -348,7 +348,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 +438,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 +463,131 @@ 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) {
+        let b = other.b;
+        for (key, value) in mem::replace(other, BTreeMap::with_b(b)) {
+            self.insert(key, value);
+        }
+    }
+
+    /// 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")]
@@ -693,16 +818,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 +842,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 +853,7 @@ mod stack {
                             }
                         } else {
                             // All done!
-                            return value;
+                            return (key, value);
                         }
                     }
                 }
@@ -739,7 +864,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 +1345,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
     }
 }
 

src/libcollections/btree/node.rs

@@ -557,6 +557,15 @@ impl<K: Ord, V> Node<K, V> {
             (index, false) => GoDown(Handle { node: node, index: index, marker: PhantomData }),
         }
     }
+
+    /// Searches for the first key greater than or equal to `key`.
+    pub fn greater_or_equal<Q: ?Sized, NodeRef: Deref<Target=Node<K, V>>>(node: NodeRef, key: &Q)
+                  -> SearchResult<NodeRef> where K: Borrow<Q>, Q: Ord {
+        match node.as_slices_internal().greater_or_equal_linear(key) {
+            (index, true) => Found(Handle { node: node, index: index, marker: PhantomData }),
+            (index, false) => GoDown(Handle { node: node, index: index, marker: PhantomData }),
+        }
+    }
 }
 
 // Public interface
@@ -1609,3 +1618,18 @@ macro_rules! node_slice_impl {
 
 node_slice_impl!(NodeSlice, Traversal, as_slices_internal, index, iter);
 node_slice_impl!(MutNodeSlice, MutTraversal, as_slices_internal_mut, index_mut, iter_mut);
+
+impl<'a, K: Ord + 'a, V: 'a> NodeSlice<'a, K, V> {
+    /// Performs linear search in a slice to find the first element greater than
+    /// or equal to `key`.
+    fn greater_or_equal_linear<Q: ?Sized>(&self, key: &Q) -> (usize, bool)
+            where K: Borrow<Q>, Q: Ord  {
+        for (i, k) in self.keys.iter().enumerate() {
+            match key.cmp(k.borrow()) {
+                Greater => {}
+                Less | Equal => return (i, true),
+            }
+        }
+        (self.keys.len(), false)
+    }
+}

src/libcollections/btree/set.rs

@@ -453,6 +453,74 @@ impl<T: Ord> BTreeSet<T> {
     pub fn remove<Q: ?Sized>(&mut self, value: &Q) -> bool where T: Borrow<Q>, Q: Ord {
         self.map.remove(value).is_some()
     }
+
+    /// Moves all elements from `other` into `Self`, leaving `other` empty.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #![feature(btree_append_split_off)]
+    /// use std::collections::BTreeMap;
+    ///
+    /// let mut a = BTreeSet::new();
+    /// a.insert(1);
+    /// a.insert(2);
+    /// a.insert(3);
+    ///
+    /// let mut b = BTreeSet::new();
+    /// b.insert(3);
+    /// b.insert(4);
+    /// b.insert(5);
+    ///
+    /// a.append(&mut b);
+    ///
+    /// assert_eq!(a.len(), 5);
+    /// assert_eq!(b.len(), 0);
+    ///
+    /// assert!(a.contains(&1));
+    /// assert!(a.contains(&2));
+    /// assert!(a.contains(&3));
+    /// assert!(a.contains(&4));
+    /// assert!(a.contains(&5));
+    /// ```
+    #[unstable(feature = "btree_append_split_off",
+               reason = "recently added as part of collections reform 2")]
+    pub fn append(&mut self, other: &mut Self) {
+        self.map.append(&mut other.map);
+    }
+
+    /// Splits the set 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.insert(2)
+    /// a.insert(3)
+    /// a.insert(4)
+    /// a.insert(5)
+    ///
+    /// let b = a.split_off(3);
+    ///
+    /// assert_eq!(a.len(), 2);
+    /// assert_eq!(b.len(), 3);
+    ///
+    /// assert!(a.contains(&1));
+    /// assert!(a.contains(&2));
+    /// assert!(b.contains(&3));
+    /// assert!(b.contains(&4));
+    /// assert!(b.contains(&5));
+    /// ```
+    #[unstable(feature = "btree_append_split_off",
+               reason = "recently added as part of collections reform 2")]
+    pub fn split_off<Q: ?Sized>(&mut self, at: &Q) -> Self
+        where T: Borrow<Q>, Q: Ord {
+        BTreeSet { map: self.map.split_off(at) }
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]