Make RangeMut::next_unchecked() output a mutable key reference

euclio · Dec 23, 2019 · 293cdf7 · 293cdf7
1 parent a916ac2
commit 293cdf7
Showing 1 changed file with 15 additions and 7 deletions.
diff --git a/src/liballoc/collections/btree/map.rs b/src/liballoc/collections/btree/map.rs
@@ -1328,7 +1328,10 @@ impl<'a, K: 'a, V: 'a> Iterator for IterMut<'a, K, V> {
             None
         } else {
             self.length -= 1;
-            unsafe { Some(self.range.next_unchecked()) }
+            unsafe {
+                let (k, v) = self.range.next_unchecked();
+                Some((k, v)) // coerce k from `&mut K` to `&K`
+            }
         }
     }
 
@@ -1707,7 +1710,14 @@ impl<'a, K, V> Iterator for RangeMut<'a, K, V> {
     type Item = (&'a K, &'a mut V);
 
     fn next(&mut self) -> Option<(&'a K, &'a mut V)> {
-        if self.front == self.back { None } else { unsafe { Some(self.next_unchecked()) } }
+        if self.front == self.back {
+            None
+        } else {
+            unsafe {
+                let (k, v) = self.next_unchecked();
+                Some((k, v)) // coerce k from `&mut K` to `&K`
+            }
+        }
     }
 
     fn last(mut self) -> Option<(&'a K, &'a mut V)> {
@@ -1716,16 +1726,15 @@ impl<'a, K, V> Iterator for RangeMut<'a, K, V> {
 }
 
 impl<'a, K, V> RangeMut<'a, K, V> {
-    unsafe fn next_unchecked(&mut self) -> (&'a K, &'a mut V) {
+    unsafe fn next_unchecked(&mut self) -> (&'a mut K, &'a mut V) {
         let handle = ptr::read(&self.front);
 
         let mut cur_handle = match handle.right_kv() {
             Ok(kv) => {
                 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`
+                return kv.into_kv_mut();
             }
             Err(last_edge) => {
                 let next_level = last_edge.into_node().ascend().ok();
@@ -1739,8 +1748,7 @@ impl<'a, K, V> RangeMut<'a, 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`
+                    return kv.into_kv_mut();
                 }
                 Err(last_edge) => {
                     let next_level = last_edge.into_node().ascend().ok();