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

BTree: encapsulate LeafRange better & some debug asserts

Looking at iterators again, I think rust-lang#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`

Author: bors

alloc/benches/btree/map.rs

@@ -177,7 +177,7 @@ pub fn iteration_mut_100000(b: &mut Bencher) {
     bench_iteration_mut(b, 100000);
 }
 
-fn bench_first_and_last(b: &mut Bencher, size: i32) {
+fn bench_first_and_last_nightly(b: &mut Bencher, size: i32) {
     let map: BTreeMap<_, _> = (0..size).map(|i| (i, i)).collect();
     b.iter(|| {
         for _ in 0..10 {
@@ -187,19 +187,44 @@ fn bench_first_and_last(b: &mut Bencher, size: i32) {
     });
 }
 
+fn bench_first_and_last_stable(b: &mut Bencher, size: i32) {
+    let map: BTreeMap<_, _> = (0..size).map(|i| (i, i)).collect();
+    b.iter(|| {
+        for _ in 0..10 {
+            black_box(map.iter().next());
+            black_box(map.iter().next_back());
+        }
+    });
+}
+
+#[bench]
+pub fn first_and_last_0_nightly(b: &mut Bencher) {
+    bench_first_and_last_nightly(b, 0);
+}
+
+#[bench]
+pub fn first_and_last_0_stable(b: &mut Bencher) {
+    bench_first_and_last_stable(b, 0);
+}
+
+#[bench]
+pub fn first_and_last_100_nightly(b: &mut Bencher) {
+    bench_first_and_last_nightly(b, 100);
+}
+
 #[bench]
-pub fn first_and_last_0(b: &mut Bencher) {
-    bench_first_and_last(b, 0);
+pub fn first_and_last_100_stable(b: &mut Bencher) {
+    bench_first_and_last_stable(b, 100);
 }
 
 #[bench]
-pub fn first_and_last_100(b: &mut Bencher) {
-    bench_first_and_last(b, 100);
+pub fn first_and_last_10k_nightly(b: &mut Bencher) {
+    bench_first_and_last_nightly(b, 10_000);
 }
 
 #[bench]
-pub fn first_and_last_10k(b: &mut Bencher) {
-    bench_first_and_last(b, 10_000);
+pub fn first_and_last_10k_stable(b: &mut Bencher) {
+    bench_first_and_last_stable(b, 10_000);
 }
 
 const BENCH_RANGE_SIZE: i32 = 145;

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> {