Auto merge of #75071 - ssomers:btree_cleanup_5, r=Mark-Simulacrum

BTreeMap: enforce the panic rule imposed by `replace`

Also, reveal the unsafe parts in the closures fed to it.

r? @Mark-Simulacrum

1 files changed, 51 insertions, 53 deletions

diff --git a/library/alloc/src/collections/btree/navigate.rs b/library/alloc/src/collections/btree/navigate.rs
index 0dcb5930964..33b1ee003ed 100644
--- a/library/alloc/src/collections/btree/navigate.rs
+++ b/library/alloc/src/collections/btree/navigate.rs
@@ -1,3 +1,5 @@
+use core::intrinsics;
+use core::mem;
 use core::ptr;
 
 use super::node::{marker, ForceResult::*, Handle, NodeRef};
@@ -79,16 +81,24 @@ def_next_kv_uncheched_dealloc! {unsafe fn next_kv_unchecked_dealloc: right_kv}
 def_next_kv_uncheched_dealloc! {unsafe fn next_back_kv_unchecked_dealloc: left_kv}
 
 /// This replaces the value behind the `v` unique reference by calling the
-/// relevant function.
+/// relevant function, and returns a result obtained along the way.
 ///
-/// Safety: The change closure must not panic.
+/// If a panic occurs in the `change` closure, the entire process will be aborted.
 #[inline]
-unsafe fn replace<T, R>(v: &mut T, change: impl FnOnce(T) -> (T, R)) -> R {
+fn replace<T, R>(v: &mut T, change: impl FnOnce(T) -> (T, R)) -> R {
+    struct PanicGuard;
+    impl Drop for PanicGuard {
+        fn drop(&mut self) {
+            intrinsics::abort()
+        }
+    }
+    let guard = PanicGuard;
     let value = unsafe { ptr::read(v) };
     let (new_value, ret) = change(value);
     unsafe {
         ptr::write(v, new_value);
     }
+    mem::forget(guard);
     ret
 }
 
@@ -97,26 +107,22 @@ impl<'a, K, V> Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Ed
     /// key and value in between.
     /// Unsafe because the caller must ensure that the leaf edge is not the last one in the tree.
     pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a V) {
-        unsafe {
-            replace(self, |leaf_edge| {
-                let kv = leaf_edge.next_kv();
-                let kv = unwrap_unchecked(kv.ok());
-                (kv.next_leaf_edge(), kv.into_kv())
-            })
-        }
+        replace(self, |leaf_edge| {
+            let kv = leaf_edge.next_kv();
+            let kv = unsafe { unwrap_unchecked(kv.ok()) };
+            (kv.next_leaf_edge(), kv.into_kv())
+        })
     }
 
     /// Moves the leaf edge handle to the previous leaf edge and returns references to the
     /// key and value in between.
     /// Unsafe because the caller must ensure that the leaf edge is not the first one in the tree.
     pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a V) {
-        unsafe {
-            replace(self, |leaf_edge| {
-                let kv = leaf_edge.next_back_kv();
-                let kv = unwrap_unchecked(kv.ok());
-                (kv.next_back_leaf_edge(), kv.into_kv())
-            })
-        }
+        replace(self, |leaf_edge| {
+            let kv = leaf_edge.next_back_kv();
+            let kv = unsafe { unwrap_unchecked(kv.ok()) };
+            (kv.next_back_leaf_edge(), kv.into_kv())
+        })
     }
 }
 
@@ -127,16 +133,14 @@ impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge
     /// - The caller must ensure that the leaf edge is not the last one in the tree.
     /// - Using the updated handle may well invalidate the returned references.
     pub unsafe fn next_unchecked(&mut self) -> (&'a mut K, &'a mut V) {
-        unsafe {
-            let kv = replace(self, |leaf_edge| {
-                let kv = leaf_edge.next_kv();
-                let kv = unwrap_unchecked(kv.ok());
-                (ptr::read(&kv).next_leaf_edge(), kv)
-            });
-            // Doing the descend (and perhaps another move) invalidates the references
-            // returned by `into_kv_mut`, so we have to do this last.
-            kv.into_kv_mut()
-        }
+        let kv = replace(self, |leaf_edge| {
+            let kv = leaf_edge.next_kv();
+            let kv = unsafe { unwrap_unchecked(kv.ok()) };
+            (unsafe { ptr::read(&kv) }.next_leaf_edge(), kv)
+        });
+        // Doing the descend (and perhaps another move) invalidates the references
+        // returned by `into_kv_mut`, so we have to do this last.
+        kv.into_kv_mut()
     }
 
     /// Moves the leaf edge handle to the previous leaf and returns references to the
@@ -145,16 +149,14 @@ impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge
     /// - The caller must ensure that the leaf edge is not the first one in the tree.
     /// - Using the updated handle may well invalidate the returned references.
     pub unsafe fn next_back_unchecked(&mut self) -> (&'a mut K, &'a mut V) {
-        unsafe {
-            let kv = replace(self, |leaf_edge| {
-                let kv = leaf_edge.next_back_kv();
-                let kv = unwrap_unchecked(kv.ok());
-                (ptr::read(&kv).next_back_leaf_edge(), kv)
-            });
-            // Doing the descend (and perhaps another move) invalidates the references
-            // returned by `into_kv_mut`, so we have to do this last.
-            kv.into_kv_mut()
-        }
+        let kv = replace(self, |leaf_edge| {
+            let kv = leaf_edge.next_back_kv();
+            let kv = unsafe { unwrap_unchecked(kv.ok()) };
+            (unsafe { ptr::read(&kv) }.next_back_leaf_edge(), kv)
+        });
+        // Doing the descend (and perhaps another move) invalidates the references
+        // returned by `into_kv_mut`, so we have to do this last.
+        kv.into_kv_mut()
     }
 }
 
@@ -172,14 +174,12 @@ impl<K, V> Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge> {
     ///   call this method again subject to both preconditions listed in the first point,
     ///   or call counterpart `next_back_unchecked` subject to its preconditions.
     pub unsafe fn next_unchecked(&mut self) -> (K, V) {
-        unsafe {
-            replace(self, |leaf_edge| {
-                let kv = next_kv_unchecked_dealloc(leaf_edge);
-                let k = ptr::read(kv.reborrow().into_kv().0);
-                let v = ptr::read(kv.reborrow().into_kv().1);
-                (kv.next_leaf_edge(), (k, v))
-            })
-        }
+        replace(self, |leaf_edge| {
+            let kv = unsafe { next_kv_unchecked_dealloc(leaf_edge) };
+            let k = unsafe { ptr::read(kv.reborrow().into_kv().0) };
+            let v = unsafe { ptr::read(kv.reborrow().into_kv().1) };
+            (kv.next_leaf_edge(), (k, v))
+        })
     }
 
     /// Moves the leaf edge handle to the previous leaf edge and returns the key
@@ -195,14 +195,12 @@ impl<K, V> Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge> {
     ///   call this method again subject to both preconditions listed in the first point,
     ///   or call counterpart `next_unchecked` subject to its preconditions.
     pub unsafe fn next_back_unchecked(&mut self) -> (K, V) {
-        unsafe {
-            replace(self, |leaf_edge| {
-                let kv = next_back_kv_unchecked_dealloc(leaf_edge);
-                let k = ptr::read(kv.reborrow().into_kv().0);
-                let v = ptr::read(kv.reborrow().into_kv().1);
-                (kv.next_back_leaf_edge(), (k, v))
-            })
-        }
+        replace(self, |leaf_edge| {
+            let kv = unsafe { next_back_kv_unchecked_dealloc(leaf_edge) };
+            let k = unsafe { ptr::read(kv.reborrow().into_kv().0) };
+            let v = unsafe { ptr::read(kv.reborrow().into_kv().1) };
+            (kv.next_back_leaf_edge(), (k, v))
+        })
     }
 }