Rollup merge of rust-lang#58431 - RalfJung:btree, r=Mark-Simulacrum

fix overlapping references in BTree

This fixes two kinds of overlapping references in BTree (both found by running the BTree test suite in Miri).

In `into_slices_mut`, we did `k.into_key_slice_mut()` followed by `self.into_val_slice_mut()` (where `k` is a copy of `self`). Calling `into_val_slice_mut` calls `self.len()`, which creates a shared reference to `NodeHeader`, which unfortunately (due to padding) overlaps with the mutable reference returned by `into_key_slice_mut`.  Hence the key slice got (partially) invalidated.  The fix is to avoid creating an `&NodeHeader` after the first slice got created.

In the iterators, we used to first create the references that will be returned, and then perform the walk on the tree.  Walking the tree creates references (such as `&mut InternalNode`) that overlap with all of the keys and values stored in a pointer; in particular, they overlap with the references the iterator will later return. This is fixed by reordering the operations of walking the tree and obtaining the inner references.

The test suite still passes (and it passes in Miri now!), but there is a lot of code here that I do not understand...

Author: Centril

src/liballoc/collections/btree/map.rs

@@ -1634,9 +1634,11 @@ impl<'a, K, V> RangeMut<'a, K, V> {
 
         let mut cur_handle = match handle.right_kv() {
             Ok(kv) => {
-                let (k, v) = ptr::read(&kv).into_kv_mut();
-                self.front = kv.right_edge();
-                return (k, v);
+                self.front = ptr::read(&kv).right_edge();
+                // Doing the descend invalidates the references returned by `into_kv_mut`,
+                // so we have to do this last.
+                let (k, v) = kv.into_kv_mut();
+                return (k, v); // coerce k from `&mut K` to `&K`
             }
             Err(last_edge) => {
                 let next_level = last_edge.into_node().ascend().ok();
@@ -1647,9 +1649,11 @@ impl<'a, K, V> RangeMut<'a, K, V> {
         loop {
             match cur_handle.right_kv() {
                 Ok(kv) => {
-                    let (k, v) = ptr::read(&kv).into_kv_mut();
-                    self.front = first_leaf_edge(kv.right_edge().descend());
-                    return (k, v);
+                    self.front = first_leaf_edge(ptr::read(&kv).right_edge().descend());
+                    // Doing the descend invalidates the references returned by `into_kv_mut`,
+                    // so we have to do this last.
+                    let (k, v) = kv.into_kv_mut();
+                    return (k, v); // coerce k from `&mut K` to `&K`
                 }
                 Err(last_edge) => {
                     let next_level = last_edge.into_node().ascend().ok();
@@ -1680,9 +1684,11 @@ impl<'a, K, V> RangeMut<'a, K, V> {
 
         let mut cur_handle = match handle.left_kv() {
             Ok(kv) => {
-                let (k, v) = ptr::read(&kv).into_kv_mut();
-                self.back = kv.left_edge();
-                return (k, v);
+                self.back = ptr::read(&kv).left_edge();
+                // Doing the descend invalidates the references returned by `into_kv_mut`,
+                // so we have to do this last.
+                let (k, v) = kv.into_kv_mut();
+                return (k, v); // coerce k from `&mut K` to `&K`
             }
             Err(last_edge) => {
                 let next_level = last_edge.into_node().ascend().ok();
@@ -1693,9 +1699,11 @@ impl<'a, K, V> RangeMut<'a, K, V> {
         loop {
             match cur_handle.left_kv() {
                 Ok(kv) => {
-                    let (k, v) = ptr::read(&kv).into_kv_mut();
-                    self.back = last_leaf_edge(kv.left_edge().descend());
-                    return (k, v);
+                    self.back = last_leaf_edge(ptr::read(&kv).left_edge().descend());
+                    // Doing the descend invalidates the references returned by `into_kv_mut`,
+                    // so we have to do this last.
+                    let (k, v) = kv.into_kv_mut();
+                    return (k, v); // coerce k from `&mut K` to `&K`
                 }
                 Err(last_edge) => {
                     let next_level = last_edge.into_node().ascend().ok();

src/liballoc/collections/btree/node.rs

@@ -645,6 +645,8 @@ impl<'a, K: 'a, V: 'a, Type> NodeRef<marker::Mut<'a>, K, V, Type> {
     }
 
     fn into_key_slice_mut(mut self) -> &'a mut [K] {
+        // Same as for `into_key_slice` above, we try to avoid a run-time check
+        // (the alignment comparison will usually be performed at compile-time).
         if mem::align_of::<K>() > mem::align_of::<LeafNode<(), ()>>() && self.is_shared_root() {
             &mut []
         } else {
@@ -667,9 +669,26 @@ impl<'a, K: 'a, V: 'a, Type> NodeRef<marker::Mut<'a>, K, V, Type> {
         }
     }
 
-    fn into_slices_mut(self) -> (&'a mut [K], &'a mut [V]) {
-        let k = unsafe { ptr::read(&self) };
-        (k.into_key_slice_mut(), self.into_val_slice_mut())
+    fn into_slices_mut(mut self) -> (&'a mut [K], &'a mut [V]) {
+        debug_assert!(!self.is_shared_root());
+        // 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,
+        // because that creates a shared reference to the header, which *can*
+        // overlap with the keys (and even the values, for ZST keys).
+        unsafe {
+            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
+            );
+            (keys, vals)
+        }
     }
 }