diff --git a/src/liballoc/collections/btree/map.rs b/src/liballoc/collections/btree/map.rs index 5ec5064b73515..18fdc8ca2e418 100644 --- a/src/liballoc/collections/btree/map.rs +++ b/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(); diff --git a/src/liballoc/collections/btree/node.rs b/src/liballoc/collections/btree/node.rs index eb0667228d1ff..1cdf4fdd2db20 100644 --- a/src/liballoc/collections/btree/node.rs +++ b/src/liballoc/collections/btree/node.rs @@ -645,6 +645,8 @@ impl<'a, K: 'a, V: 'a, Type> NodeRef, 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::() > mem::align_of::>() && self.is_shared_root() { &mut [] } else { @@ -667,9 +669,26 @@ impl<'a, K: 'a, V: 'a, Type> NodeRef, 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) + } } }