Auto merge of #100214 - scottmcm:strict-range, r=thomcc

Optimize `array::IntoIter`

`.into_iter()` on arrays was slower than it needed to be (especially compared to slice iterator) since it uses `Range<usize>`, which needs to handle degenerate ranges like `10..4`.

This PR adds an internal `IndexRange` type that's like `Range<usize>` but with a safety invariant that means it doesn't need to worry about those cases -- it only handles `start <= end` -- and thus can give LLVM more information to optimize better.

I added one simple demonstration of the improvement as a codegen test.

(`vec::IntoIter` uses pointers instead of indexes, so doesn't have this problem, but that only works because its elements are boxed.  `array::IntoIter` can't use pointers because that would keep it from being movable.)

Author: bors

library/core/src/array/iter.rs

@@ -1,10 +1,10 @@
 //! Defines the `IntoIter` owned iterator for arrays.
 
 use crate::{
-    cmp, fmt,
+    fmt,
     iter::{self, ExactSizeIterator, FusedIterator, TrustedLen},
     mem::{self, MaybeUninit},
-    ops::Range,
+    ops::{IndexRange, Range},
     ptr,
 };
 
@@ -29,9 +29,10 @@ pub struct IntoIter<T, const N: usize> {
     /// The elements in `data` that have not been yielded yet.
     ///
     /// Invariants:
-    /// - `alive.start <= alive.end`
     /// - `alive.end <= N`
-    alive: Range<usize>,
+    ///
+    /// (And the `IndexRange` type requires `alive.start <= alive.end`.)
+    alive: IndexRange,
 }
 
 // Note: the `#[rustc_skip_array_during_method_dispatch]` on `trait IntoIterator`
@@ -69,7 +70,7 @@ impl<T, const N: usize> IntoIterator for [T; N] {
         // Until then, we can use `mem::transmute_copy` to create a bitwise copy
         // as a different type, then forget `array` so that it is not dropped.
         unsafe {
-            let iter = IntoIter { data: mem::transmute_copy(&self), alive: 0..N };
+            let iter = IntoIter { data: mem::transmute_copy(&self), alive: IndexRange::zero_to(N) };
             mem::forget(self);
             iter
         }
@@ -147,7 +148,9 @@ impl<T, const N: usize> IntoIter<T, N> {
         buffer: [MaybeUninit<T>; N],
         initialized: Range<usize>,
     ) -> Self {
-        Self { data: buffer, alive: initialized }
+        // SAFETY: one of our safety conditions is that the range is canonical.
+        let alive = unsafe { IndexRange::new_unchecked(initialized.start, initialized.end) };
+        Self { data: buffer, alive }
     }
 
     /// Creates an iterator over `T` which returns no elements.
@@ -283,24 +286,19 @@ impl<T, const N: usize> Iterator for IntoIter<T, N> {
     }
 
     fn advance_by(&mut self, n: usize) -> Result<(), usize> {
-        let len = self.len();
-
-        // The number of elements to drop.  Always in-bounds by construction.
-        let delta = cmp::min(n, len);
-
-        let range_to_drop = self.alive.start..(self.alive.start + delta);
+        let original_len = self.len();
 
-        // Moving the start marks them as conceptually "dropped", so if anything
-        // goes bad then our drop impl won't double-free them.
-        self.alive.start += delta;
+        // This also moves the start, which marks them as conceptually "dropped",
+        // so if anything goes bad then our drop impl won't double-free them.
+        let range_to_drop = self.alive.take_prefix(n);
 
         // SAFETY: These elements are currently initialized, so it's fine to drop them.
         unsafe {
             let slice = self.data.get_unchecked_mut(range_to_drop);
             ptr::drop_in_place(MaybeUninit::slice_assume_init_mut(slice));
         }
 
-        if n > len { Err(len) } else { Ok(()) }
+        if n > original_len { Err(original_len) } else { Ok(()) }
     }
 }
 
@@ -338,24 +336,19 @@ impl<T, const N: usize> DoubleEndedIterator for IntoIter<T, N> {
     }
 
     fn advance_back_by(&mut self, n: usize) -> Result<(), usize> {
-        let len = self.len();
-
-        // The number of elements to drop.  Always in-bounds by construction.
-        let delta = cmp::min(n, len);
-
-        let range_to_drop = (self.alive.end - delta)..self.alive.end;
+        let original_len = self.len();
 
-        // Moving the end marks them as conceptually "dropped", so if anything
-        // goes bad then our drop impl won't double-free them.
-        self.alive.end -= delta;
+        // This also moves the end, which marks them as conceptually "dropped",
+        // so if anything goes bad then our drop impl won't double-free them.
+        let range_to_drop = self.alive.take_suffix(n);
 
         // SAFETY: These elements are currently initialized, so it's fine to drop them.
         unsafe {
             let slice = self.data.get_unchecked_mut(range_to_drop);
             ptr::drop_in_place(MaybeUninit::slice_assume_init_mut(slice));
         }
 
-        if n > len { Err(len) } else { Ok(()) }
+        if n > original_len { Err(original_len) } else { Ok(()) }
     }
 }
 
@@ -372,9 +365,7 @@ impl<T, const N: usize> Drop for IntoIter<T, N> {
 #[stable(feature = "array_value_iter_impls", since = "1.40.0")]
 impl<T, const N: usize> ExactSizeIterator for IntoIter<T, N> {
     fn len(&self) -> usize {
-        // Will never underflow due to the invariant `alive.start <=
-        // alive.end`.
-        self.alive.end - self.alive.start
+        self.alive.len()
     }
     fn is_empty(&self) -> bool {
         self.alive.is_empty()
@@ -396,14 +387,15 @@ impl<T: Clone, const N: usize> Clone for IntoIter<T, N> {
     fn clone(&self) -> Self {
         // Note, we don't really need to match the exact same alive range, so
         // we can just clone into offset 0 regardless of where `self` is.
-        let mut new = Self { data: MaybeUninit::uninit_array(), alive: 0..0 };
+        let mut new = Self { data: MaybeUninit::uninit_array(), alive: IndexRange::zero_to(0) };
 
         // Clone all alive elements.
         for (src, dst) in iter::zip(self.as_slice(), &mut new.data) {
             // Write a clone into the new array, then update its alive range.
             // If cloning panics, we'll correctly drop the previous items.
             dst.write(src.clone());
-            new.alive.end += 1;
+            // This addition cannot overflow as we're iterating a slice
+            new.alive = IndexRange::zero_to(new.alive.end() + 1);
         }
 
         new

library/core/src/lib.rs

@@ -114,6 +114,7 @@
 #![feature(const_fmt_arguments_new)]
 #![feature(const_heap)]
 #![feature(const_convert)]
+#![feature(const_index_range_slice_index)]
 #![feature(const_inherent_unchecked_arith)]
 #![feature(const_int_unchecked_arith)]
 #![feature(const_intrinsic_forget)]

library/core/src/ops/index_range.rs

@@ -0,0 +1,166 @@
+use crate::intrinsics::{assert_unsafe_precondition, unchecked_add, unchecked_sub};
+use crate::iter::{FusedIterator, TrustedLen};
+
+/// Like a `Range<usize>`, but with a safety invariant that `start <= end`.
+///
+/// This means that `end - start` cannot overflow, allowing some μoptimizations.
+///
+/// (Normal `Range` code needs to handle degenerate ranges like `10..0`,
+///  which takes extra checks compared to only handling the canonical form.)
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub(crate) struct IndexRange {
+    start: usize,
+    end: usize,
+}
+
+impl IndexRange {
+    /// # Safety
+    /// - `start <= end`
+    #[inline]
+    pub const unsafe fn new_unchecked(start: usize, end: usize) -> Self {
+        // SAFETY: comparisons on usize are pure
+        unsafe { assert_unsafe_precondition!((start: usize, end: usize) => start <= end) };
+        IndexRange { start, end }
+    }
+
+    #[inline]
+    pub const fn zero_to(end: usize) -> Self {
+        IndexRange { start: 0, end }
+    }
+
+    #[inline]
+    pub const fn start(&self) -> usize {
+        self.start
+    }
+
+    #[inline]
+    pub const fn end(&self) -> usize {
+        self.end
+    }
+
+    #[inline]
+    pub const fn len(&self) -> usize {
+        // SAFETY: By invariant, this cannot wrap
+        unsafe { unchecked_sub(self.end, self.start) }
+    }
+
+    /// # Safety
+    /// - Can only be called when `start < end`, aka when `len > 0`.
+    #[inline]
+    unsafe fn next_unchecked(&mut self) -> usize {
+        debug_assert!(self.start < self.end);
+
+        let value = self.start;
+        // SAFETY: The range isn't empty, so this cannot overflow
+        self.start = unsafe { unchecked_add(value, 1) };
+        value
+    }
+
+    /// # Safety
+    /// - Can only be called when `start < end`, aka when `len > 0`.
+    #[inline]
+    unsafe fn next_back_unchecked(&mut self) -> usize {
+        debug_assert!(self.start < self.end);
+
+        // SAFETY: The range isn't empty, so this cannot overflow
+        let value = unsafe { unchecked_sub(self.end, 1) };
+        self.end = value;
+        value
+    }
+
+    /// Removes the first `n` items from this range, returning them as an `IndexRange`.
+    /// If there are fewer than `n`, then the whole range is returned and
+    /// `self` is left empty.
+    ///
+    /// This is designed to help implement `Iterator::advance_by`.
+    #[inline]
+    pub fn take_prefix(&mut self, n: usize) -> Self {
+        let mid = if n <= self.len() {
+            // SAFETY: We just checked that this will be between start and end,
+            // and thus the addition cannot overflow.
+            unsafe { unchecked_add(self.start, n) }
+        } else {
+            self.end
+        };
+        let prefix = Self { start: self.start, end: mid };
+        self.start = mid;
+        prefix
+    }
+
+    /// Removes the last `n` items from this range, returning them as an `IndexRange`.
+    /// If there are fewer than `n`, then the whole range is returned and
+    /// `self` is left empty.
+    ///
+    /// This is designed to help implement `Iterator::advance_back_by`.
+    #[inline]
+    pub fn take_suffix(&mut self, n: usize) -> Self {
+        let mid = if n <= self.len() {
+            // SAFETY: We just checked that this will be between start and end,
+            // and thus the addition cannot overflow.
+            unsafe { unchecked_sub(self.end, n) }
+        } else {
+            self.start
+        };
+        let suffix = Self { start: mid, end: self.end };
+        self.end = mid;
+        suffix
+    }
+}
+
+impl Iterator for IndexRange {
+    type Item = usize;
+
+    #[inline]
+    fn next(&mut self) -> Option<usize> {
+        if self.len() > 0 {
+            // SAFETY: We just checked that the range is non-empty
+            unsafe { Some(self.next_unchecked()) }
+        } else {
+            None
+        }
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let len = self.len();
+        (len, Some(len))
+    }
+
+    #[inline]
+    fn advance_by(&mut self, n: usize) -> Result<(), usize> {
+        let original_len = self.len();
+        self.take_prefix(n);
+        if n > original_len { Err(original_len) } else { Ok(()) }
+    }
+}
+
+impl DoubleEndedIterator for IndexRange {
+    #[inline]
+    fn next_back(&mut self) -> Option<usize> {
+        if self.len() > 0 {
+            // SAFETY: We just checked that the range is non-empty
+            unsafe { Some(self.next_back_unchecked()) }
+        } else {
+            None
+        }
+    }
+
+    #[inline]
+    fn advance_back_by(&mut self, n: usize) -> Result<(), usize> {
+        let original_len = self.len();
+        self.take_suffix(n);
+        if n > original_len { Err(original_len) } else { Ok(()) }
+    }
+}
+
+impl ExactSizeIterator for IndexRange {
+    #[inline]
+    fn len(&self) -> usize {
+        self.len()
+    }
+}
+
+// SAFETY: Because we only deal in `usize`, our `len` is always perfect.
+unsafe impl TrustedLen for IndexRange {}
+
+impl FusedIterator for IndexRange {}

library/core/src/ops/mod.rs

@@ -146,6 +146,7 @@ mod drop;
 mod function;
 mod generator;
 mod index;
+mod index_range;
 mod range;
 mod try_trait;
 mod unsize;
@@ -178,6 +179,8 @@ pub use self::index::{Index, IndexMut};
 #[stable(feature = "rust1", since = "1.0.0")]
 pub use self::range::{Range, RangeFrom, RangeFull, RangeTo};
 
+pub(crate) use self::index_range::IndexRange;
+
 #[stable(feature = "inclusive_range", since = "1.26.0")]
 pub use self::range::{Bound, RangeBounds, RangeInclusive, RangeToInclusive};