2 files changed, 47 insertions, 20 deletions

diff --git a/src/liballoc/collections/btree/map.rs b/src/liballoc/collections/btree/map.rs
index d2f4278d0d0..24d1f61fa68 100644
--- a/src/liballoc/collections/btree/map.rs
+++ b/src/liballoc/collections/btree/map.rs
@@ -1672,19 +1672,12 @@ impl<'a, K: 'a, V: 'a> DrainFilterInner<'a, K, V> {
         edge.reborrow().next_kv().ok().map(|kv| kv.into_kv())
     }
 
-    unsafe fn next_kv(
-        &mut self,
-    ) -> Option<Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::KV>> {
-        let edge = self.cur_leaf_edge.as_ref()?;
-        unsafe { ptr::read(edge).next_kv().ok() }
-    }
-
     /// Implementation of a typical `DrainFilter::next` method, given the predicate.
     pub(super) fn next<F>(&mut self, pred: &mut F) -> Option<(K, V)>
     where
         F: FnMut(&K, &mut V) -> bool,
     {
-        while let Some(mut kv) = unsafe { self.next_kv() } {
+        while let Ok(mut kv) = self.cur_leaf_edge.take()?.next_kv() {
             let (k, v) = kv.kv_mut();
             if pred(k, v) {
                 *self.length -= 1;
diff --git a/src/liballoc/tests/btree/map.rs b/src/liballoc/tests/btree/map.rs
index 5cd3ce0a2d8..f9f81716e35 100644
--- a/src/liballoc/tests/btree/map.rs
+++ b/src/liballoc/tests/btree/map.rs
@@ -887,10 +887,8 @@ mod test_drain_filter {
             }
         }
 
-        let mut map = BTreeMap::new();
-        map.insert(0, D);
-        map.insert(4, D);
-        map.insert(8, D);
+        // Keys are multiples of 4, so that each key is counted by a hexadecimal digit.
+        let mut map = (0..3).map(|i| (i * 4, D)).collect::<BTreeMap<_, _>>();
 
         catch_unwind(move || {
             drop(map.drain_filter(|i, _| {
@@ -898,7 +896,7 @@ mod test_drain_filter {
                 true
             }))
         })
-        .ok();
+        .unwrap_err();
 
         assert_eq!(PREDS.load(Ordering::SeqCst), 0x011);
         assert_eq!(DROPS.load(Ordering::SeqCst), 3);
@@ -916,10 +914,8 @@ mod test_drain_filter {
             }
         }
 
-        let mut map = BTreeMap::new();
-        map.insert(0, D);
-        map.insert(4, D);
-        map.insert(8, D);
+        // Keys are multiples of 4, so that each key is counted by a hexadecimal digit.
+        let mut map = (0..3).map(|i| (i * 4, D)).collect::<BTreeMap<_, _>>();
 
         catch_unwind(AssertUnwindSafe(|| {
             drop(map.drain_filter(|i, _| {
@@ -930,7 +926,45 @@ mod test_drain_filter {
                 }
             }))
         }))
-        .ok();
+        .unwrap_err();
+
+        assert_eq!(PREDS.load(Ordering::SeqCst), 0x011);
+        assert_eq!(DROPS.load(Ordering::SeqCst), 1);
+        assert_eq!(map.len(), 2);
+        assert_eq!(map.first_entry().unwrap().key(), &4);
+        assert_eq!(map.last_entry().unwrap().key(), &8);
+    }
+
+    // Same as above, but attempt to use the iterator again after the panic in the predicate
+    #[test]
+    fn pred_panic_reuse() {
+        static PREDS: AtomicUsize = AtomicUsize::new(0);
+        static DROPS: AtomicUsize = AtomicUsize::new(0);
+
+        struct D;
+        impl Drop for D {
+            fn drop(&mut self) {
+                DROPS.fetch_add(1, Ordering::SeqCst);
+            }
+        }
+
+        // Keys are multiples of 4, so that each key is counted by a hexadecimal digit.
+        let mut map = (0..3).map(|i| (i * 4, D)).collect::<BTreeMap<_, _>>();
+
+        {
+            let mut it = map.drain_filter(|i, _| {
+                PREDS.fetch_add(1usize << i, Ordering::SeqCst);
+                match i {
+                    0 => true,
+                    _ => panic!(),
+                }
+            });
+            catch_unwind(AssertUnwindSafe(|| while it.next().is_some() {})).unwrap_err();
+            // Iterator behaviour after a panic is explicitly unspecified,
+            // so this is just the current implementation:
+            let result = catch_unwind(AssertUnwindSafe(|| it.next()));
+            assert!(matches!(result, Ok(None)));
+        }
 
         assert_eq!(PREDS.load(Ordering::SeqCst), 0x011);
         assert_eq!(DROPS.load(Ordering::SeqCst), 1);
@@ -1399,7 +1433,7 @@ fn test_into_iter_drop_leak_height_0() {
     map.insert("d", D);
     map.insert("e", D);
 
-    catch_unwind(move || drop(map.into_iter())).ok();
+    catch_unwind(move || drop(map.into_iter())).unwrap_err();
 
     assert_eq!(DROPS.load(Ordering::SeqCst), 5);
 }
@@ -1423,7 +1457,7 @@ fn test_into_iter_drop_leak_height_1() {
         DROPS.store(0, Ordering::SeqCst);
         PANIC_POINT.store(panic_point, Ordering::SeqCst);
         let map: BTreeMap<_, _> = (0..size).map(|i| (i, D)).collect();
-        catch_unwind(move || drop(map.into_iter())).ok();
+        catch_unwind(move || drop(map.into_iter())).unwrap_err();
         assert_eq!(DROPS.load(Ordering::SeqCst), size);
     }
 }