fix invalidating references in BTree iterators

2 files changed, 21 insertions, 13 deletions

diff --git a/src/liballoc/collections/btree/map.rs b/src/liballoc/collections/btree/map.rs
index 5ec5064b735..18fdc8ca2e4 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 cb38e6470ab..1cdf4fdd2db 100644
--- a/src/liballoc/collections/btree/node.rs
+++ b/src/liballoc/collections/btree/node.rs
@@ -675,7 +675,7 @@ impl<'a, K: 'a, V: 'a, Type> NodeRef<marker::Mut<'a>, K, V, Type> {
         // 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.
+        // overlap with the keys (and even the values, for ZST keys).
         unsafe {
             let len = self.len();
             let leaf = self.as_leaf_mut();