Skip to content

Commit f04bbc6

Browse files
borsbors
committed
Auto merge of #136771 - scottmcm:poke-slice-iter-next, r=joboet
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 #136735 to get the codegen tests to pass.~~
2 parents 28b83ee + 7add358 commit f04bbc6

7 files changed

+763
-152
lines changed

‎library/core/src/slice/iter/macros.rs‎

Lines changed: 29 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -154,16 +154,39 @@ macro_rules! iterator {
154154

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

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

‎tests/codegen/slice-iter-nonnull.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,11 @@
1414
// CHECK-LABEL: @slice_iter_next(
1515
#[no_mangle]
1616
pub fn slice_iter_next<'a>(it: &mut std::slice::Iter<'a, u32>) -> Option<&'a u32> {
17-
// CHECK: %[[ENDP:.+]] = getelementptr inbounds{{( nuw)?}} i8, ptr %it, {{i32 4|i64 8}}
18-
// CHECK: %[[END:.+]] = load ptr, ptr %[[ENDP]]
17+
// CHECK: %[[START:.+]] = load ptr, ptr %it,
1918
// CHECK-SAME: !nonnull
2019
// CHECK-SAME: !noundef
21-
// CHECK: %[[START:.+]] = load ptr, ptr %it,
20+
// CHECK: %[[ENDP:.+]] = getelementptr inbounds{{( nuw)?}} i8, ptr %it, {{i32 4|i64 8}}
21+
// CHECK: %[[END:.+]] = load ptr, ptr %[[ENDP]]
2222
// CHECK-SAME: !nonnull
2323
// CHECK-SAME: !noundef
2424
// CHECK: icmp eq ptr %[[START]], %[[END]]

‎tests/mir-opt/pre-codegen/slice_iter.enumerated_loop.PreCodegen.after.panic-abort.mir‎

Lines changed: 180 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -4,28 +4,30 @@ fn enumerated_loop(_1: &[T], _2: impl Fn(usize, &T)) -> () {
44
debug slice => _1;
55
debug f => _2;
66
let mut _0: ();
7-
let mut _11: std::slice::Iter<'_, T>;
8-
let mut _12: std::iter::Enumerate<std::slice::Iter<'_, T>>;
9-
let mut _13: std::iter::Enumerate<std::slice::Iter<'_, T>>;
10-
let mut _21: std::option::Option<(usize, &T)>;
11-
let mut _24: &impl Fn(usize, &T);
12-
let mut _25: (usize, &T);
13-
let _26: ();
7+
let mut _11: std::ptr::NonNull<T>;
8+
let mut _12: *const T;
9+
let mut _13: usize;
10+
let mut _32: std::option::Option<(usize, &T)>;
11+
let mut _35: &impl Fn(usize, &T);
12+
let mut _36: (usize, &T);
13+
let _37: ();
1414
scope 1 {
15-
debug iter => _13;
16-
let _22: usize;
17-
let _23: &T;
15+
debug (((iter: Enumerate<std::slice::Iter<'_, T>>).0: std::slice::Iter<'_, T>).0: std::ptr::NonNull<T>) => _11;
16+
debug (((iter: Enumerate<std::slice::Iter<'_, T>>).0: std::slice::Iter<'_, T>).1: *const T) => _12;
17+
debug (((iter: Enumerate<std::slice::Iter<'_, T>>).0: std::slice::Iter<'_, T>).2: std::marker::PhantomData<&T>) => const ZeroSized: PhantomData<&T>;
18+
debug ((iter: Enumerate<std::slice::Iter<'_, T>>).1: usize) => _13;
19+
let _33: usize;
20+
let _34: &T;
1821
scope 2 {
19-
debug i => _22;
20-
debug x => _23;
22+
debug i => _33;
23+
debug x => _34;
2124
}
2225
scope 19 (inlined <Enumerate<std::slice::Iter<'_, T>> as Iterator>::next) {
23-
let mut _14: &mut std::slice::Iter<'_, T>;
24-
let mut _15: std::option::Option<&T>;
25-
let mut _19: (usize, bool);
26-
let mut _20: (usize, &T);
26+
let mut _27: std::option::Option<&T>;
27+
let mut _30: (usize, bool);
28+
let mut _31: (usize, &T);
2729
scope 20 {
28-
let _18: usize;
30+
let _29: usize;
2931
scope 25 {
3032
}
3133
}
@@ -40,11 +42,59 @@ fn enumerated_loop(_1: &[T], _2: impl Fn(usize, &T)) -> () {
4042
}
4143
}
4244
scope 26 (inlined <Option<&T> as Try>::branch) {
43-
let mut _16: isize;
44-
let _17: &T;
45+
let _28: &T;
4546
scope 27 {
4647
}
4748
}
49+
scope 29 (inlined <std::slice::Iter<'_, T> as Iterator>::next) {
50+
let _14: std::ptr::NonNull<T>;
51+
let _16: std::ptr::NonNull<T>;
52+
let mut _19: bool;
53+
let mut _22: std::ptr::NonNull<T>;
54+
let mut _24: usize;
55+
let _26: &T;
56+
scope 30 {
57+
let _15: *const T;
58+
scope 31 {
59+
let _23: usize;
60+
scope 32 {
61+
scope 35 (inlined core::num::<impl usize>::unchecked_sub) {
62+
scope 36 (inlined core::ub_checks::check_language_ub) {
63+
scope 37 (inlined core::ub_checks::check_language_ub::runtime) {
64+
}
65+
}
66+
}
67+
scope 38 (inlined without_provenance_mut::<T>) {
68+
}
69+
}
70+
scope 33 (inlined std::ptr::const_ptr::<impl *const T>::addr) {
71+
scope 34 (inlined std::ptr::const_ptr::<impl *const T>::cast::<()>) {
72+
}
73+
}
74+
scope 39 (inlined <NonNull<T> as PartialEq>::eq) {
75+
let mut _17: *mut T;
76+
let mut _18: *mut T;
77+
scope 40 (inlined NonNull::<T>::as_ptr) {
78+
}
79+
scope 41 (inlined NonNull::<T>::as_ptr) {
80+
}
81+
}
82+
scope 42 (inlined NonNull::<T>::add) {
83+
let mut _20: *const T;
84+
let mut _21: *const T;
85+
scope 43 (inlined NonNull::<T>::as_ptr) {
86+
}
87+
}
88+
scope 44 (inlined NonNull::<T>::as_ref::<'_>) {
89+
let _25: *const T;
90+
scope 45 (inlined NonNull::<T>::as_ptr) {
91+
}
92+
scope 46 (inlined std::ptr::mut_ptr::<impl *mut T>::cast_const) {
93+
}
94+
}
95+
}
96+
}
97+
}
4898
}
4999
}
50100
scope 3 (inlined core::slice::<impl [T]>::iter) {
@@ -89,9 +139,7 @@ fn enumerated_loop(_1: &[T], _2: impl Fn(usize, &T)) -> () {
89139
}
90140

91141
bb0: {
92-
StorageLive(_11);
93142
StorageLive(_3);
94-
StorageLive(_6);
95143
StorageLive(_4);
96144
_3 = PtrMetadata(copy _1);
97145
_4 = &raw const (*_1);
@@ -120,86 +168,150 @@ fn enumerated_loop(_1: &[T], _2: impl Fn(usize, &T)) -> () {
120168
}
121169

122170
bb3: {
123-
StorageLive(_10);
124171
_10 = copy _9;
125-
_11 = std::slice::Iter::<'_, T> { ptr: copy _6, end_or_len: move _10, _marker: const ZeroSized: PhantomData<&T> };
126-
StorageDead(_10);
127172
StorageDead(_9);
128173
StorageDead(_4);
129-
StorageDead(_6);
130174
StorageDead(_3);
131-
_12 = Enumerate::<std::slice::Iter<'_, T>> { iter: copy _11, count: const 0_usize };
132-
StorageDead(_11);
175+
StorageLive(_11);
176+
StorageLive(_12);
133177
StorageLive(_13);
134-
_13 = copy _12;
178+
_11 = copy _6;
179+
_12 = copy _10;
180+
_13 = const 0_usize;
135181
goto -> bb4;
136182
}
137183

138184
bb4: {
139-
StorageLive(_21);
140-
StorageLive(_18);
141-
StorageLive(_19);
142-
StorageLive(_15);
185+
StorageLive(_32);
186+
StorageLive(_29);
187+
StorageLive(_30);
188+
StorageLive(_27);
143189
StorageLive(_14);
144-
_14 = &mut (_13.0: std::slice::Iter<'_, T>);
145-
_15 = <std::slice::Iter<'_, T> as Iterator>::next(move _14) -> [return: bb5, unwind unreachable];
190+
StorageLive(_15);
191+
StorageLive(_23);
192+
StorageLive(_24);
193+
StorageLive(_16);
194+
StorageLive(_26);
195+
_14 = copy _11;
196+
_15 = copy _12;
197+
switchInt(const <T as std::mem::SizedTypeProperties>::IS_ZST) -> [0: bb5, otherwise: bb8];
146198
}
147199

148200
bb5: {
149-
StorageDead(_14);
150-
StorageLive(_16);
151-
_16 = discriminant(_15);
152-
switchInt(move _16) -> [0: bb6, 1: bb8, otherwise: bb11];
201+
StorageLive(_19);
202+
_16 = copy _15 as std::ptr::NonNull<T> (Transmute);
203+
StorageLive(_17);
204+
_17 = copy _14 as *mut T (Transmute);
205+
StorageLive(_18);
206+
_18 = copy _16 as *mut T (Transmute);
207+
_19 = Eq(move _17, move _18);
208+
StorageDead(_18);
209+
StorageDead(_17);
210+
switchInt(move _19) -> [0: bb6, otherwise: bb7];
153211
}
154212

155213
bb6: {
156-
StorageDead(_16);
157-
StorageDead(_15);
158214
StorageDead(_19);
159-
StorageDead(_18);
215+
StorageLive(_22);
216+
StorageLive(_21);
217+
StorageLive(_20);
218+
_20 = copy _14 as *const T (Transmute);
219+
_21 = Offset(move _20, const 1_usize);
220+
StorageDead(_20);
221+
_22 = NonNull::<T> { pointer: move _21 };
160222
StorageDead(_21);
161-
StorageDead(_13);
162-
drop(_2) -> [return: bb7, unwind unreachable];
223+
_11 = move _22;
224+
StorageDead(_22);
225+
goto -> bb13;
163226
}
164227

165228
bb7: {
166-
return;
229+
StorageDead(_19);
230+
StorageDead(_26);
231+
StorageDead(_16);
232+
StorageDead(_24);
233+
StorageDead(_23);
234+
StorageDead(_15);
235+
StorageDead(_14);
236+
goto -> bb10;
167237
}
168238

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

178244
bb9: {
179-
(_13.1: usize) = move (_19.0: usize);
180-
StorageLive(_20);
181-
_20 = (copy _18, copy _17);
182-
_21 = Option::<(usize, &T)>::Some(move _20);
183-
StorageDead(_20);
184-
StorageDead(_19);
185-
StorageDead(_18);
186-
_22 = copy (((_21 as Some).0: (usize, &T)).0: usize);
187-
_23 = copy (((_21 as Some).0: (usize, &T)).1: &T);
188-
StorageLive(_24);
189-
_24 = &_2;
190-
StorageLive(_25);
191-
_25 = (copy _22, copy _23);
192-
_26 = <impl Fn(usize, &T) as Fn<(usize, &T)>>::call(move _24, move _25) -> [return: bb10, unwind unreachable];
245+
StorageDead(_26);
246+
StorageDead(_16);
247+
StorageDead(_24);
248+
StorageDead(_23);
249+
StorageDead(_15);
250+
StorageDead(_14);
251+
goto -> bb10;
193252
}
194253

195254
bb10: {
255+
StorageDead(_27);
256+
StorageDead(_30);
257+
StorageDead(_29);
258+
StorageDead(_32);
259+
StorageDead(_11);
260+
StorageDead(_12);
261+
StorageDead(_13);
262+
drop(_2) -> [return: bb11, unwind unreachable];
263+
}
264+
265+
bb11: {
266+
return;
267+
}
268+
269+
bb12: {
270+
_24 = SubUnchecked(copy _23, const 1_usize);
271+
_12 = copy _24 as *const T (Transmute);
272+
goto -> bb13;
273+
}
274+
275+
bb13: {
276+
StorageLive(_25);
277+
_25 = copy _14 as *const T (Transmute);
278+
_26 = &(*_25);
196279
StorageDead(_25);
280+
_27 = Option::<&T>::Some(copy _26);
281+
StorageDead(_26);
282+
StorageDead(_16);
197283
StorageDead(_24);
198-
StorageDead(_21);
199-
goto -> bb4;
284+
StorageDead(_23);
285+
StorageDead(_15);
286+
StorageDead(_14);
287+
_28 = move ((_27 as Some).0: &T);
288+
StorageDead(_27);
289+
_29 = copy _13;
290+
_30 = AddWithOverflow(copy _13, const 1_usize);
291+
assert(!move (_30.1: bool), "attempt to compute `{} + {}`, which would overflow", copy _13, const 1_usize) -> [success: bb14, unwind unreachable];
200292
}
201293

202-
bb11: {
203-
unreachable;
294+
bb14: {
295+
_13 = move (_30.0: usize);
296+
StorageLive(_31);
297+
_31 = (copy _29, copy _28);
298+
_32 = Option::<(usize, &T)>::Some(move _31);
299+
StorageDead(_31);
300+
StorageDead(_30);
301+
StorageDead(_29);
302+
_33 = copy (((_32 as Some).0: (usize, &T)).0: usize);
303+
_34 = copy (((_32 as Some).0: (usize, &T)).1: &T);
304+
StorageLive(_35);
305+
_35 = &_2;
306+
StorageLive(_36);
307+
_36 = (copy _33, copy _34);
308+
_37 = <impl Fn(usize, &T) as Fn<(usize, &T)>>::call(move _35, move _36) -> [return: bb15, unwind unreachable];
309+
}
310+
311+
bb15: {
312+
StorageDead(_36);
313+
StorageDead(_35);
314+
StorageDead(_32);
315+
goto -> bb4;
204316
}
205317
}

0 commit comments

Comments
 (0)