Skip to content

Commit fb1fab5

Browse files
mbrubeckmbrubeck
committed
Tests for HashMap/HashSet::drain_filter
1 parent 49aef96 commit fb1fab5

File tree

2 files changed

+232
-0
lines changed

2 files changed

+232
-0
lines changed

‎library/std/src/collections/hash/map/tests.rs‎

Lines changed: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -924,3 +924,164 @@ fn test_raw_entry() {
924924
}
925925
}
926926
}
927+
928+
mod test_drain_filter {
929+
use super::*;
930+
931+
use crate::panic::{catch_unwind, AssertUnwindSafe};
932+
use crate::sync::atomic::{AtomicUsize, Ordering};
933+
934+
trait EqSorted: Iterator {
935+
fn eq_sorted<I: IntoIterator<Item = Self::Item>>(self, other: I) -> bool;
936+
}
937+
938+
impl<T: Iterator> EqSorted for T
939+
where
940+
T::Item: Eq + Ord,
941+
{
942+
fn eq_sorted<I: IntoIterator<Item = Self::Item>>(self, other: I) -> bool {
943+
let mut v: Vec<_> = self.collect();
944+
v.sort_unstable();
945+
v.into_iter().eq(other)
946+
}
947+
}
948+
949+
#[test]
950+
fn empty() {
951+
let mut map: HashMap<i32, i32> = HashMap::new();
952+
map.drain_filter(|_, _| unreachable!("there's nothing to decide on"));
953+
assert!(map.is_empty());
954+
}
955+
956+
#[test]
957+
fn consuming_nothing() {
958+
let pairs = (0..3).map(|i| (i, i));
959+
let mut map: HashMap<_, _> = pairs.collect();
960+
assert!(map.drain_filter(|_, _| false).eq_sorted(crate::iter::empty()));
961+
assert_eq!(map.len(), 3);
962+
}
963+
964+
#[test]
965+
fn consuming_all() {
966+
let pairs = (0..3).map(|i| (i, i));
967+
let mut map: HashMap<_, _> = pairs.clone().collect();
968+
assert!(map.drain_filter(|_, _| true).eq_sorted(pairs));
969+
assert!(map.is_empty());
970+
}
971+
972+
#[test]
973+
fn mutating_and_keeping() {
974+
let pairs = (0..3).map(|i| (i, i));
975+
let mut map: HashMap<_, _> = pairs.collect();
976+
assert!(
977+
map.drain_filter(|_, v| {
978+
*v += 6;
979+
false
980+
})
981+
.eq_sorted(crate::iter::empty())
982+
);
983+
assert!(map.keys().copied().eq_sorted(0..3));
984+
assert!(map.values().copied().eq_sorted(6..9));
985+
}
986+
987+
#[test]
988+
fn mutating_and_removing() {
989+
let pairs = (0..3).map(|i| (i, i));
990+
let mut map: HashMap<_, _> = pairs.collect();
991+
assert!(
992+
map.drain_filter(|_, v| {
993+
*v += 6;
994+
true
995+
})
996+
.eq_sorted((0..3).map(|i| (i, i + 6)))
997+
);
998+
assert!(map.is_empty());
999+
}
1000+
1001+
#[test]
1002+
fn drop_panic_leak() {
1003+
static PREDS: AtomicUsize = AtomicUsize::new(0);
1004+
static DROPS: AtomicUsize = AtomicUsize::new(0);
1005+
1006+
struct D;
1007+
impl Drop for D {
1008+
fn drop(&mut self) {
1009+
if DROPS.fetch_add(1, Ordering::SeqCst) == 1 {
1010+
panic!("panic in `drop`");
1011+
}
1012+
}
1013+
}
1014+
1015+
let mut map = (0..3).map(|i| (i, D)).collect::<HashMap<_, _>>();
1016+
1017+
catch_unwind(move || {
1018+
drop(map.drain_filter(|_, _| {
1019+
PREDS.fetch_add(1, Ordering::SeqCst);
1020+
true
1021+
}))
1022+
})
1023+
.unwrap_err();
1024+
1025+
assert_eq!(PREDS.load(Ordering::SeqCst), 3);
1026+
assert_eq!(DROPS.load(Ordering::SeqCst), 3);
1027+
}
1028+
1029+
#[test]
1030+
fn pred_panic_leak() {
1031+
static PREDS: AtomicUsize = AtomicUsize::new(0);
1032+
static DROPS: AtomicUsize = AtomicUsize::new(0);
1033+
1034+
struct D;
1035+
impl Drop for D {
1036+
fn drop(&mut self) {
1037+
DROPS.fetch_add(1, Ordering::SeqCst);
1038+
}
1039+
}
1040+
1041+
let mut map = (0..3).map(|i| (i, D)).collect::<HashMap<_, _>>();
1042+
1043+
catch_unwind(AssertUnwindSafe(|| {
1044+
drop(map.drain_filter(|_, _| match PREDS.fetch_add(1, Ordering::SeqCst) {
1045+
0 => true,
1046+
_ => panic!(),
1047+
}))
1048+
}))
1049+
.unwrap_err();
1050+
1051+
assert_eq!(PREDS.load(Ordering::SeqCst), 2);
1052+
assert_eq!(DROPS.load(Ordering::SeqCst), 1);
1053+
assert_eq!(map.len(), 2);
1054+
}
1055+
1056+
// Same as above, but attempt to use the iterator again after the panic in the predicate
1057+
#[test]
1058+
fn pred_panic_reuse() {
1059+
static PREDS: AtomicUsize = AtomicUsize::new(0);
1060+
static DROPS: AtomicUsize = AtomicUsize::new(0);
1061+
1062+
struct D;
1063+
impl Drop for D {
1064+
fn drop(&mut self) {
1065+
DROPS.fetch_add(1, Ordering::SeqCst);
1066+
}
1067+
}
1068+
1069+
let mut map = (0..3).map(|i| (i, D)).collect::<HashMap<_, _>>();
1070+
1071+
{
1072+
let mut it = map.drain_filter(|_, _| match PREDS.fetch_add(1, Ordering::SeqCst) {
1073+
0 => true,
1074+
_ => panic!(),
1075+
});
1076+
catch_unwind(AssertUnwindSafe(|| while it.next().is_some() {})).unwrap_err();
1077+
// Iterator behaviour after a panic is explicitly unspecified,
1078+
// so this is just the current implementation:
1079+
let result = catch_unwind(AssertUnwindSafe(|| it.next()));
1080+
assert!(result.is_err());
1081+
}
1082+
1083+
assert_eq!(PREDS.load(Ordering::SeqCst), 3);
1084+
assert_eq!(DROPS.load(Ordering::SeqCst), 1);
1085+
assert_eq!(map.len(), 2);
1086+
}
1087+
}

‎library/std/src/collections/hash/set/tests.rs‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
use super::super::map::RandomState;
22
use super::HashSet;
33

4+
use crate::panic::{catch_unwind, AssertUnwindSafe};
5+
use crate::sync::atomic::{AtomicU32, Ordering};
6+
47
#[test]
58
fn test_zero_capacities() {
69
type HS = HashSet<i32>;
@@ -413,3 +416,71 @@ fn test_retain() {
413416
assert!(set.contains(&4));
414417
assert!(set.contains(&6));
415418
}
419+
420+
#[test]
421+
fn test_drain_filter() {
422+
let mut x: HashSet<_> = [1].iter().copied().collect();
423+
let mut y: HashSet<_> = [1].iter().copied().collect();
424+
425+
x.drain_filter(|_| true);
426+
y.drain_filter(|_| false);
427+
assert_eq!(x.len(), 0);
428+
assert_eq!(y.len(), 1);
429+
}
430+
431+
#[test]
432+
fn test_drain_filter_drop_panic_leak() {
433+
static PREDS: AtomicU32 = AtomicU32::new(0);
434+
static DROPS: AtomicU32 = AtomicU32::new(0);
435+
436+
#[derive(PartialEq, Eq, PartialOrd, Hash)]
437+
struct D(i32);
438+
impl Drop for D {
439+
fn drop(&mut self) {
440+
if DROPS.fetch_add(1, Ordering::SeqCst) == 1 {
441+
panic!("panic in `drop`");
442+
}
443+
}
444+
}
445+
446+
let mut set = (0..3).map(|i| D(i)).collect::<HashSet<_>>();
447+
448+
catch_unwind(move || {
449+
drop(set.drain_filter(|_| {
450+
PREDS.fetch_add(1, Ordering::SeqCst);
451+
true
452+
}))
453+
})
454+
.ok();
455+
456+
assert_eq!(PREDS.load(Ordering::SeqCst), 3);
457+
assert_eq!(DROPS.load(Ordering::SeqCst), 3);
458+
}
459+
460+
#[test]
461+
fn test_drain_filter_pred_panic_leak() {
462+
static PREDS: AtomicU32 = AtomicU32::new(0);
463+
static DROPS: AtomicU32 = AtomicU32::new(0);
464+
465+
#[derive(PartialEq, Eq, PartialOrd, Hash)]
466+
struct D;
467+
impl Drop for D {
468+
fn drop(&mut self) {
469+
DROPS.fetch_add(1, Ordering::SeqCst);
470+
}
471+
}
472+
473+
let mut set: HashSet<_> = (0..3).map(|_| D).collect();
474+
475+
catch_unwind(AssertUnwindSafe(|| {
476+
drop(set.drain_filter(|_| match PREDS.fetch_add(1, Ordering::SeqCst) {
477+
0 => true,
478+
_ => panic!(),
479+
}))
480+
}))
481+
.ok();
482+
483+
assert_eq!(PREDS.load(Ordering::SeqCst), 1);
484+
assert_eq!(DROPS.load(Ordering::SeqCst), 3);
485+
assert_eq!(set.len(), 0);
486+
}

0 commit comments

Comments
 (0)