Make functions dependent only on shared root avoidance safe

diff --git a/src/liballoc/collections/btree/map.rs b/src/liballoc/collections/btree/map.rs
index 478fa8e15268e4818d441e8f58b4215aace2ba17..21f99257d81d0fea68778e113a4285dfc510ee61 100644 (file)
--- a/src/liballoc/collections/btree/map.rs
+++ b/src/liballoc/collections/btree/map.rs
@@ -1349,7 +1349,8 @@ fn fix_left_border(&mut self) {
         self.fix_top();
     }
 
-    /// If the root node is the shared root node, allocate our own node.
+    /// If the root node is the empty (non-allocated) root node, allocate our
+    /// own node.
     fn ensure_root_is_owned(&mut self) {
         if self.root.is_none() {
             self.root = Some(node::Root::new_leaf());
@@ -1509,7 +1510,6 @@ fn drop(&mut self) {
                 // don't have to care about panics this time (they'll abort).
                 while let Some(_) = self.0.next() {}
 
-                // No need to avoid the shared root, because the tree was definitely not empty.
                 unsafe {
                     let mut node =
                         unwrap_unchecked(ptr::read(&self.0.front)).into_node().forget_type();

diff --git a/src/liballoc/collections/btree/node.rs b/src/liballoc/collections/btree/node.rs
index 5530c9593a51a9a9fe64bc8d522b944450c2b750..99f531264ba2076bbe951ff5a97e49f8de24102e 100644 (file)
--- a/src/liballoc/collections/btree/node.rs
+++ b/src/liballoc/collections/btree/node.rs
@@ -169,8 +169,9 @@ fn as_ptr(&self) -> NonNull<LeafNode<K, V>> {
     }
 }
 
-/// Either an owned tree or a shared, empty tree.  Note that this does not have a destructor,
-/// and must be cleaned up manually if it is an owned tree.
+/// An owned tree.
+///
+/// Note that this does not have a destructor, and must be cleaned up manually.
 pub struct Root<K, V> {
     node: BoxedNode<K, V>,
     /// The number of levels below the root node.
@@ -278,10 +279,7 @@ pub fn pop_level(&mut self) {
 ///   `Leaf`, the `NodeRef` points to a leaf node, when this is `Internal` the
 ///   `NodeRef` points to an internal node, and when this is `LeafOrInternal` the
 ///   `NodeRef` could be pointing to either type of node.
-///   Note that in case of a leaf node, this might still be the shared root!
-///   Only turn this into a `LeafNode` reference if you know it is not the shared root!
-///   Shared references must be dereferenceable *for the entire size of their pointee*,
-///   so '&LeafNode` or `&InternalNode` pointing to the shared root is undefined behavior.
+///
 ///   Turning this into a `NodeHeader` reference is always safe.
 pub struct NodeRef<BorrowType, K, V, Type> {
     /// The number of levels below the node.
@@ -344,14 +342,15 @@ fn reborrow(&self) -> NodeRef<marker::Immut<'_>, K, V, Type> {
         NodeRef { height: self.height, node: self.node, root: self.root, _marker: PhantomData }
     }
 
-    /// Exposes the leaf "portion" of any leaf or internal node that is not the shared root.
+    /// Exposes the leaf "portion" of any leaf or internal node.
     /// If the node is a leaf, this function simply opens up its data.
     /// If the node is an internal node, so not a leaf, it does have all the data a leaf has
     /// (header, keys and values), and this function exposes that.
-    /// Unsafe because the node must not be the shared root. For more information,
-    /// see the `NodeRef` comments.
-    unsafe fn as_leaf(&self) -> &LeafNode<K, V> {
-        self.node.as_ref()
+    fn as_leaf(&self) -> &LeafNode<K, V> {
+        // The node must be valid for at least the LeafNode portion.
+        // This is not a reference in the NodeRef type because we don't know if
+        // it should be unique or shared.
+        unsafe { self.node.as_ref() }
     }
 
     fn as_header(&self) -> &NodeHeader<K, V> {
@@ -359,14 +358,12 @@ fn as_header(&self) -> &NodeHeader<K, V> {
     }
 
     /// Borrows a view into the keys stored in the node.
-    /// Unsafe because the caller must ensure that the node is not the shared root.
-    pub unsafe fn keys(&self) -> &[K] {
+    pub fn keys(&self) -> &[K] {
         self.reborrow().into_key_slice()
     }
 
     /// Borrows a view into the values stored in the node.
-    /// Unsafe because the caller must ensure that the node is not the shared root.
-    unsafe fn vals(&self) -> &[V] {
+    fn vals(&self) -> &[V] {
         self.reborrow().into_val_slice()
     }
 
@@ -470,39 +467,37 @@ unsafe fn reborrow_mut(&mut self) -> NodeRef<marker::Mut<'_>, K, V, Type> {
     /// (header, keys and values), and this function exposes that.
     ///
     /// Returns a raw ptr to avoid asserting exclusive access to the entire node.
-    /// This also implies you can invoke this member on the shared root, but the resulting pointer
-    /// might not be properly aligned and definitely would not allow accessing keys and values.
     fn as_leaf_mut(&mut self) -> *mut LeafNode<K, V> {
         self.node.as_ptr()
     }
 
-    /// Unsafe because the caller must ensure that the node is not the shared root.
-    unsafe fn keys_mut(&mut self) -> &mut [K] {
-        self.reborrow_mut().into_key_slice_mut()
+    fn keys_mut(&mut self) -> &mut [K] {
+        // SAFETY: the caller will not be able to call further methods on self
+        // until the key slice reference is dropped, as we have unique access
+        // for the lifetime of the borrow.
+        unsafe { self.reborrow_mut().into_key_slice_mut() }
     }
 
-    /// Unsafe because the caller must ensure that the node is not the shared root.
-    unsafe fn vals_mut(&mut self) -> &mut [V] {
-        self.reborrow_mut().into_val_slice_mut()
+    fn vals_mut(&mut self) -> &mut [V] {
+        // SAFETY: the caller will not be able to call further methods on self
+        // until the value slice reference is dropped, as we have unique access
+        // for the lifetime of the borrow.
+        unsafe { self.reborrow_mut().into_val_slice_mut() }
     }
 }
 
 impl<'a, K: 'a, V: 'a, Type> NodeRef<marker::Immut<'a>, K, V, Type> {
-    /// Unsafe because the caller must ensure that the node is not the shared root.
-    unsafe fn into_key_slice(self) -> &'a [K] {
-        // We cannot be the shared root, so `as_leaf` is okay.
-        slice::from_raw_parts(MaybeUninit::first_ptr(&self.as_leaf().keys), self.len())
+    fn into_key_slice(self) -> &'a [K] {
+        unsafe { slice::from_raw_parts(MaybeUninit::first_ptr(&self.as_leaf().keys), self.len()) }
     }
 
-    /// Unsafe because the caller must ensure that the node is not the shared root.
-    unsafe fn into_val_slice(self) -> &'a [V] {
-        // We cannot be the shared root, so `as_leaf` is okay.
-        slice::from_raw_parts(MaybeUninit::first_ptr(&self.as_leaf().vals), self.len())
+    fn into_val_slice(self) -> &'a [V] {
+        unsafe { slice::from_raw_parts(MaybeUninit::first_ptr(&self.as_leaf().vals), self.len()) }
     }
 
-    /// Unsafe because the caller must ensure that the node is not the shared root.
-    unsafe fn into_slices(self) -> (&'a [K], &'a [V]) {
-        let k = ptr::read(&self);
+    fn into_slices(self) -> (&'a [K], &'a [V]) {
+        // SAFETY: equivalent to reborrow() except not requiring Type: 'a
+        let k = unsafe { ptr::read(&self) };
         (k.into_key_slice(), self.into_val_slice())
     }
 }
@@ -514,25 +509,27 @@ pub fn into_root_mut(self) -> &'a mut Root<K, V> {
         unsafe { &mut *(self.root as *mut Root<K, V>) }
     }
 
-    /// Unsafe because the caller must ensure that the node is not the shared root.
-    unsafe fn into_key_slice_mut(mut self) -> &'a mut [K] {
-        // We cannot be the shared root, so `as_leaf_mut` is okay.
-        slice::from_raw_parts_mut(
-            MaybeUninit::first_ptr_mut(&mut (*self.as_leaf_mut()).keys),
-            self.len(),
-        )
+    fn into_key_slice_mut(mut self) -> &'a mut [K] {
+        // SAFETY: The keys of a node must always be initialized up to length.
+        unsafe {
+            slice::from_raw_parts_mut(
+                MaybeUninit::first_ptr_mut(&mut (*self.as_leaf_mut()).keys),
+                self.len(),
+            )
+        }
     }
 
-    /// Unsafe because the caller must ensure that the node is not the shared root.
-    unsafe fn into_val_slice_mut(mut self) -> &'a mut [V] {
-        slice::from_raw_parts_mut(
-            MaybeUninit::first_ptr_mut(&mut (*self.as_leaf_mut()).vals),
-            self.len(),
-        )
+    fn into_val_slice_mut(mut self) -> &'a mut [V] {
+        // SAFETY: The values of a node must always be initialized up to length.
+        unsafe {
+            slice::from_raw_parts_mut(
+                MaybeUninit::first_ptr_mut(&mut (*self.as_leaf_mut()).vals),
+                self.len(),
+            )
+        }
     }
 
-    /// Unsafe because the caller must ensure that the node is not the shared root.
-    unsafe fn into_slices_mut(mut self) -> (&'a mut [K], &'a mut [V]) {
+    fn into_slices_mut(mut self) -> (&'a mut [K], &'a mut [V]) {
         // We cannot use the getters here, because calling the second one
         // invalidates the reference returned by the first.
         // More precisely, it is the call to `len` that is the culprit,
@@ -540,8 +537,13 @@ unsafe fn into_slices_mut(mut self) -> (&'a mut [K], &'a mut [V]) {
         // overlap with the keys (and even the values, for ZST keys).
         let len = self.len();
         let leaf = self.as_leaf_mut();
-        let keys = slice::from_raw_parts_mut(MaybeUninit::first_ptr_mut(&mut (*leaf).keys), len);
-        let vals = slice::from_raw_parts_mut(MaybeUninit::first_ptr_mut(&mut (*leaf).vals), len);
+        // SAFETY: The keys and values of a node must always be initialized up to length.
+        let keys = unsafe {
+            slice::from_raw_parts_mut(MaybeUninit::first_ptr_mut(&mut (*leaf).keys), len)
+        };
+        let vals = unsafe {
+            slice::from_raw_parts_mut(MaybeUninit::first_ptr_mut(&mut (*leaf).vals), len)
+        };
         (keys, vals)
     }
 }
@@ -698,8 +700,7 @@ pub fn pop_front(&mut self) -> (K, V, Option<Root<K, V>>) {
         }
     }
 
-    /// Unsafe because the caller must ensure that the node is not the shared root.
-    unsafe fn into_kv_pointers_mut(mut self) -> (*mut K, *mut V) {
+    fn into_kv_pointers_mut(mut self) -> (*mut K, *mut V) {
         (self.keys_mut().as_mut_ptr(), self.vals_mut().as_mut_ptr())
     }
 }

diff --git a/src/liballoc/collections/btree/search.rs b/src/liballoc/collections/btree/search.rs
index f9354263864e5e7eecb799c69d69784c1d3d52f3..4e80f7f21ebff5c6817a82f9e9436acf3ea5d9d2 100644 (file)
--- a/src/liballoc/collections/btree/search.rs
+++ b/src/liballoc/collections/btree/search.rs
@@ -67,12 +67,11 @@ fn search_linear<BorrowType, K, V, Type, Q: ?Sized>(
     Q: Ord,
     K: Borrow<Q>,
 {
-    // This function is defined over all borrow types (immutable, mutable, owned),
-    // and may be called on the shared root in each case.
+    // This function is defined over all borrow types (immutable, mutable, owned).
     // Using `keys()` is fine here even if BorrowType is mutable, as all we return
     // is an index -- not a reference.
     let len = node.len();
-    let keys = unsafe { node.keys() }; // safe because a non-empty node cannot be the shared root
+    let keys = node.keys();
     for (i, k) in keys.iter().enumerate() {
         match key.cmp(k.borrow()) {
             Ordering::Greater => {}