Rollup merge of #74677 - ssomers:btree_cleanup_2, r=Amanieu

Remove needless unsafety from BTreeMap::drain_filter

Remove one piece of unsafe code in the iteration over the iterator returned by BTreeMap::drain_filter.
- Changes an explicitly unspecified part of the API: when the user-supplied predicate (or some of BTreeMap's code) panicked, and the caller tries to use the iterator again, we no longer offer the same key/value pair to the predicate again but pretend the iterator has finished. Note that Miri does not find UB in the test case added here with the unsafe code (or without).
- Makes the code a little easier on the eyes.
- Makes the code a little harder on the CPU:
```
benchcmp c0 c2 --threshold 3
 name                                         c0 ns/iter  c2 ns/iter  diff ns/iter  diff %  speedup
 btree::set::clone_100_and_drain_all          2,794       2,900                106   3.79%   x 0.96
 btree::set::clone_100_and_drain_half         2,604       2,964                360  13.82%   x 0.88
 btree::set::clone_10k_and_drain_half         287,770     322,755           34,985  12.16%   x 0.89
```
r? @Amanieu

diff --git a/src/liballoc/collections/btree/map.rs b/src/liballoc/collections/btree/map.rs
index d2f4278d0d0e0a817d51246c971265c4056c6103..24d1f61fa68c48dd773ec161e2d1f98356b4f4bd 100644 (file)
--- a/src/liballoc/collections/btree/map.rs
+++ b/src/liballoc/collections/btree/map.rs
@@ -1672,19 +1672,12 @@ pub(super) fn peek(&self) -> Option<(&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 5cd3ce0a2d80dd86a45dffc75a2a54c3542f9b48..f9f81716e357c243b3c7ec9add10b11f6dabff5e 100644 (file)
--- a/src/liballoc/tests/btree/map.rs
+++ b/src/liballoc/tests/btree/map.rs
@@ -887,10 +887,8 @@ fn drop(&mut self) {
             }
         }
 
-        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 @@ fn drop(&mut self) {
                 true
             }))
         })
-        .ok();
+        .unwrap_err();
 
         assert_eq!(PREDS.load(Ordering::SeqCst), 0x011);
         assert_eq!(DROPS.load(Ordering::SeqCst), 3);
@@ -916,10 +914,8 @@ fn drop(&mut self) {
             }
         }
 
-        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 @@ fn drop(&mut self) {
                 }
             }))
         }))
-        .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 drop(&mut self) {
     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 drop(&mut self) {
         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);
     }
 }