Skip to content

Commit

Permalink
Auto merge of #100214 - scottmcm:strict-range, r=thomcc
Browse files Browse the repository at this point in the history
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.)
  • Loading branch information
bors committed Sep 21, 2022
2 parents 7f8a5a4 + d9216c5 commit 94c5afe
Show file tree
Hide file tree
Showing 5 changed files with 268 additions and 31 deletions.
54 changes: 23 additions & 31 deletions core/src/array/iter.rs
Original file line number Diff line number Diff line change
@@ -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,
};

Expand All @@ -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`
Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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(()) }
}
}

Expand Down Expand Up @@ -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(()) }
}
}

Expand All @@ -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()
Expand All @@ -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
Expand Down
1 change: 1 addition & 0 deletions core/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down
166 changes: 166 additions & 0 deletions core/src/ops/index_range.rs
Original file line number Diff line number Diff line change
@@ -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 {}
3 changes: 3 additions & 0 deletions core/src/ops/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -146,6 +146,7 @@ mod drop;
mod function;
mod generator;
mod index;
mod index_range;
mod range;
mod try_trait;
mod unsize;
Expand Down Expand Up @@ -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};

Expand Down
Loading

0 comments on commit 94c5afe

Please sign in to comment.