Skip to content

Commit 88b3b52

Browse files
committed
Auto merge of #141086 - a1phyr:spec_advance_by, r=jhpratt
Implement `advance_by` via `try_fold` for `Sized` iterators When `try_fold` is overriden, it is usually easier for compilers to optimize. Example difference: https://iter.godbolt.org/z/z8cEfnKro
2 parents aa57e46 + 5240cd3 commit 88b3b52

File tree

1 file changed

+31
-5
lines changed

1 file changed

+31
-5
lines changed

library/core/src/iter/traits/iterator.rs

Lines changed: 31 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -294,13 +294,39 @@ pub trait Iterator {
294294
#[inline]
295295
#[unstable(feature = "iter_advance_by", reason = "recently added", issue = "77404")]
296296
fn advance_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
297-
for i in 0..n {
298-
if self.next().is_none() {
299-
// SAFETY: `i` is always less than `n`.
300-
return Err(unsafe { NonZero::new_unchecked(n - i) });
297+
/// Helper trait to specialize `advance_by` via `try_fold` for `Sized` iterators.
298+
trait SpecAdvanceBy {
299+
fn spec_advance_by(&mut self, n: usize) -> Result<(), NonZero<usize>>;
300+
}
301+
302+
impl<I: Iterator + ?Sized> SpecAdvanceBy for I {
303+
default fn spec_advance_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
304+
for i in 0..n {
305+
if self.next().is_none() {
306+
// SAFETY: `i` is always less than `n`.
307+
return Err(unsafe { NonZero::new_unchecked(n - i) });
308+
}
309+
}
310+
Ok(())
311+
}
312+
}
313+
314+
impl<I: Iterator> SpecAdvanceBy for I {
315+
fn spec_advance_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
316+
let Some(n) = NonZero::new(n) else {
317+
return Ok(());
318+
};
319+
320+
let res = self.try_fold(n, |n, _| NonZero::new(n.get() - 1));
321+
322+
match res {
323+
None => Ok(()),
324+
Some(n) => Err(n),
325+
}
301326
}
302327
}
303-
Ok(())
328+
329+
self.spec_advance_by(n)
304330
}
305331

306332
/// Returns the `n`th element of the iterator.

0 commit comments

Comments
 (0)