diff --git a/src/liballoc/benches/btree/map.rs b/src/liballoc/benches/btree/map.rs index eb5f51d9adc58..ea69769279f12 100644 --- a/src/liballoc/benches/btree/map.rs +++ b/src/liballoc/benches/btree/map.rs @@ -146,6 +146,36 @@ pub fn iter_100000(b: &mut Bencher) { bench_iter(b, 100000); } +fn bench_iter_mut(b: &mut Bencher, size: i32) { + let mut map = BTreeMap::::new(); + let mut rng = thread_rng(); + + for _ in 0..size { + map.insert(rng.gen(), rng.gen()); + } + + b.iter(|| { + for kv in map.iter_mut() { + black_box(kv); + } + }); +} + +#[bench] +pub fn iter_mut_20(b: &mut Bencher) { + bench_iter_mut(b, 20); +} + +#[bench] +pub fn iter_mut_1000(b: &mut Bencher) { + bench_iter_mut(b, 1000); +} + +#[bench] +pub fn iter_mut_100000(b: &mut Bencher) { + bench_iter_mut(b, 100000); +} + fn bench_first_and_last(b: &mut Bencher, size: i32) { let map: BTreeMap<_, _> = (0..size).map(|i| (i, i)).collect(); b.iter(|| { diff --git a/src/liballoc/benches/btree/set.rs b/src/liballoc/benches/btree/set.rs index 18502ded308c7..3f0c3cc3d7d71 100644 --- a/src/liballoc/benches/btree/set.rs +++ b/src/liballoc/benches/btree/set.rs @@ -14,19 +14,13 @@ fn random(n: usize) -> BTreeSet { } fn neg(n: usize) -> BTreeSet { - let mut set = BTreeSet::new(); - for i in -(n as i32)..=-1 { - set.insert(i); - } + let set: BTreeSet = (-(n as i32)..=-1).collect(); assert_eq!(set.len(), n); set } fn pos(n: usize) -> BTreeSet { - let mut set = BTreeSet::new(); - for i in 1..=(n as i32) { - set.insert(i); - } + let set: BTreeSet = (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)]} diff --git a/src/liballoc/benches/lib.rs b/src/liballoc/benches/lib.rs index 9acda886064ce..3bb11eb790ebe 100644 --- a/src/liballoc/benches/lib.rs +++ b/src/liballoc/benches/lib.rs @@ -1,3 +1,4 @@ +#![feature(btree_drain_retain)] #![feature(map_first_last)] #![feature(repr_simd)] #![feature(test)] diff --git a/src/liballoc/collections/btree/map.rs b/src/liballoc/collections/btree/map.rs index 5b48b594ff907..5898b40ce64a5 100644 --- a/src/liballoc/collections/btree/map.rs +++ b/src/liballoc/collections/btree/map.rs @@ -496,6 +496,50 @@ impl Debug for OccupiedEntry<'_, K, V> { } } +macro_rules! tree_hopper { + (unsafe fn $name:ident : $next_kv:ident $next_edge:ident $initial_leaf_edge:ident) => { + /// Position the given leaf edge handle on the next leaf edge, and return the KV in between. + /// Caller must ensure that there is another leaf edge in the given direction. + /// It would be more elegant to pass in the handle by value and to return an optional tuple, + /// but either of those causes a 10% performance loss. + unsafe fn $name + (leaf_edge: &mut Handle, marker::Edge>) + -> Handle, marker::KV> { + let mut cur_handle = match ptr::read(leaf_edge).$next_kv() { + Ok(leaf_kv) => { + *leaf_edge = ptr::read(&leaf_kv).$next_edge(); + return leaf_kv.forget_node_type(); + } + Err(last_edge) => { + match last_edge.into_node().ascend() { + Ok(next_level) => next_level, + Err(_) => unreachable!(), + } + } + }; + + loop { + cur_handle = match ptr::read(&cur_handle).$next_kv() { + Ok(internal_kv) => { + let next_internal_edge = ptr::read(&internal_kv).$next_edge(); + *leaf_edge = $initial_leaf_edge(next_internal_edge.descend()); + return internal_kv.forget_node_type(); + } + Err(last_edge) => { + match last_edge.into_node().ascend() { + Ok(next_level) => next_level, + Err(_) => unreachable!(), + } + } + } + } + } + }; +} + +tree_hopper! {unsafe fn leftward_next_unchecked : left_kv left_edge last_leaf_edge} +tree_hopper! {unsafe fn rightward_next_unchecked : right_kv right_edge first_leaf_edge} + // An iterator for merging two sorted sequences into one struct MergeIter> { left: Peekable, @@ -525,7 +569,39 @@ impl BTreeMap { } } - /// 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 { + 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 +918,51 @@ impl BTreeMap { } } + /// 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 = (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(&mut self, mut f: F) + where F: FnMut(&K, &mut V) -> bool, + { + let mut cur_leaf_edge = first_leaf_edge(self.root.as_mut()); + let mut n_retained = 0; + while n_retained < self.length { + let kv = unsafe { rightward_next_unchecked(&mut cur_leaf_edge) }; + // Doing the descend (or perhaps another move) invalidates the references + // returned by `into_kv_mut`, so we have to do this last. + let (k, v) = unsafe { ptr::read(&kv).into_kv_mut() }; + let retain = f(k, v); + if retain { + n_retained += 1; + } else { + self.length -= 1; + let (k, _v, hole) = OccupiedEntry::remove_kv_by_handle(kv); + // cur_leaf_edge may now be invalid + cur_leaf_edge = 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 @@ -1449,33 +1570,10 @@ impl Iterator for IntoIter { self.length -= 1; } - let handle = unsafe { ptr::read(&self.front) }; - - let mut cur_handle = match handle.right_kv() { - Ok(kv) => { - let k = unsafe { ptr::read(kv.reborrow().into_kv().0) }; - let v = unsafe { ptr::read(kv.reborrow().into_kv().1) }; - self.front = kv.right_edge(); - return Some((k, v)); - } - Err(last_edge) => unsafe { - unwrap_unchecked(last_edge.into_node().deallocate_and_ascend()) - }, - }; - - loop { - match cur_handle.right_kv() { - Ok(kv) => { - let k = unsafe { ptr::read(kv.reborrow().into_kv().0) }; - let v = unsafe { ptr::read(kv.reborrow().into_kv().1) }; - self.front = first_leaf_edge(kv.right_edge().descend()); - return Some((k, v)); - } - Err(last_edge) => unsafe { - cur_handle = unwrap_unchecked(last_edge.into_node().deallocate_and_ascend()); - }, - } - } + let kv = unsafe { rightward_next_unchecked(&mut self.front) }; + let k = unsafe { ptr::read(kv.reborrow().into_kv().0) }; + let v = unsafe { ptr::read(kv.reborrow().into_kv().1) }; + Some((k, v)) } fn size_hint(&self) -> (usize, Option) { @@ -1492,33 +1590,10 @@ impl DoubleEndedIterator for IntoIter { self.length -= 1; } - let handle = unsafe { ptr::read(&self.back) }; - - let mut cur_handle = match handle.left_kv() { - Ok(kv) => { - let k = unsafe { ptr::read(kv.reborrow().into_kv().0) }; - let v = unsafe { ptr::read(kv.reborrow().into_kv().1) }; - self.back = kv.left_edge(); - return Some((k, v)); - } - Err(last_edge) => unsafe { - unwrap_unchecked(last_edge.into_node().deallocate_and_ascend()) - }, - }; - - loop { - match cur_handle.left_kv() { - Ok(kv) => { - let k = unsafe { ptr::read(kv.reborrow().into_kv().0) }; - let v = unsafe { ptr::read(kv.reborrow().into_kv().1) }; - self.back = last_leaf_edge(kv.left_edge().descend()); - return Some((k, v)); - } - Err(last_edge) => unsafe { - cur_handle = unwrap_unchecked(last_edge.into_node().deallocate_and_ascend()); - }, - } - } + let kv = unsafe { leftward_next_unchecked(&mut self.back) }; + let k = unsafe { ptr::read(kv.reborrow().into_kv().0) }; + let v = unsafe { ptr::read(kv.reborrow().into_kv().1) }; + Some((k, v)) } } @@ -1667,33 +1742,8 @@ impl FusedIterator for ValuesMut<'_, K, V> {} impl<'a, K, V> Range<'a, K, V> { unsafe fn next_unchecked(&mut self) -> (&'a K, &'a V) { - let handle = self.front; - - let mut cur_handle = match handle.right_kv() { - Ok(kv) => { - let ret = kv.into_kv(); - self.front = kv.right_edge(); - return ret; - } - Err(last_edge) => { - let next_level = last_edge.into_node().ascend().ok(); - unwrap_unchecked(next_level) - } - }; - - loop { - match cur_handle.right_kv() { - Ok(kv) => { - let ret = kv.into_kv(); - self.front = first_leaf_edge(kv.right_edge().descend()); - return ret; - } - Err(last_edge) => { - let next_level = last_edge.into_node().ascend().ok(); - cur_handle = unwrap_unchecked(next_level); - } - } - } + let kv = rightward_next_unchecked(&mut self.front); + kv.into_kv() } } @@ -1710,33 +1760,8 @@ impl<'a, K, V> DoubleEndedIterator for Range<'a, K, V> { impl<'a, K, V> Range<'a, K, V> { unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a V) { - let handle = self.back; - - let mut cur_handle = match handle.left_kv() { - Ok(kv) => { - let ret = kv.into_kv(); - self.back = kv.left_edge(); - return ret; - } - Err(last_edge) => { - let next_level = last_edge.into_node().ascend().ok(); - unwrap_unchecked(next_level) - } - }; - - loop { - match cur_handle.left_kv() { - Ok(kv) => { - let ret = kv.into_kv(); - self.back = last_leaf_edge(kv.left_edge().descend()); - return ret; - } - Err(last_edge) => { - let next_level = last_edge.into_node().ascend().ok(); - cur_handle = unwrap_unchecked(next_level); - } - } - } + let kv = leftward_next_unchecked(&mut self.back); + kv.into_kv() } } @@ -1772,37 +1797,11 @@ impl<'a, K, V> Iterator for RangeMut<'a, K, V> { impl<'a, K, V> RangeMut<'a, K, V> { unsafe fn next_unchecked(&mut self) -> (&'a K, &'a mut V) { - let handle = ptr::read(&self.front); - - let mut cur_handle = match handle.right_kv() { - Ok(kv) => { - self.front = ptr::read(&kv).right_edge(); - // Doing the descend invalidates the references returned by `into_kv_mut`, - // so we have to do this last. - let (k, v) = kv.into_kv_mut(); - return (k, v); // coerce k from `&mut K` to `&K` - } - Err(last_edge) => { - let next_level = last_edge.into_node().ascend().ok(); - unwrap_unchecked(next_level) - } - }; - - loop { - match cur_handle.right_kv() { - Ok(kv) => { - self.front = first_leaf_edge(ptr::read(&kv).right_edge().descend()); - // Doing the descend invalidates the references returned by `into_kv_mut`, - // so we have to do this last. - let (k, v) = kv.into_kv_mut(); - return (k, v); // coerce k from `&mut K` to `&K` - } - Err(last_edge) => { - let next_level = last_edge.into_node().ascend().ok(); - cur_handle = unwrap_unchecked(next_level); - } - } - } + let kv = rightward_next_unchecked(&mut self.front); + // Doing the descend (or perhaps another move) invalidates the references + // returned by `into_kv_mut`, so we have to do this last. + let (k, v) = kv.into_kv_mut(); + (k, v) // coerce k from `&mut K` to `&K` } } @@ -1822,37 +1821,11 @@ impl FusedIterator for RangeMut<'_, K, V> {} impl<'a, K, V> RangeMut<'a, K, V> { unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a mut V) { - let handle = ptr::read(&self.back); - - let mut cur_handle = match handle.left_kv() { - Ok(kv) => { - self.back = ptr::read(&kv).left_edge(); - // Doing the descend invalidates the references returned by `into_kv_mut`, - // so we have to do this last. - let (k, v) = kv.into_kv_mut(); - return (k, v); // coerce k from `&mut K` to `&K` - } - Err(last_edge) => { - let next_level = last_edge.into_node().ascend().ok(); - unwrap_unchecked(next_level) - } - }; - - loop { - match cur_handle.left_kv() { - Ok(kv) => { - self.back = last_leaf_edge(ptr::read(&kv).left_edge().descend()); - // Doing the descend invalidates the references returned by `into_kv_mut`, - // so we have to do this last. - let (k, v) = kv.into_kv_mut(); - return (k, v); // coerce k from `&mut K` to `&K` - } - Err(last_edge) => { - let next_level = last_edge.into_node().ascend().ok(); - cur_handle = unwrap_unchecked(next_level); - } - } - } + let kv = leftward_next_unchecked(&mut self.back); + // Doing the descend (or perhaps another move) invalidates the references + // returned by `into_kv_mut`, so we have to do this last. + let (k, v) = kv.into_kv_mut(); + (k, v) // coerce k from `&mut K` to `&K` } } @@ -2354,7 +2327,6 @@ impl<'a, K: Ord, V: Default> Entry<'a, K, V> { Vacant(entry) => entry.insert(Default::default()), } } - } impl<'a, K: Ord, V> VacantEntry<'a, K, V> { @@ -2618,9 +2590,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, K, V, marker::LeafOrInternal>, marker::KV>) + -> (K, V, Option, 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 +2625,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 +2642,7 @@ impl<'a, K: Ord, V> OccupiedEntry<'a, K, V> { } } - (old_key, old_val) + (old_key, old_val, None) } } diff --git a/src/liballoc/collections/btree/node.rs b/src/liballoc/collections/btree/node.rs index ab010b35f6ad1..886c597d63e07 100644 --- a/src/liballoc/collections/btree/node.rs +++ b/src/liballoc/collections/btree/node.rs @@ -1245,8 +1245,8 @@ impl<'a, K, V> Handle, 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, K, V, marker::Leaf>, marker::Edge>, K, V) { debug_assert!(!self.node.is_shared_root()); @@ -1405,7 +1405,7 @@ impl<'a, K, V> Handle, 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, 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( dest.correct_childrens_parent_links(dest_offset, dest_offset + count); } +impl Handle, marker::KV> { + #[unstable(feature = "btree_drain_retain", issue = "42849")] + pub fn forget_node_type(self) + -> Handle, marker::KV> { + Handle::new_kv(self.node.forget_type(), self.idx) + } +} + +impl Handle, marker::KV> { + #[unstable(feature = "btree_drain_retain", issue = "42849")] + pub fn forget_node_type(self) + -> Handle, marker::KV> { + Handle::new_kv(self.node.forget_type(), self.idx) + } +} + impl Handle, HandleType> { diff --git a/src/liballoc/collections/btree/set.rs b/src/liballoc/collections/btree/set.rs index 85b93e0eda45b..c7b3514614d62 100644 --- a/src/liballoc/collections/btree/set.rs +++ b/src/liballoc/collections/btree/set.rs @@ -515,6 +515,31 @@ impl BTreeSet { 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 { + IntoIter { iter: self.map.drain() } + } + /// Clears the set, removing all values. /// /// # Examples @@ -884,6 +909,27 @@ impl BTreeSet { 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 = (0..8).collect(); + /// set.retain(|&v| v % 2 == 0); + /// assert_eq!(set.len(), 4); + /// ``` + #[unstable(feature = "btree_drain_retain", issue = "42849")] + pub fn retain(&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 diff --git a/src/liballoc/tests/btree/map.rs b/src/liballoc/tests/btree/map.rs index 27843aeaeb0c8..2d09020218071 100644 --- a/src/liballoc/tests/btree/map.rs +++ b/src/liballoc/tests/btree/map.rs @@ -363,6 +363,282 @@ fn test_range_mut() { } } +const DEPTH_1_CAPACITY: usize = 11; // node::CAPACITY +const DEPTH_2_MIN_SIZE: usize = 11 + 1; // just bigger than node::CAPACITY +#[cfg(not(miri))] // Miri is too slow +const DEPTH_3_MIN_SIZE: usize = 11 + (12 * 11) + 1; + +#[test] +fn test_drain_empty() { + let mut map: BTreeMap = BTreeMap::new(); + { + let mut iter = map.drain(); + assert_eq!(iter.size_hint(), (0, Some(0))); + assert_eq!(iter.next(), None); + assert_eq!(iter.size_hint(), (0, Some(0))); + } + assert!(map.is_empty()); +} + +#[test] +fn test_drain_underfull_consuming_none() { + let pairs = (0..3).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.collect(); + { + let iter = map.drain(); + assert_eq!(iter.size_hint(), (3, Some(3))); + } + assert!(map.is_empty()); +} + +#[test] +fn test_drain_underfull_consuming_some() { + let pairs = (0..3).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.collect(); + { + let mut iter = map.drain(); + assert_eq!(iter.size_hint(), (3, Some(3))); + assert_eq!(iter.next(), Some((0, 0))); + assert_eq!(iter.size_hint(), (2, Some(2))); + assert_eq!(iter.next(), Some((1, 1))); + assert_eq!(iter.size_hint(), (1, Some(1))); + } + assert!(map.is_empty()); +} + +#[test] +fn test_drain_underfull_consuming_all() { + let pairs = (0..3).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.clone().collect(); + assert_eq!(map.drain().collect::>(), pairs.collect::>()); + assert!(map.is_empty()); +} + +#[test] +fn test_drain_depth_2_consuming_none() { + let pairs = (0..DEPTH_2_MIN_SIZE).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.collect(); + map.drain(); + assert!(map.is_empty()); +} + +#[test] +fn test_drain_depth_2_consuming_all() { + let pairs = (0..DEPTH_2_MIN_SIZE).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.clone().collect(); + assert_eq!(map.drain().collect::>(), pairs.collect::>()); + assert!(map.is_empty()); +} + +#[cfg(not(miri))] // Miri is too slow +#[test] +fn test_drain_depth_3_consuming_none() { + let pairs = (0..DEPTH_3_MIN_SIZE).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.collect(); + map.drain(); + assert!(map.is_empty()); +} + +#[cfg(not(miri))] // Miri is too slow +#[test] +fn test_drain_depth_3_consuming_all() { + let pairs = (0..DEPTH_3_MIN_SIZE).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.clone().collect(); + assert_eq!(map.drain().collect::>(), pairs.collect::>()); + assert!(map.is_empty()); +} + +#[test] +fn test_retain_empty() { + let mut map: BTreeMap = BTreeMap::new(); + map.retain(|_, _| unreachable!("there's nothing to retain")); + assert!(map.is_empty()); +} + +#[test] +fn test_retain_underfull_removing_none() { + let pairs = (0..3).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.collect(); + map.retain(|_, _| true); + assert_eq!(map.keys().copied().collect::>(), vec![0, 1, 2]); +} + +#[test] +fn test_retain_underfull_removing_first() { + let pairs = (0..3).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.collect(); + map.retain(|i, _| *i != 0); + assert_eq!(map.keys().copied().collect::>(), vec![1, 2]); +} + +#[test] +fn test_retain_underfull_removing_middle() { + let pairs = (0..3).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.collect(); + map.retain(|i, _| *i != 1); + assert_eq!(map.keys().copied().collect::>(), vec![0, 2]); +} + +#[test] +fn test_retain_underfull_removing_last() { + let pairs = (0..3).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.collect(); + map.retain(|i, _| *i != 2); + assert_eq!(map.keys().copied().collect::>(), vec![0, 1]); +} + +#[test] +fn test_retain_underfull_retaining_first() { + let pairs = (0..3).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.collect(); + map.retain(|i, _| *i == 0); + assert_eq!(map.keys().copied().collect::>(), vec![0]); +} + +#[test] +fn test_retain_underfull_retaining_middle() { + let pairs = (0..3).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.collect(); + map.retain(|i, _| *i == 1); + assert_eq!(map.keys().copied().collect::>(), vec![1]); +} + +#[test] +fn test_retain_underfull_retaining_last() { + let pairs = (0..3).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.collect(); + map.retain(|i, _| *i == 2); + assert_eq!(map.keys().copied().collect::>(), vec![2]); +} + +#[test] +fn test_retain_underfull_removing_all() { + let pairs = (0..3).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.clone().collect(); + map.retain(|_, _| false); + assert!(map.is_empty()); +} + +#[test] +fn test_retain_depth_1_removing_none() { + let pairs = (0..DEPTH_1_CAPACITY).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.collect(); + map.retain(|_, _| true); + assert_eq!(map.keys().copied().collect::>(), (0..DEPTH_1_CAPACITY).collect::>()); +} + +#[test] +fn test_retain_depth_1_muting_only() { + let pairs = (0..DEPTH_1_CAPACITY).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.collect(); + map.retain(|_, v| { *v += 1; true }); + assert_eq!(map.keys().copied().collect::>(), (0..DEPTH_1_CAPACITY).collect::>()); + assert_eq!(map.values().copied().collect::>(), + (1..=DEPTH_1_CAPACITY).collect::>()); +} + +#[test] +fn test_retain_depth_1_removing_first() { + let pairs = (0..DEPTH_1_CAPACITY).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.collect(); + map.retain(|i, _| *i > 0); + assert_eq!(map.keys().copied().collect::>(), + (1..DEPTH_1_CAPACITY).collect::>()); +} + +#[test] +fn test_retain_depth_1_removing_middle() { + let pairs = (0..DEPTH_1_CAPACITY).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.collect(); + map.retain(|i, _| *i != DEPTH_1_CAPACITY / 2); + assert_eq!(map.keys().copied().collect::>(), + (0..DEPTH_1_CAPACITY).filter(|i| *i != DEPTH_1_CAPACITY / 2).collect::>()); +} + +#[test] +fn test_retain_depth_1_removing_last() { + let pairs = (0..DEPTH_1_CAPACITY).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.collect(); + map.retain(|i, _| *i < DEPTH_1_CAPACITY - 1); + assert_eq!(map.keys().copied().collect::>(), + (0..(DEPTH_1_CAPACITY-1)).collect::>()); +} + +#[test] +fn test_retain_depth_1_retaining_first() { + let pairs = (0..DEPTH_1_CAPACITY).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.collect(); + map.retain(|i, _| *i == 0); + assert_eq!(map.keys().copied().collect::>(), vec![0]); +} + +#[test] +fn test_retain_depth_1_retaining_middle() { + let pairs = (0..DEPTH_1_CAPACITY).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.collect(); + map.retain(|i, _| *i == DEPTH_1_CAPACITY / 2); + assert_eq!(map.keys().copied().collect::>(), vec![DEPTH_1_CAPACITY / 2]); +} + +#[test] +fn test_retain_depth_1_retaining_last() { + let pairs = (0..DEPTH_1_CAPACITY).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.collect(); + map.retain(|i, _| *i == DEPTH_1_CAPACITY - 1); + assert_eq!(map.keys().copied().collect::>(), vec![DEPTH_1_CAPACITY - 1]); +} + +#[test] +fn test_retain_depth_1_removing_all() { + let pairs = (0..DEPTH_1_CAPACITY).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.clone().collect(); + map.retain(|_, _| false); + assert!(map.is_empty()); +} + +#[test] +fn test_retain_depth_2_removing_none() { + let pairs = (0..DEPTH_2_MIN_SIZE).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.collect(); + map.retain(|_, _| true); + assert_eq!(map.keys().copied().collect::>(), (0..DEPTH_2_MIN_SIZE).collect::>()); +} + +#[test] +fn test_retain_depth_2_removing_all() { + let pairs = (0..DEPTH_2_MIN_SIZE).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.clone().collect(); + map.retain(|_, _| false); + assert!(map.is_empty()); +} + +#[cfg(not(miri))] // Miri is too slow +#[test] +fn test_retain_depth_3_removing_none() { + let pairs = (0..DEPTH_3_MIN_SIZE).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.collect(); + map.retain(|_, _| true); + assert_eq!(map.keys().copied().collect::>(), (0..DEPTH_3_MIN_SIZE).collect::>()); +} + +#[cfg(not(miri))] // Miri is too slow +#[test] +fn test_retain_depth_3_removing_all() { + let pairs = (0..DEPTH_3_MIN_SIZE).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.clone().collect(); + map.retain(|_, _| false); + assert!(map.is_empty()); +} + +#[test] +fn test_retain_mutating_only() { + let pairs = (0..3).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.collect(); + map.retain(|_, v| { *v += 1; true }); + assert_eq!(map.keys().copied().collect::>(), vec![0, 1, 2]); + assert_eq!(map.values().copied().collect::>(), vec![1, 2, 3]); +} + #[test] fn test_borrow() { // make sure these compile -- using the Borrow trait diff --git a/src/liballoc/tests/btree/set.rs b/src/liballoc/tests/btree/set.rs index 13cd26280227e..3f3ac8e9df704 100644 --- a/src/liballoc/tests/btree/set.rs +++ b/src/liballoc/tests/btree/set.rs @@ -305,6 +305,43 @@ fn test_is_subset() { assert_eq!(is_subset(&[99, 100], &large), false); } +#[test] +fn test_drain() { + let mut x: BTreeSet<_> = [1].iter().copied().collect(); + let mut y: BTreeSet<_> = [1].iter().copied().collect(); + + let v: Vec<_> = x.drain().collect(); + y.drain(); + assert_eq!(v, [1]); + assert!(x.is_empty()); + assert!(y.is_empty()); + x.drain(); + y.drain(); + assert!(x.is_empty()); + assert!(y.is_empty()); +} + +#[test] +fn test_retain() { + let mut x: BTreeSet<_> = [1].iter().copied().collect(); + let mut y: BTreeSet<_> = [1].iter().copied().collect(); + + x.retain(|_| false); + y.retain(|_| true); + assert!(x.is_empty()); + assert_eq!(y.iter().copied().collect::>(), vec![1]); +} + +#[test] +fn test_clear() { + let mut x: BTreeSet<_> = [1].iter().copied().collect(); + + x.clear(); + assert!(x.is_empty()); + x.clear(); + assert!(x.is_empty()); +} + #[test] fn test_zip() { let mut x = BTreeSet::new(); diff --git a/src/liballoc/tests/lib.rs b/src/liballoc/tests/lib.rs index 605e0ef55d707..34e6cf59995fa 100644 --- a/src/liballoc/tests/lib.rs +++ b/src/liballoc/tests/lib.rs @@ -1,5 +1,6 @@ #![feature(allocator_api)] #![feature(box_syntax)] +#![feature(btree_drain_retain)] #![feature(drain_filter)] #![feature(exact_size_is_empty)] #![feature(map_first_last)]