Skip to content

Commit 7849be2

Browse files
ssomersssomers
committed
BTreeMap/BTreeSet drain & retain
1 parent 582a4ea commit 7849be2

File tree

8 files changed

+555
-16
lines changed

8 files changed

+555
-16
lines changed

‎src/liballoc/benches/btree/set.rs

+50-8
Original file line numberDiff line numberDiff line change
@@ -14,19 +14,13 @@ fn random(n: usize) -> BTreeSet<usize> {
1414
}
1515

1616
fn neg(n: usize) -> BTreeSet<i32> {
17-
let mut set = BTreeSet::new();
18-
for i in -(n as i32)..=-1 {
19-
set.insert(i);
20-
}
17+
let set: BTreeSet<i32> = (-(n as i32)..=-1).collect();
2118
assert_eq!(set.len(), n);
2219
set
2320
}
2421

2522
fn pos(n: usize) -> BTreeSet<i32> {
26-
let mut set = BTreeSet::new();
27-
for i in 1..=(n as i32) {
28-
set.insert(i);
29-
}
23+
let set: BTreeSet<i32> = (1..=(n as i32)).collect();
3024
assert_eq!(set.len(), n);
3125
set
3226
}
@@ -56,6 +50,54 @@ macro_rules! set_bench {
5650
};
5751
}
5852

53+
const BUILD_SET_SIZE: usize = 100;
54+
55+
#[bench]
56+
pub fn build_and_clear(b: &mut Bencher) {
57+
b.iter(|| pos(BUILD_SET_SIZE).clear())
58+
}
59+
60+
#[bench]
61+
pub fn build_and_drain(b: &mut Bencher) {
62+
b.iter(|| pos(BUILD_SET_SIZE).drain().count())
63+
}
64+
65+
#[bench]
66+
pub fn build_and_drop(b: &mut Bencher) {
67+
b.iter(|| pos(BUILD_SET_SIZE))
68+
}
69+
70+
#[bench]
71+
pub fn build_and_into_iter(b: &mut Bencher) {
72+
b.iter(|| pos(BUILD_SET_SIZE).into_iter().count())
73+
}
74+
75+
#[bench]
76+
pub fn build_and_pop_all(b: &mut Bencher) {
77+
b.iter(|| {
78+
let mut s = pos(BUILD_SET_SIZE);
79+
while s.pop_first().is_some() {
80+
}
81+
s
82+
});
83+
}
84+
85+
#[bench]
86+
pub fn build_and_remove_all(b: &mut Bencher) {
87+
b.iter(|| {
88+
let mut s = pos(BUILD_SET_SIZE);
89+
while let Some(elt) = s.iter().copied().next() {
90+
s.remove(&elt);
91+
}
92+
s
93+
});
94+
}
95+
96+
#[bench]
97+
pub fn build_and_retain_nothing(b: &mut Bencher) {
98+
b.iter(|| pos(BUILD_SET_SIZE).retain(|_| false))
99+
}
100+
59101
set_bench! {intersection_100_neg_vs_100_pos, intersection, count, [neg(100), pos(100)]}
60102
set_bench! {intersection_100_neg_vs_10k_pos, intersection, count, [neg(100), pos(10_000)]}
61103
set_bench! {intersection_100_pos_vs_100_neg, intersection, count, [pos(100), neg(100)]}

‎src/liballoc/benches/lib.rs

+1
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
#![feature(btree_drain_retain)]
12
#![feature(map_first_last)]
23
#![feature(repr_simd)]
34
#![feature(test)]

‎src/liballoc/collections/btree/map.rs

+124-4
Original file line numberDiff line numberDiff line change
@@ -525,7 +525,39 @@ impl<K: Ord, V> BTreeMap<K, V> {
525525
}
526526
}
527527

528-
/// Clears the map, removing all values.
528+
/// Clears the map, returning all key-value pairs as an iterator.
529+
///
530+
/// # Examples
531+
///
532+
/// Basic usage:
533+
///
534+
/// ```
535+
/// #![feature(btree_drain_retain)]
536+
/// use std::collections::BTreeMap;
537+
///
538+
/// let mut a = BTreeMap::new();
539+
/// a.insert(1, "a");
540+
/// a.insert(2, "b");
541+
///
542+
/// for (k, v) in a.drain().take(1) {
543+
/// assert!(k == 1 || k == 2);
544+
/// assert!(v == "a" || v == "b");
545+
/// }
546+
///
547+
/// assert!(a.is_empty());
548+
/// ```
549+
#[unstable(feature = "btree_drain_retain", issue = "42849")]
550+
pub fn drain(&mut self) -> IntoIter<K, V> {
551+
let length = mem::replace(&mut self.length, 0);
552+
let root = mem::replace(&mut self.root, node::Root::shared_empty_root());
553+
let root1 = unsafe { ptr::read(&root).into_ref() };
554+
let root2 = unsafe { ptr::read(&root).into_ref() };
555+
let front = first_leaf_edge(root1);
556+
let back = last_leaf_edge(root2);
557+
IntoIter { front, back, length }
558+
}
559+
560+
/// Clears the map, removing all key-value pairs.
529561
///
530562
/// # Examples
531563
///
@@ -842,6 +874,47 @@ impl<K: Ord, V> BTreeMap<K, V> {
842874
}
843875
}
844876

877+
/// Retains only the elements specified by the predicate.
878+
///
879+
/// In other words, remove all pairs `(k, v)` such that `f(&k, &mut v)` returns `false`.
880+
///
881+
/// # Examples
882+
///
883+
/// ```
884+
/// #![feature(btree_drain_retain)]
885+
/// use std::collections::BTreeMap;
886+
///
887+
/// let mut map: BTreeMap<i32, i32> = (0..8).map(|x|(x, x)).collect();
888+
/// map.retain(|&k, _| k % 2 == 0);
889+
/// assert_eq!(map.len(), 4);
890+
/// ```
891+
#[unstable(feature = "btree_drain_retain", issue = "42849")]
892+
pub fn retain<F>(&mut self, mut f: F)
893+
where F: FnMut(&K, &mut V) -> bool,
894+
{
895+
let mut cur_leaf_edge = first_leaf_edge(self.root.as_mut());
896+
while let Some((kv, next_leaf_edge)) = RangeMut::next_handles(cur_leaf_edge) {
897+
let (k, v) = unsafe { ptr::read(&kv).into_kv_mut() };
898+
let retain = f(k, v);
899+
cur_leaf_edge = if retain {
900+
next_leaf_edge
901+
} else {
902+
let (k, _v, hole) = OccupiedEntry::remove_kv_by_handle(kv);
903+
// next_leaf_edge is now invalid or wrong
904+
self.length -= 1;
905+
match hole {
906+
Some(next_leaf_edge) => next_leaf_edge,
907+
None => {
908+
let root1 = self.root.as_mut();
909+
let root2 = unsafe { ptr::read(&root1) };
910+
let (front, _back) = range_search(root1, root2, k..);
911+
front
912+
}
913+
}
914+
}
915+
}
916+
}
917+
845918
/// Moves all elements from `other` into `Self`, leaving `other` empty.
846919
///
847920
/// # Examples
@@ -1804,6 +1877,41 @@ impl<'a, K, V> RangeMut<'a, K, V> {
18041877
}
18051878
}
18061879
}
1880+
1881+
// Transform given leaf edge handle into handles of next kv and next leaf edge
1882+
fn next_handles(
1883+
leaf_edge: Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>
1884+
) -> Option<(Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::KV>,
1885+
Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>)> {
1886+
let mut cur_edge = match leaf_edge.right_kv() {
1887+
Ok(kv) => {
1888+
let next_leaf_edge = unsafe { ptr::read(&kv).right_edge() };
1889+
return Some((kv.forget_node_type(), next_leaf_edge));
1890+
}
1891+
Err(last_edge) => {
1892+
match last_edge.into_node().ascend() {
1893+
Ok(next_level) => next_level,
1894+
Err(_) => return None,
1895+
}
1896+
}
1897+
};
1898+
1899+
loop {
1900+
cur_edge = match cur_edge.right_kv() {
1901+
Ok(kv) => {
1902+
let right_edge = unsafe { ptr::read(&kv).right_edge() };
1903+
let next_leaf_edge = first_leaf_edge(right_edge.descend());
1904+
return Some((kv.forget_node_type(), next_leaf_edge));
1905+
}
1906+
Err(last_edge) => {
1907+
match last_edge.into_node().ascend() {
1908+
Ok(next_level) => next_level,
1909+
Err(_) => return None,
1910+
}
1911+
}
1912+
}
1913+
}
1914+
}
18071915
}
18081916

18091917
#[stable(feature = "btree_range", since = "1.17.0")]
@@ -2618,9 +2726,21 @@ impl<'a, K: Ord, V> OccupiedEntry<'a, K, V> {
26182726
fn remove_kv(self) -> (K, V) {
26192727
*self.length -= 1;
26202728

2621-
let (small_leaf, old_key, old_val) = match self.handle.force() {
2729+
let (k, v, _) = OccupiedEntry::remove_kv_by_handle(self.handle);
2730+
(k, v)
2731+
}
2732+
2733+
// Removes and returns a key/value-pair, and optionally (if cheap), returns the resulting edge
2734+
// corresponding to the former left and right edges of the entry.
2735+
fn remove_kv_by_handle(
2736+
handle: Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::KV>)
2737+
-> (K, V, Option<Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>>) {
2738+
let (small_leaf, old_key, old_val) = match handle.force() {
26222739
Leaf(leaf) => {
26232740
let (hole, old_key, old_val) = leaf.remove();
2741+
if hole.reborrow().into_node().len() >= node::MIN_LEN {
2742+
return (old_key, old_val, Some(hole))
2743+
}
26242744
(hole.into_node(), old_key, old_val)
26252745
}
26262746
Internal(mut internal) => {
@@ -2641,7 +2761,7 @@ impl<'a, K: Ord, V> OccupiedEntry<'a, K, V> {
26412761

26422762
// Handle underflow
26432763
let mut cur_node = small_leaf.forget_type();
2644-
while cur_node.len() < node::CAPACITY / 2 {
2764+
while cur_node.len() < node::MIN_LEN {
26452765
match handle_underfull_node(cur_node) {
26462766
AtRoot => break,
26472767
EmptyParent(_) => unreachable!(),
@@ -2658,7 +2778,7 @@ impl<'a, K: Ord, V> OccupiedEntry<'a, K, V> {
26582778
}
26592779
}
26602780

2661-
(old_key, old_val)
2781+
(old_key, old_val, None)
26622782
}
26632783
}
26642784

‎src/liballoc/collections/btree/node.rs

+20-4
Original file line numberDiff line numberDiff line change
@@ -1245,8 +1245,8 @@ impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::KV>
12451245
}
12461246
}
12471247

1248-
/// Removes the key/value pair pointed to by this handle, returning the edge between the
1249-
/// now adjacent key/value pairs to the left and right of this handle.
1248+
/// Removes the key/value pair pointed to by this handle and returns it, along with the edge
1249+
/// between the now adjacent key/value pairs (if any) to the left and right of this handle.
12501250
pub fn remove(mut self)
12511251
-> (Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>, K, V) {
12521252
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::
14051405
}
14061406
}
14071407

1408-
/// This removes a key/value pair from the left child and replaces it with the key/value pair
1408+
/// This removes a key/value pair from the left child and places it in the key/value storage
14091409
/// pointed to by this handle while pushing the old key/value pair of this handle into the right
14101410
/// child.
14111411
pub fn steal_left(&mut self) {
@@ -1422,7 +1422,7 @@ impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::
14221422
}
14231423
}
14241424

1425-
/// This removes a key/value pair from the right child and replaces it with the key/value pair
1425+
/// This removes a key/value pair from the right child and places it in the key/value storage
14261426
/// pointed to by this handle while pushing the old key/value pair of this handle into the left
14271427
/// child.
14281428
pub fn steal_right(&mut self) {
@@ -1589,6 +1589,22 @@ unsafe fn move_edges<K, V>(
15891589
dest.correct_childrens_parent_links(dest_offset, dest_offset + count);
15901590
}
15911591

1592+
impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::KV> {
1593+
#[unstable(feature = "btree_drain_retain", issue = "42849")]
1594+
pub fn forget_node_type(self)
1595+
-> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV> {
1596+
Handle::new_kv(self.node.forget_type(), self.idx)
1597+
}
1598+
}
1599+
1600+
impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::KV> {
1601+
#[unstable(feature = "btree_drain_retain", issue = "42849")]
1602+
pub fn forget_node_type(self)
1603+
-> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV> {
1604+
Handle::new_kv(self.node.forget_type(), self.idx)
1605+
}
1606+
}
1607+
15921608
impl<BorrowType, K, V, HandleType>
15931609
Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, HandleType> {
15941610

‎src/liballoc/collections/btree/set.rs

+46
Original file line numberDiff line numberDiff line change
@@ -515,6 +515,31 @@ impl<T: Ord> BTreeSet<T> {
515515
Union(MergeIterInner::new(self.iter(), other.iter()))
516516
}
517517

518+
/// Clears the set, returning all values as an iterator.
519+
///
520+
/// # Examples
521+
///
522+
/// Basic usage:
523+
///
524+
/// ```
525+
/// #![feature(btree_drain_retain)]
526+
/// use std::collections::BTreeSet;
527+
///
528+
/// let mut a = BTreeSet::new();
529+
/// a.insert(1);
530+
/// a.insert(2);
531+
///
532+
/// for k in a.drain().take(1) {
533+
/// assert!(k == 1 || k == 2);
534+
/// }
535+
///
536+
/// assert!(a.is_empty());
537+
/// ```
538+
#[unstable(feature = "btree_drain_retain", issue = "42849")]
539+
pub fn drain(&mut self) -> IntoIter<T> {
540+
IntoIter { iter: self.map.drain() }
541+
}
542+
518543
/// Clears the set, removing all values.
519544
///
520545
/// # Examples
@@ -884,6 +909,27 @@ impl<T: Ord> BTreeSet<T> {
884909
Recover::take(&mut self.map, value)
885910
}
886911

912+
/// Retains only the value specified by the predicate.
913+
///
914+
/// In other words, remove all value `v` such that `f(&v)` returns `false`.
915+
///
916+
/// # Examples
917+
///
918+
/// ```
919+
/// #![feature(btree_drain_retain)]
920+
/// use std::collections::BTreeSet;
921+
///
922+
/// let mut set: BTreeSet<i32> = (0..8).collect();
923+
/// set.retain(|&v| v % 2 == 0);
924+
/// assert_eq!(set.len(), 4);
925+
/// ```
926+
#[unstable(feature = "btree_drain_retain", issue = "42849")]
927+
pub fn retain<F>(&mut self, mut f: F)
928+
where F: FnMut(&T) -> bool,
929+
{
930+
self.map.retain(|k, _| f(k));
931+
}
932+
887933
/// Moves all elements from `other` into `Self`, leaving `other` empty.
888934
///
889935
/// # Examples

0 commit comments

Comments
 (0)