Auto merge of #85980 - ssomers:btree_cleanup_LeafRange, r=Mark-Simulacrum

BTree: encapsulate LeafRange better & some debug asserts

Looking at iterators again, I think #81937 didn't house enough code in `LeafRange`. Moving the API boundary a little makes things more local in navigate.rs and less complicated in map.rs.

r? `@Mark-Simulacrum`

2 files changed, 111 insertions, 51 deletions

diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs
index 30194aa446f..d7519c6d093 100644
--- a/library/alloc/src/collections/btree/map.rs
+++ b/library/alloc/src/collections/btree/map.rs
@@ -1299,7 +1299,7 @@ impl<'a, K: 'a, V: 'a> Iterator for Iter<'a, K, V> {
             None
         } else {
             self.length -= 1;
-            Some(unsafe { self.range.next_unchecked() })
+            Some(unsafe { self.range.inner.next_unchecked() })
         }
     }
 
@@ -1330,7 +1330,7 @@ impl<'a, K: 'a, V: 'a> DoubleEndedIterator for Iter<'a, K, V> {
             None
         } else {
             self.length -= 1;
-            Some(unsafe { self.range.next_back_unchecked() })
+            Some(unsafe { self.range.inner.next_back_unchecked() })
         }
     }
 }
@@ -1368,7 +1368,7 @@ impl<'a, K: 'a, V: 'a> Iterator for IterMut<'a, K, V> {
             None
         } else {
             self.length -= 1;
-            Some(unsafe { self.range.next_unchecked() })
+            Some(unsafe { self.range.inner.next_unchecked() })
         }
     }
 
@@ -1396,7 +1396,7 @@ impl<'a, K: 'a, V: 'a> DoubleEndedIterator for IterMut<'a, K, V> {
             None
         } else {
             self.length -= 1;
-            Some(unsafe { self.range.next_back_unchecked() })
+            Some(unsafe { self.range.inner.next_back_unchecked() })
         }
     }
 }
@@ -1475,7 +1475,7 @@ impl<K, V> Drop for Dropper<K, V> {
 #[stable(feature = "btree_drop", since = "1.7.0")]
 impl<K, V> Drop for IntoIter<K, V> {
     fn drop(&mut self) {
-        if let Some(front) = self.range.front.take() {
+        if let Some(front) = self.range.take_front() {
             Dropper { front, remaining_length: self.length };
         }
     }
@@ -1490,8 +1490,7 @@ impl<K, V> Iterator for IntoIter<K, V> {
             None
         } else {
             self.length -= 1;
-            let front = self.range.front.as_mut().unwrap();
-            let kv = unsafe { front.deallocating_next_unchecked() };
+            let kv = unsafe { self.range.deallocating_next_unchecked() };
             Some(kv.into_key_val())
         }
     }
@@ -1508,8 +1507,7 @@ impl<K, V> DoubleEndedIterator for IntoIter<K, V> {
             None
         } else {
             self.length -= 1;
-            let back = self.range.back.as_mut().unwrap();
-            let kv = unsafe { back.deallocating_next_back_unchecked() };
+            let kv = unsafe { self.range.deallocating_next_back_unchecked() };
             Some(kv.into_key_val())
         }
     }
@@ -1727,7 +1725,7 @@ impl<'a, K, V> Iterator for Range<'a, K, V> {
     type Item = (&'a K, &'a V);
 
     fn next(&mut self) -> Option<(&'a K, &'a V)> {
-        if self.inner.is_empty() { None } else { Some(unsafe { self.next_unchecked() }) }
+        self.inner.next_checked()
     }
 
     fn last(mut self) -> Option<(&'a K, &'a V)> {
@@ -1777,12 +1775,6 @@ impl<K, V> ExactSizeIterator for ValuesMut<'_, K, V> {
 #[stable(feature = "fused", since = "1.26.0")]
 impl<K, V> FusedIterator for ValuesMut<'_, K, V> {}
 
-impl<'a, K, V> Range<'a, K, V> {
-    unsafe fn next_unchecked(&mut self) -> (&'a K, &'a V) {
-        unsafe { self.inner.front.as_mut().unwrap_unchecked().next_unchecked() }
-    }
-}
-
 #[stable(feature = "map_into_keys_values", since = "1.54.0")]
 impl<K, V> Iterator for IntoKeys<K, V> {
     type Item = K;
@@ -1862,13 +1854,7 @@ impl<K, V> FusedIterator for IntoValues<K, V> {}
 #[stable(feature = "btree_range", since = "1.17.0")]
 impl<'a, K, V> DoubleEndedIterator for Range<'a, K, V> {
     fn next_back(&mut self) -> Option<(&'a K, &'a V)> {
-        if self.inner.is_empty() { None } else { Some(unsafe { self.next_back_unchecked() }) }
-    }
-}
-
-impl<'a, K, V> Range<'a, K, V> {
-    unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a V) {
-        unsafe { self.inner.back.as_mut().unwrap_unchecked().next_back_unchecked() }
+        self.inner.next_back_checked()
     }
 }
 
@@ -1878,7 +1864,7 @@ impl<K, V> FusedIterator for Range<'_, K, V> {}
 #[stable(feature = "btree_range", since = "1.17.0")]
 impl<K, V> Clone for Range<'_, K, V> {
     fn clone(&self) -> Self {
-        Range { inner: LeafRange { front: self.inner.front, back: self.inner.back } }
+        Range { inner: self.inner.clone() }
     }
 }
 
@@ -1887,7 +1873,7 @@ 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.inner.is_empty() { None } else { Some(unsafe { self.next_unchecked() }) }
+        self.inner.next_checked()
     }
 
     fn last(mut self) -> Option<(&'a K, &'a mut V)> {
@@ -1904,10 +1890,6 @@ 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 { self.inner.front.as_mut().unwrap_unchecked().next_unchecked() }
-    }
-
     /// Returns an iterator of references over the remaining items.
     #[inline]
     pub(super) fn iter(&self) -> Range<'_, K, V> {
@@ -1918,19 +1900,13 @@ impl<'a, K, V> RangeMut<'a, K, V> {
 #[stable(feature = "btree_range", since = "1.17.0")]
 impl<'a, K, V> DoubleEndedIterator for RangeMut<'a, K, V> {
     fn next_back(&mut self) -> Option<(&'a K, &'a mut V)> {
-        if self.inner.is_empty() { None } else { Some(unsafe { self.next_back_unchecked() }) }
+        self.inner.next_back_checked()
     }
 }
 
 #[stable(feature = "fused", since = "1.26.0")]
 impl<K, V> FusedIterator for RangeMut<'_, K, V> {}
 
-impl<'a, K, V> RangeMut<'a, K, V> {
-    unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a mut V) {
-        unsafe { self.inner.back.as_mut().unwrap_unchecked().next_back_unchecked() }
-    }
-}
-
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<K: Ord, V> FromIterator<(K, V)> for BTreeMap<K, V> {
     fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> BTreeMap<K, V> {
diff --git a/library/alloc/src/collections/btree/navigate.rs b/library/alloc/src/collections/btree/navigate.rs
index 563c070dd0f..17568f7b005 100644
--- a/library/alloc/src/collections/btree/navigate.rs
+++ b/library/alloc/src/collections/btree/navigate.rs
@@ -4,9 +4,16 @@ use core::ptr;
 
 use super::node::{marker, ForceResult::*, Handle, NodeRef};
 
+// `front` and `back` are always both `None` or both `Some`.
 pub struct LeafRange<BorrowType, K, V> {
-    pub front: Option<Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>>,
-    pub back: Option<Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>>,
+    front: Option<Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>>,
+    back: Option<Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>>,
+}
+
+impl<'a, K: 'a, V: 'a> Clone for LeafRange<marker::Immut<'a>, K, V> {
+    fn clone(&self) -> Self {
+        LeafRange { front: self.front.clone(), back: self.back.clone() }
+    }
 }
 
 impl<BorrowType, K, V> LeafRange<BorrowType, K, V> {
@@ -14,7 +21,7 @@ impl<BorrowType, K, V> LeafRange<BorrowType, K, V> {
         LeafRange { front: None, back: None }
     }
 
-    pub fn is_empty(&self) -> bool {
+    fn is_empty(&self) -> bool {
         self.front == self.back
     }
 
@@ -27,6 +34,81 @@ impl<BorrowType, K, V> LeafRange<BorrowType, K, V> {
     }
 }
 
+impl<'a, K, V> LeafRange<marker::Immut<'a>, K, V> {
+    #[inline]
+    pub fn next_checked(&mut self) -> Option<(&'a K, &'a V)> {
+        if self.is_empty() { None } else { Some(unsafe { self.next_unchecked() }) }
+    }
+
+    #[inline]
+    pub fn next_back_checked(&mut self) -> Option<(&'a K, &'a V)> {
+        if self.is_empty() { None } else { Some(unsafe { self.next_back_unchecked() }) }
+    }
+
+    #[inline]
+    pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a V) {
+        debug_assert!(self.front.is_some());
+        unsafe { self.front.as_mut().unwrap_unchecked().next_unchecked() }
+    }
+
+    #[inline]
+    pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a V) {
+        debug_assert!(self.back.is_some());
+        unsafe { self.back.as_mut().unwrap_unchecked().next_back_unchecked() }
+    }
+}
+
+impl<'a, K, V> LeafRange<marker::ValMut<'a>, K, V> {
+    #[inline]
+    pub fn next_checked(&mut self) -> Option<(&'a K, &'a mut V)> {
+        if self.is_empty() { None } else { Some(unsafe { self.next_unchecked() }) }
+    }
+
+    #[inline]
+    pub fn next_back_checked(&mut self) -> Option<(&'a K, &'a mut V)> {
+        if self.is_empty() { None } else { Some(unsafe { self.next_back_unchecked() }) }
+    }
+
+    #[inline]
+    pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a mut V) {
+        debug_assert!(self.front.is_some());
+        unsafe { self.front.as_mut().unwrap_unchecked().next_unchecked() }
+    }
+
+    #[inline]
+    pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a mut V) {
+        debug_assert!(self.back.is_some());
+        unsafe { self.back.as_mut().unwrap_unchecked().next_back_unchecked() }
+    }
+}
+
+impl<K, V> LeafRange<marker::Dying, K, V> {
+    #[inline]
+    pub fn take_front(
+        &mut self,
+    ) -> Option<Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge>> {
+        self.front.take()
+    }
+
+    #[inline]
+    pub unsafe fn deallocating_next_unchecked(
+        &mut self,
+    ) -> Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV> {
+        debug_assert!(self.front.is_some());
+        let front = self.front.as_mut().unwrap();
+        unsafe { front.deallocating_next_unchecked() }
+    }
+
+    #[inline]
+    pub unsafe fn deallocating_next_back_unchecked(
+        &mut self,
+    ) -> Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV> {
+        debug_assert!(self.back.is_some());
+        let back = self.back.as_mut().unwrap();
+        unsafe { back.deallocating_next_back_unchecked() }
+    }
+}
+
 impl<BorrowType: marker::BorrowType, K, V> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
     /// Finds the distinct leaf edges delimiting a specified range in a tree.
     ///
@@ -36,7 +118,7 @@ impl<BorrowType: marker::BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Lea
     /// result will eventually reach the same edge.
     ///
     /// If there are no such edges, i.e., if the tree contains no key within
-    /// the range, returns a pair of empty options.
+    /// the range, returns an empty `front` and `back`.
     ///
     /// # Safety
     /// Unless `BorrowType` is `Immut`, do not use the handles to visit the same
@@ -191,7 +273,7 @@ impl<BorrowType: marker::BorrowType, K, V>
     /// Given a leaf edge handle, returns [`Result::Ok`] with a handle to the neighboring KV
     /// on the left side, which is either in the same leaf node or in an ancestor node.
     /// If the leaf edge is the first one in the tree, returns [`Result::Err`] with the root node.
-    pub fn next_back_kv(
+    fn next_back_kv(
         self,
     ) -> Result<
         Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV>,
@@ -216,7 +298,7 @@ impl<BorrowType: marker::BorrowType, K, V>
     /// Given an internal edge handle, returns [`Result::Ok`] with a handle to the neighboring KV
     /// on the right side, which is either in the same internal node or in an ancestor node.
     /// If the internal edge is the last one in the tree, returns [`Result::Err`] with the root node.
-    pub fn next_kv(
+    fn next_kv(
         self,
     ) -> Result<
         Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::KV>,
@@ -250,7 +332,7 @@ impl<K, V> Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge> {
     /// - The returned KV handle is only valid to access the key and value,
     ///   and only valid until the next call to this method or counterpart
     ///   `deallocating_next_back`.
-    pub unsafe fn deallocating_next(
+    unsafe fn deallocating_next(
         self,
     ) -> Option<(Self, Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV>)>
     {
@@ -316,9 +398,10 @@ impl<'a, K, V> Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Ed
     ///
     /// # Safety
     /// There must be another KV in the direction travelled.
-    pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a V) {
+    unsafe fn next_unchecked(&mut self) -> (&'a K, &'a V) {
         super::mem::replace(self, |leaf_edge| {
             let kv = leaf_edge.next_kv();
+            debug_assert!(kv.is_ok());
             let kv = unsafe { kv.ok().unwrap_unchecked() };
             (kv.next_leaf_edge(), kv.into_kv())
         })
@@ -329,9 +412,10 @@ impl<'a, K, V> Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Ed
     ///
     /// # Safety
     /// There must be another KV in the direction travelled.
-    pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a V) {
+    unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a V) {
         super::mem::replace(self, |leaf_edge| {
             let kv = leaf_edge.next_back_kv();
+            debug_assert!(kv.is_ok());
             let kv = unsafe { kv.ok().unwrap_unchecked() };
             (kv.next_back_leaf_edge(), kv.into_kv())
         })
@@ -344,9 +428,10 @@ impl<'a, K, V> Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::E
     ///
     /// # Safety
     /// There must be another KV in the direction travelled.
-    pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a mut V) {
+    unsafe fn next_unchecked(&mut self) -> (&'a K, &'a mut V) {
         let kv = super::mem::replace(self, |leaf_edge| {
             let kv = leaf_edge.next_kv();
+            debug_assert!(kv.is_ok());
             let kv = unsafe { kv.ok().unwrap_unchecked() };
             (unsafe { ptr::read(&kv) }.next_leaf_edge(), kv)
         });
@@ -359,9 +444,10 @@ impl<'a, K, V> Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::E
     ///
     /// # Safety
     /// There must be another KV in the direction travelled.
-    pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a mut V) {
+    unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a mut V) {
         let kv = super::mem::replace(self, |leaf_edge| {
             let kv = leaf_edge.next_back_kv();
+            debug_assert!(kv.is_ok());
             let kv = unsafe { kv.ok().unwrap_unchecked() };
             (unsafe { ptr::read(&kv) }.next_back_leaf_edge(), kv)
         });
@@ -403,7 +489,7 @@ impl<K, V> Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge> {
     ///
     /// The only safe way to proceed with the updated handle is to compare it, drop it,
     /// or call this method or counterpart `deallocating_next_unchecked` again.
-    pub unsafe fn deallocating_next_back_unchecked(
+    unsafe fn deallocating_next_back_unchecked(
         &mut self,
     ) -> Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV> {
         super::mem::replace(self, |leaf_edge| unsafe {
@@ -508,9 +594,7 @@ impl<BorrowType: marker::BorrowType, K, V>
     }
 
     /// Returns the leaf edge closest to a KV for backward navigation.
-    pub fn next_back_leaf_edge(
-        self,
-    ) -> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
+    fn next_back_leaf_edge(self) -> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
         match self.force() {
             Leaf(leaf_kv) => leaf_kv.left_edge(),
             Internal(internal_kv) => {