BTreeMap/BTreeSet drain & retain

Author: ssomers

src/liballoc/benches/btree/set.rs

@@ -14,19 +14,13 @@ fn random(n: usize) -> BTreeSet<usize> {
 }
 
 fn neg(n: usize) -> BTreeSet<i32> {
-    let mut set = BTreeSet::new();
-    for i in -(n as i32)..=-1 {
-        set.insert(i);
-    }
+    let set: BTreeSet<i32> = (-(n as i32)..=-1).collect();
     assert_eq!(set.len(), n);
     set
 }
 
 fn pos(n: usize) -> BTreeSet<i32> {
-    let mut set = BTreeSet::new();
-    for i in 1..=(n as i32) {
-        set.insert(i);
-    }
+    let set: BTreeSet<i32> = (1..=(n as i32)).collect();
     assert_eq!(set.len(), n);
     set
 }
@@ -56,6 +50,54 @@ macro_rules! set_bench {
     };
 }
 
+const BUILD_SET_SIZE: usize = 100;
+
+#[bench]
+pub fn build_and_clear(b: &mut Bencher) {
+    b.iter(|| pos(BUILD_SET_SIZE).clear())
+}
+
+#[bench]
+pub fn build_and_drain(b: &mut Bencher) {
+    b.iter(|| pos(BUILD_SET_SIZE).drain().count())
+}
+
+#[bench]
+pub fn build_and_drop(b: &mut Bencher) {
+    b.iter(|| pos(BUILD_SET_SIZE))
+}
+
+#[bench]
+pub fn build_and_into_iter(b: &mut Bencher) {
+    b.iter(|| pos(BUILD_SET_SIZE).into_iter().count())
+}
+
+#[bench]
+pub fn build_and_pop_all(b: &mut Bencher) {
+    b.iter(|| {
+        let mut s = pos(BUILD_SET_SIZE);
+        while s.pop_first().is_some() {
+        }
+        s
+    });
+}
+
+#[bench]
+pub fn build_and_remove_all(b: &mut Bencher) {
+    b.iter(|| {
+        let mut s = pos(BUILD_SET_SIZE);
+        while let Some(elt) = s.iter().copied().next() {
+            s.remove(&elt);
+        }
+        s
+    });
+}
+
+#[bench]
+pub fn build_and_retain_nothing(b: &mut Bencher) {
+    b.iter(|| pos(BUILD_SET_SIZE).retain(|_| false))
+}
+
 set_bench! {intersection_100_neg_vs_100_pos, intersection, count, [neg(100), pos(100)]}
 set_bench! {intersection_100_neg_vs_10k_pos, intersection, count, [neg(100), pos(10_000)]}
 set_bench! {intersection_100_pos_vs_100_neg, intersection, count, [pos(100), neg(100)]}

src/liballoc/benches/lib.rs

@@ -1,3 +1,4 @@
+#![feature(btree_drain_retain)]
 #![feature(map_first_last)]
 #![feature(repr_simd)]
 #![feature(test)]

src/liballoc/collections/btree/map.rs

@@ -525,7 +525,39 @@ impl<K: Ord, V> BTreeMap<K, V> {
         }
     }
 
-    /// Clears the map, removing all values.
+    /// Clears the map, returning all key-value pairs as an iterator.
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// #![feature(btree_drain_retain)]
+    /// use std::collections::BTreeMap;
+    ///
+    /// let mut a = BTreeMap::new();
+    /// a.insert(1, "a");
+    /// a.insert(2, "b");
+    ///
+    /// for (k, v) in a.drain().take(1) {
+    ///     assert!(k == 1 || k == 2);
+    ///     assert!(v == "a" || v == "b");
+    /// }
+    ///
+    /// assert!(a.is_empty());
+    /// ```
+    #[unstable(feature = "btree_drain_retain", issue = "42849")]
+    pub fn drain(&mut self) -> IntoIter<K, V> {
+        let length = mem::replace(&mut self.length, 0);
+        let root = mem::replace(&mut self.root, node::Root::shared_empty_root());
+        let root1 = unsafe { ptr::read(&root).into_ref() };
+        let root2 = unsafe { ptr::read(&root).into_ref() };
+        let front = first_leaf_edge(root1);
+        let back = last_leaf_edge(root2);
+        IntoIter { front, back, length }
+    }
+
+    /// Clears the map, removing all key-value pairs.
     ///
     /// # Examples
     ///
@@ -842,6 +874,47 @@ impl<K: Ord, V> BTreeMap<K, V> {
         }
     }
 
+    /// Retains only the elements specified by the predicate.
+    ///
+    /// In other words, remove all pairs `(k, v)` such that `f(&k, &mut v)` returns `false`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(btree_drain_retain)]
+    /// use std::collections::BTreeMap;
+    ///
+    /// let mut map: BTreeMap<i32, i32> = (0..8).map(|x|(x, x)).collect();
+    /// map.retain(|&k, _| k % 2 == 0);
+    /// assert_eq!(map.len(), 4);
+    /// ```
+    #[unstable(feature = "btree_drain_retain", issue = "42849")]
+    pub fn retain<F>(&mut self, mut f: F)
+        where F: FnMut(&K, &mut V) -> bool,
+    {
+        let mut cur_leaf_edge = first_leaf_edge(self.root.as_mut());
+        while let Some((kv, next_leaf_edge)) = RangeMut::next_handles(cur_leaf_edge) {
+            let (k, v) = unsafe { ptr::read(&kv).into_kv_mut() };
+            let retain = f(k, v);
+            cur_leaf_edge = if retain {
+                next_leaf_edge
+            } else {
+                let (k, _v, hole) = OccupiedEntry::remove_kv_by_handle(kv);
+                // next_leaf_edge is now invalid or wrong
+                self.length -= 1;
+                match hole {
+                    Some(next_leaf_edge) => next_leaf_edge,
+                    None => {
+                        let root1 = self.root.as_mut();
+                        let root2 = unsafe { ptr::read(&root1) };
+                        let (front, _back) = range_search(root1, root2, k..);
+                        front
+                    }
+                }
+            }
+        }
+    }
+
     /// Moves all elements from `other` into `Self`, leaving `other` empty.
     ///
     /// # Examples
@@ -1804,6 +1877,41 @@ impl<'a, K, V> RangeMut<'a, K, V> {
             }
         }
     }
+
+    // Transform given leaf edge handle into handles of next kv and next leaf edge
+    fn next_handles(
+        leaf_edge: Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>
+        ) -> Option<(Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::KV>,
+                     Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>)> {
+        let mut cur_edge = match leaf_edge.right_kv() {
+            Ok(kv) => {
+                let next_leaf_edge = unsafe { ptr::read(&kv).right_edge() };
+                return Some((kv.forget_node_type(), next_leaf_edge));
+            }
+            Err(last_edge) => {
+                match last_edge.into_node().ascend() {
+                    Ok(next_level) => next_level,
+                    Err(_) => return None,
+                }
+            }
+        };
+
+        loop {
+            cur_edge = match cur_edge.right_kv() {
+                Ok(kv) => {
+                    let right_edge = unsafe { ptr::read(&kv).right_edge() };
+                    let next_leaf_edge = first_leaf_edge(right_edge.descend());
+                    return Some((kv.forget_node_type(), next_leaf_edge));
+                }
+                Err(last_edge) => {
+                    match last_edge.into_node().ascend() {
+                        Ok(next_level) => next_level,
+                        Err(_) => return None,
+                    }
+                }
+            }
+        }
+    }
 }
 
 #[stable(feature = "btree_range", since = "1.17.0")]
@@ -2618,9 +2726,21 @@ impl<'a, K: Ord, V> OccupiedEntry<'a, K, V> {
     fn remove_kv(self) -> (K, V) {
         *self.length -= 1;
 
-        let (small_leaf, old_key, old_val) = match self.handle.force() {
+        let (k, v, _) = OccupiedEntry::remove_kv_by_handle(self.handle);
+        (k, v)
+    }
+
+    // Removes and returns a key/value-pair, and optionally (if cheap), returns the resulting edge
+    // corresponding to the former left and right edges of the entry.
+    fn remove_kv_by_handle(
+            handle: Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::KV>)
+            -> (K, V, Option<Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>>) {
+        let (small_leaf, old_key, old_val) = match handle.force() {
             Leaf(leaf) => {
                 let (hole, old_key, old_val) = leaf.remove();
+                if hole.reborrow().into_node().len() >= node::MIN_LEN {
+                    return (old_key, old_val, Some(hole))
+                }
                 (hole.into_node(), old_key, old_val)
             }
             Internal(mut internal) => {
@@ -2641,7 +2761,7 @@ impl<'a, K: Ord, V> OccupiedEntry<'a, K, V> {
 
         // Handle underflow
         let mut cur_node = small_leaf.forget_type();
-        while cur_node.len() < node::CAPACITY / 2 {
+        while cur_node.len() < node::MIN_LEN {
             match handle_underfull_node(cur_node) {
                 AtRoot => break,
                 EmptyParent(_) => unreachable!(),
@@ -2658,7 +2778,7 @@ impl<'a, K: Ord, V> OccupiedEntry<'a, K, V> {
             }
         }
 
-        (old_key, old_val)
+        (old_key, old_val, None)
     }
 }
 

src/liballoc/collections/btree/node.rs

@@ -1245,8 +1245,8 @@ impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::KV>
         }
     }
 
-    /// Removes the key/value pair pointed to by this handle, returning the edge between the
-    /// now adjacent key/value pairs to the left and right of this handle.
+    /// Removes the key/value pair pointed to by this handle and returns it, along with the edge
+    /// between the now adjacent key/value pairs (if any) to the left and right of this handle.
     pub fn remove(mut self)
             -> (Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>, K, V) {
         debug_assert!(!self.node.is_shared_root());
@@ -1405,7 +1405,7 @@ impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::
         }
     }
 
-    /// This removes a key/value pair from the left child and replaces it with the key/value pair
+    /// This removes a key/value pair from the left child and places it in the key/value storage
     /// pointed to by this handle while pushing the old key/value pair of this handle into the right
     /// child.
     pub fn steal_left(&mut self) {
@@ -1422,7 +1422,7 @@ impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::
         }
     }
 
-    /// This removes a key/value pair from the right child and replaces it with the key/value pair
+    /// This removes a key/value pair from the right child and places it in the key/value storage
     /// pointed to by this handle while pushing the old key/value pair of this handle into the left
     /// child.
     pub fn steal_right(&mut self) {
@@ -1589,6 +1589,22 @@ unsafe fn move_edges<K, V>(
     dest.correct_childrens_parent_links(dest_offset, dest_offset + count);
 }
 
+impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::KV> {
+    #[unstable(feature = "btree_drain_retain", issue = "42849")]
+    pub fn forget_node_type(self)
+        -> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV> {
+        Handle::new_kv(self.node.forget_type(), self.idx)
+    }
+}
+
+impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::KV> {
+    #[unstable(feature = "btree_drain_retain", issue = "42849")]
+    pub fn forget_node_type(self)
+        -> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV> {
+        Handle::new_kv(self.node.forget_type(), self.idx)
+    }
+}
+
 impl<BorrowType, K, V, HandleType>
         Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, HandleType> {
 

src/liballoc/collections/btree/set.rs

@@ -515,6 +515,31 @@ impl<T: Ord> BTreeSet<T> {
         Union(MergeIterInner::new(self.iter(), other.iter()))
     }
 
+    /// Clears the set, returning all values as an iterator.
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// #![feature(btree_drain_retain)]
+    /// use std::collections::BTreeSet;
+    ///
+    /// let mut a = BTreeSet::new();
+    /// a.insert(1);
+    /// a.insert(2);
+    ///
+    /// for k in a.drain().take(1) {
+    ///     assert!(k == 1 || k == 2);
+    /// }
+    ///
+    /// assert!(a.is_empty());
+    /// ```
+    #[unstable(feature = "btree_drain_retain", issue = "42849")]
+    pub fn drain(&mut self) -> IntoIter<T> {
+        IntoIter { iter: self.map.drain() }
+    }
+
     /// Clears the set, removing all values.
     ///
     /// # Examples
@@ -884,6 +909,27 @@ impl<T: Ord> BTreeSet<T> {
         Recover::take(&mut self.map, value)
     }
 
+    /// Retains only the value specified by the predicate.
+    ///
+    /// In other words, remove all value `v` such that `f(&v)` returns `false`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(btree_drain_retain)]
+    /// use std::collections::BTreeSet;
+    ///
+    /// let mut set: BTreeSet<i32> = (0..8).collect();
+    /// set.retain(|&v| v % 2 == 0);
+    /// assert_eq!(set.len(), 4);
+    /// ```
+    #[unstable(feature = "btree_drain_retain", issue = "42849")]
+    pub fn retain<F>(&mut self, mut f: F)
+        where F: FnMut(&T) -> bool,
+    {
+        self.map.retain(|k, _| f(k));
+    }
+
     /// Moves all elements from `other` into `Self`, leaving `other` empty.
     ///
     /// # Examples