From 734fc0477c277c4aed121704a043cd943fdfe1d0 Mon Sep 17 00:00:00 2001 From: Stein Somers Date: Wed, 22 Jul 2020 22:37:08 +0200 Subject: [PATCH] BTreeMap: enforce the panic rule imposed by `replace` --- .../alloc/src/collections/btree/navigate.rs | 104 +++++++++--------- 1 file 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(v: &mut T, change: impl FnOnce(T) -> (T, R)) -> R { +fn replace(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, 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, 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, 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 Handle, 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 Handle, 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)) + }) } }