Skip to content

Commit

Permalink
Auto merge of rust-lang#136771 - scottmcm:poke-slice-iter-next, r=joboet
Browse files Browse the repository at this point in the history 
Simplify `slice::Iter::next` enough that it inlines

Inspired by this zulip conversation: <https://rust-lang.zulipchat.com/#narrow/channel/189540-t-compiler.2Fwg-mir-opt/topic/Feedback.20on.20a.20MIR.20optimization.20idea/near/498579990>

~~Draft for now because it needs rust-lang#136735 to get the codegen tests to pass.~~
  • Loading branch information
@bors
bors committed Feb 20, 2025
2 parents 28b83ee + 7add358 commit f04bbc6
Show file tree
Hide file tree
Showing 7 changed files with 763 additions and 152 deletions.
35 changes: 29 additions & 6 deletions library/core/src/slice/iter/macros.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -154,16 +154,39 @@ macro_rules! iterator {

#[inline]
fn next(&mut self) -> Option<$elem> {
// could be implemented with slices, but this avoids bounds checks
// intentionally not using the helpers because this is
// one of the most mono'd things in the library.

// SAFETY: The call to `next_unchecked` is
// safe since we check if the iterator is empty first.
let ptr = self.ptr;
let end_or_len = self.end_or_len;
// SAFETY: See inner comments. (For some reason having multiple
// block breaks inlining this -- if you can fix that please do!)
unsafe {
if is_empty!(self) {
None
if T::IS_ZST {
let len = end_or_len.addr();
if len == 0 {
return None;
}
// SAFETY: just checked that it's not zero, so subtracting one
// cannot wrap. (Ideally this would be `checked_sub`, which
// does the same thing internally, but as of 2025-02 that
// doesn't optimize quite as small in MIR.)
self.end_or_len = without_provenance_mut(len.unchecked_sub(1));
} else {
Some(self.next_unchecked())
// SAFETY: by type invariant, the `end_or_len` field is always
// non-null for a non-ZST pointee. (This transmute ensures we
// get `!nonnull` metadata on the load of the field.)
if ptr == crate::intrinsics::transmute::<$ptr, NonNull<T>>(end_or_len) {
return None;
}
// SAFETY: since it's not empty, per the check above, moving
// forward one keeps us inside the slice, and this is valid.
self.ptr = ptr.add(1);
}
// SAFETY: Now that we know it wasn't empty and we've moved past
// the first one (to avoid giving a duplicate `&mut` next time),
// we can give out a reference to it.
Some({ptr}.$into_ref())
}
}

Expand Down
6 changes: 3 additions & 3 deletions tests/codegen/slice-iter-nonnull.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,11 +14,11 @@
// CHECK-LABEL: @slice_iter_next(
#[no_mangle]
pub fn slice_iter_next<'a>(it: &mut std::slice::Iter<'a, u32>) -> Option<&'a u32> {
// CHECK: %[[ENDP:.+]] = getelementptr inbounds{{( nuw)?}} i8, ptr %it, {{i32 4|i64 8}}
// CHECK: %[[END:.+]] = load ptr, ptr %[[ENDP]]
// CHECK: %[[START:.+]] = load ptr, ptr %it,
// CHECK-SAME: !nonnull
// CHECK-SAME: !noundef
// CHECK: %[[START:.+]] = load ptr, ptr %it,
// CHECK: %[[ENDP:.+]] = getelementptr inbounds{{( nuw)?}} i8, ptr %it, {{i32 4|i64 8}}
// CHECK: %[[END:.+]] = load ptr, ptr %[[ENDP]]
// CHECK-SAME: !nonnull
// CHECK-SAME: !noundef
// CHECK: icmp eq ptr %[[START]], %[[END]]
Expand Down
248 changes: 180 additions & 68 deletions tests/mir-opt/pre-codegen/slice_iter.enumerated_loop.PreCodegen.after.panic-abort.mir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,28 +4,30 @@ fn enumerated_loop(_1: &[T], _2: impl Fn(usize, &T)) -> () {
debug slice => _1;
debug f => _2;
let mut _0: ();
let mut _11: std::slice::Iter<'_, T>;
let mut _12: std::iter::Enumerate<std::slice::Iter<'_, T>>;
let mut _13: std::iter::Enumerate<std::slice::Iter<'_, T>>;
let mut _21: std::option::Option<(usize, &T)>;
let mut _24: &impl Fn(usize, &T);
let mut _25: (usize, &T);
let _26: ();
let mut _11: std::ptr::NonNull<T>;
let mut _12: *const T;
let mut _13: usize;
let mut _32: std::option::Option<(usize, &T)>;
let mut _35: &impl Fn(usize, &T);
let mut _36: (usize, &T);
let _37: ();
scope 1 {
debug iter => _13;
let _22: usize;
let _23: &T;
debug (((iter: Enumerate<std::slice::Iter<'_, T>>).0: std::slice::Iter<'_, T>).0: std::ptr::NonNull<T>) => _11;
debug (((iter: Enumerate<std::slice::Iter<'_, T>>).0: std::slice::Iter<'_, T>).1: *const T) => _12;
debug (((iter: Enumerate<std::slice::Iter<'_, T>>).0: std::slice::Iter<'_, T>).2: std::marker::PhantomData<&T>) => const ZeroSized: PhantomData<&T>;
debug ((iter: Enumerate<std::slice::Iter<'_, T>>).1: usize) => _13;
let _33: usize;
let _34: &T;
scope 2 {
debug i => _22;
debug x => _23;
debug i => _33;
debug x => _34;
}
scope 19 (inlined <Enumerate<std::slice::Iter<'_, T>> as Iterator>::next) {
let mut _14: &mut std::slice::Iter<'_, T>;
let mut _15: std::option::Option<&T>;
let mut _19: (usize, bool);
let mut _20: (usize, &T);
let mut _27: std::option::Option<&T>;
let mut _30: (usize, bool);
let mut _31: (usize, &T);
scope 20 {
let _18: usize;
let _29: usize;
scope 25 {
}
}
Expand All @@ -40,11 +42,59 @@ fn enumerated_loop(_1: &[T], _2: impl Fn(usize, &T)) -> () {
}
}
scope 26 (inlined <Option<&T> as Try>::branch) {
let mut _16: isize;
let _17: &T;
let _28: &T;
scope 27 {
}
}
scope 29 (inlined <std::slice::Iter<'_, T> as Iterator>::next) {
let _14: std::ptr::NonNull<T>;
let _16: std::ptr::NonNull<T>;
let mut _19: bool;
let mut _22: std::ptr::NonNull<T>;
let mut _24: usize;
let _26: &T;
scope 30 {
let _15: *const T;
scope 31 {
let _23: usize;
scope 32 {
scope 35 (inlined core::num::<impl usize>::unchecked_sub) {
scope 36 (inlined core::ub_checks::check_language_ub) {
scope 37 (inlined core::ub_checks::check_language_ub::runtime) {
}
}
}
scope 38 (inlined without_provenance_mut::<T>) {
}
}
scope 33 (inlined std::ptr::const_ptr::<impl *const T>::addr) {
scope 34 (inlined std::ptr::const_ptr::<impl *const T>::cast::<()>) {
}
}
scope 39 (inlined <NonNull<T> as PartialEq>::eq) {
let mut _17: *mut T;
let mut _18: *mut T;
scope 40 (inlined NonNull::<T>::as_ptr) {
}
scope 41 (inlined NonNull::<T>::as_ptr) {
}
}
scope 42 (inlined NonNull::<T>::add) {
let mut _20: *const T;
let mut _21: *const T;
scope 43 (inlined NonNull::<T>::as_ptr) {
}
}
scope 44 (inlined NonNull::<T>::as_ref::<'_>) {
let _25: *const T;
scope 45 (inlined NonNull::<T>::as_ptr) {
}
scope 46 (inlined std::ptr::mut_ptr::<impl *mut T>::cast_const) {
}
}
}
}
}
}
}
scope 3 (inlined core::slice::<impl [T]>::iter) {
Expand Down Expand Up @@ -89,9 +139,7 @@ fn enumerated_loop(_1: &[T], _2: impl Fn(usize, &T)) -> () {
}

bb0: {
StorageLive(_11);
StorageLive(_3);
StorageLive(_6);
StorageLive(_4);
_3 = PtrMetadata(copy _1);
_4 = &raw const (*_1);
Expand Down Expand Up @@ -120,86 +168,150 @@ fn enumerated_loop(_1: &[T], _2: impl Fn(usize, &T)) -> () {
}

bb3: {
StorageLive(_10);
_10 = copy _9;
_11 = std::slice::Iter::<'_, T> { ptr: copy _6, end_or_len: move _10, _marker: const ZeroSized: PhantomData<&T> };
StorageDead(_10);
StorageDead(_9);
StorageDead(_4);
StorageDead(_6);
StorageDead(_3);
_12 = Enumerate::<std::slice::Iter<'_, T>> { iter: copy _11, count: const 0_usize };
StorageDead(_11);
StorageLive(_11);
StorageLive(_12);
StorageLive(_13);
_13 = copy _12;
_11 = copy _6;
_12 = copy _10;
_13 = const 0_usize;
goto -> bb4;
}

bb4: {
StorageLive(_21);
StorageLive(_18);
StorageLive(_19);
StorageLive(_15);
StorageLive(_32);
StorageLive(_29);
StorageLive(_30);
StorageLive(_27);
StorageLive(_14);
_14 = &mut (_13.0: std::slice::Iter<'_, T>);
_15 = <std::slice::Iter<'_, T> as Iterator>::next(move _14) -> [return: bb5, unwind unreachable];
StorageLive(_15);
StorageLive(_23);
StorageLive(_24);
StorageLive(_16);
StorageLive(_26);
_14 = copy _11;
_15 = copy _12;
switchInt(const <T as std::mem::SizedTypeProperties>::IS_ZST) -> [0: bb5, otherwise: bb8];
}

bb5: {
StorageDead(_14);
StorageLive(_16);
_16 = discriminant(_15);
switchInt(move _16) -> [0: bb6, 1: bb8, otherwise: bb11];
StorageLive(_19);
_16 = copy _15 as std::ptr::NonNull<T> (Transmute);
StorageLive(_17);
_17 = copy _14 as *mut T (Transmute);
StorageLive(_18);
_18 = copy _16 as *mut T (Transmute);
_19 = Eq(move _17, move _18);
StorageDead(_18);
StorageDead(_17);
switchInt(move _19) -> [0: bb6, otherwise: bb7];
}

bb6: {
StorageDead(_16);
StorageDead(_15);
StorageDead(_19);
StorageDead(_18);
StorageLive(_22);
StorageLive(_21);
StorageLive(_20);
_20 = copy _14 as *const T (Transmute);
_21 = Offset(move _20, const 1_usize);
StorageDead(_20);
_22 = NonNull::<T> { pointer: move _21 };
StorageDead(_21);
StorageDead(_13);
drop(_2) -> [return: bb7, unwind unreachable];
_11 = move _22;
StorageDead(_22);
goto -> bb13;
}

bb7: {
return;
StorageDead(_19);
StorageDead(_26);
StorageDead(_16);
StorageDead(_24);
StorageDead(_23);
StorageDead(_15);
StorageDead(_14);
goto -> bb10;
}

bb8: {
_17 = move ((_15 as Some).0: &T);
StorageDead(_16);
StorageDead(_15);
_18 = copy (_13.1: usize);
_19 = AddWithOverflow(copy (_13.1: usize), const 1_usize);
assert(!move (_19.1: bool), "attempt to compute `{} + {}`, which would overflow", copy (_13.1: usize), const 1_usize) -> [success: bb9, unwind unreachable];
_23 = copy _15 as usize (Transmute);
switchInt(copy _23) -> [0: bb9, otherwise: bb12];
}

bb9: {
(_13.1: usize) = move (_19.0: usize);
StorageLive(_20);
_20 = (copy _18, copy _17);
_21 = Option::<(usize, &T)>::Some(move _20);
StorageDead(_20);
StorageDead(_19);
StorageDead(_18);
_22 = copy (((_21 as Some).0: (usize, &T)).0: usize);
_23 = copy (((_21 as Some).0: (usize, &T)).1: &T);
StorageLive(_24);
_24 = &_2;
StorageLive(_25);
_25 = (copy _22, copy _23);
_26 = <impl Fn(usize, &T) as Fn<(usize, &T)>>::call(move _24, move _25) -> [return: bb10, unwind unreachable];
StorageDead(_26);
StorageDead(_16);
StorageDead(_24);
StorageDead(_23);
StorageDead(_15);
StorageDead(_14);
goto -> bb10;
}

bb10: {
StorageDead(_27);
StorageDead(_30);
StorageDead(_29);
StorageDead(_32);
StorageDead(_11);
StorageDead(_12);
StorageDead(_13);
drop(_2) -> [return: bb11, unwind unreachable];
}

bb11: {
return;
}

bb12: {
_24 = SubUnchecked(copy _23, const 1_usize);
_12 = copy _24 as *const T (Transmute);
goto -> bb13;
}

bb13: {
StorageLive(_25);
_25 = copy _14 as *const T (Transmute);
_26 = &(*_25);
StorageDead(_25);
_27 = Option::<&T>::Some(copy _26);
StorageDead(_26);
StorageDead(_16);
StorageDead(_24);
StorageDead(_21);
goto -> bb4;
StorageDead(_23);
StorageDead(_15);
StorageDead(_14);
_28 = move ((_27 as Some).0: &T);
StorageDead(_27);
_29 = copy _13;
_30 = AddWithOverflow(copy _13, const 1_usize);
assert(!move (_30.1: bool), "attempt to compute `{} + {}`, which would overflow", copy _13, const 1_usize) -> [success: bb14, unwind unreachable];
}

bb11: {
unreachable;
bb14: {
_13 = move (_30.0: usize);
StorageLive(_31);
_31 = (copy _29, copy _28);
_32 = Option::<(usize, &T)>::Some(move _31);
StorageDead(_31);
StorageDead(_30);
StorageDead(_29);
_33 = copy (((_32 as Some).0: (usize, &T)).0: usize);
_34 = copy (((_32 as Some).0: (usize, &T)).1: &T);
StorageLive(_35);
_35 = &_2;
StorageLive(_36);
_36 = (copy _33, copy _34);
_37 = <impl Fn(usize, &T) as Fn<(usize, &T)>>::call(move _35, move _36) -> [return: bb15, unwind unreachable];
}

bb15: {
StorageDead(_36);
StorageDead(_35);
StorageDead(_32);
goto -> bb4;
}
}
Loading

0 comments on commit f04bbc6

Please sign in to comment.