Auto merge of #111379 - nyurik:intersperse-speed-up, r=cuviper

Boost iterator intersperse(_with) performance

I did some benchmark digging into the `intersperse` and `intersperse_with` code as part of [this discussion](https://internals.rust-lang.org/t/add-iterate-with-separators-iterator-function/18781/13), and as a result I optimized them a bit, without relying on the peekable iterator.

See also [full benchmark repo](https://github.com/nyurik/intersperse_perf)

Benchmarks show near 2x performance improvements with the simple `sum` [benchmarks](https://gist.github.com/nyurik/68b6c9b3d90f0d14746d4186bf8fa1e2):
![image](https://user-images.githubusercontent.com/1641515/237005195-16aebef4-9eed-4514-8b7c-da1d1f5bd9e0.png)
This commit is contained in:
bors 2024-01-27 08:16:49 +00:00
commit 8b6a431b3d

View File

@ -1,4 +1,5 @@
use super::Peekable; use crate::fmt;
use crate::iter::{Fuse, FusedIterator};
/// An iterator adapter that places a separator between all elements. /// An iterator adapter that places a separator between all elements.
/// ///
@ -10,9 +11,18 @@ pub struct Intersperse<I: Iterator>
where where
I::Item: Clone, I::Item: Clone,
{ {
started: bool,
separator: I::Item, separator: I::Item,
iter: Peekable<I>, next_item: Option<I::Item>,
needs_sep: bool, iter: Fuse<I>,
}
#[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")]
impl<I> FusedIterator for Intersperse<I>
where
I: FusedIterator,
I::Item: Clone,
{
} }
impl<I: Iterator> Intersperse<I> impl<I: Iterator> Intersperse<I>
@ -20,7 +30,7 @@ impl<I: Iterator> Intersperse<I>
I::Item: Clone, I::Item: Clone,
{ {
pub(in crate::iter) fn new(iter: I, separator: I::Item) -> Self { pub(in crate::iter) fn new(iter: I, separator: I::Item) -> Self {
Self { iter: iter.peekable(), separator, needs_sep: false } Self { started: false, separator, next_item: None, iter: iter.fuse() }
} }
} }
@ -33,27 +43,43 @@ impl<I> Iterator for Intersperse<I>
type Item = I::Item; type Item = I::Item;
#[inline] #[inline]
fn next(&mut self) -> Option<I::Item> { fn next(&mut self) -> Option<Self::Item> {
if self.needs_sep && self.iter.peek().is_some() { if self.started {
self.needs_sep = false; if let Some(v) = self.next_item.take() {
Some(v)
} else {
let next_item = self.iter.next();
if next_item.is_some() {
self.next_item = next_item;
Some(self.separator.clone()) Some(self.separator.clone())
} else { } else {
self.needs_sep = true; None
}
}
} else {
self.started = true;
self.iter.next() self.iter.next()
} }
} }
fn size_hint(&self) -> (usize, Option<usize>) {
intersperse_size_hint(&self.iter, self.started, self.next_item.is_some())
}
fn fold<B, F>(self, init: B, f: F) -> B fn fold<B, F>(self, init: B, f: F) -> B
where where
Self: Sized, Self: Sized,
F: FnMut(B, Self::Item) -> B, F: FnMut(B, Self::Item) -> B,
{ {
let separator = self.separator; let separator = self.separator;
intersperse_fold(self.iter, init, f, move || separator.clone(), self.needs_sep) intersperse_fold(
} self.iter,
init,
fn size_hint(&self) -> (usize, Option<usize>) { f,
intersperse_size_hint(&self.iter, self.needs_sep) move || separator.clone(),
self.started,
self.next_item,
)
} }
} }
@ -66,39 +92,50 @@ pub struct IntersperseWith<I, G>
where where
I: Iterator, I: Iterator,
{ {
started: bool,
separator: G, separator: G,
iter: Peekable<I>, next_item: Option<I::Item>,
needs_sep: bool, iter: Fuse<I>,
} }
#[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")] #[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")]
impl<I, G> crate::fmt::Debug for IntersperseWith<I, G> impl<I, G> FusedIterator for IntersperseWith<I, G>
where where
I: Iterator + crate::fmt::Debug, I: FusedIterator,
I::Item: crate::fmt::Debug, G: FnMut() -> I::Item,
G: crate::fmt::Debug,
{ {
fn fmt(&self, f: &mut crate::fmt::Formatter<'_>) -> crate::fmt::Result { }
#[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")]
impl<I, G> fmt::Debug for IntersperseWith<I, G>
where
I: Iterator + fmt::Debug,
I::Item: fmt::Debug,
G: fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("IntersperseWith") f.debug_struct("IntersperseWith")
.field("started", &self.started)
.field("separator", &self.separator) .field("separator", &self.separator)
.field("iter", &self.iter) .field("iter", &self.iter)
.field("needs_sep", &self.needs_sep) .field("next_item", &self.next_item)
.finish() .finish()
} }
} }
#[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")] #[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")]
impl<I, G> crate::clone::Clone for IntersperseWith<I, G> impl<I, G> Clone for IntersperseWith<I, G>
where where
I: Iterator + crate::clone::Clone, I: Iterator + Clone,
I::Item: crate::clone::Clone, I::Item: Clone,
G: Clone, G: Clone,
{ {
fn clone(&self) -> Self { fn clone(&self) -> Self {
IntersperseWith { Self {
started: self.started,
separator: self.separator.clone(), separator: self.separator.clone(),
iter: self.iter.clone(), iter: self.iter.clone(),
needs_sep: self.needs_sep.clone(), next_item: self.next_item.clone(),
} }
} }
} }
@ -109,7 +146,7 @@ impl<I, G> IntersperseWith<I, G>
G: FnMut() -> I::Item, G: FnMut() -> I::Item,
{ {
pub(in crate::iter) fn new(iter: I, separator: G) -> Self { pub(in crate::iter) fn new(iter: I, separator: G) -> Self {
Self { iter: iter.peekable(), separator, needs_sep: false } Self { started: false, separator, next_item: None, iter: iter.fuse() }
} }
} }
@ -122,38 +159,52 @@ impl<I, G> Iterator for IntersperseWith<I, G>
type Item = I::Item; type Item = I::Item;
#[inline] #[inline]
fn next(&mut self) -> Option<I::Item> { fn next(&mut self) -> Option<Self::Item> {
if self.needs_sep && self.iter.peek().is_some() { if self.started {
self.needs_sep = false; if let Some(v) = self.next_item.take() {
Some(v)
} else {
let next_item = self.iter.next();
if next_item.is_some() {
self.next_item = next_item;
Some((self.separator)()) Some((self.separator)())
} else { } else {
self.needs_sep = true; None
}
}
} else {
self.started = true;
self.iter.next() self.iter.next()
} }
} }
fn size_hint(&self) -> (usize, Option<usize>) {
intersperse_size_hint(&self.iter, self.started, self.next_item.is_some())
}
fn fold<B, F>(self, init: B, f: F) -> B fn fold<B, F>(self, init: B, f: F) -> B
where where
Self: Sized, Self: Sized,
F: FnMut(B, Self::Item) -> B, F: FnMut(B, Self::Item) -> B,
{ {
intersperse_fold(self.iter, init, f, self.separator, self.needs_sep) intersperse_fold(self.iter, init, f, self.separator, self.started, self.next_item)
}
fn size_hint(&self) -> (usize, Option<usize>) {
intersperse_size_hint(&self.iter, self.needs_sep)
} }
} }
fn intersperse_size_hint<I>(iter: &I, needs_sep: bool) -> (usize, Option<usize>) fn intersperse_size_hint<I>(iter: &I, started: bool, next_is_some: bool) -> (usize, Option<usize>)
where where
I: Iterator, I: Iterator,
{ {
let (lo, hi) = iter.size_hint(); let (lo, hi) = iter.size_hint();
let next_is_elem = !needs_sep;
( (
lo.saturating_sub(next_is_elem as usize).saturating_add(lo), lo.saturating_sub(!started as usize)
hi.and_then(|hi| hi.saturating_sub(next_is_elem as usize).checked_add(hi)), .saturating_add(next_is_some as usize)
.saturating_add(lo),
hi.and_then(|hi| {
hi.saturating_sub(!started as usize)
.saturating_add(next_is_some as usize)
.checked_add(hi)
}),
) )
} }
@ -162,7 +213,8 @@ fn intersperse_fold<I, B, F, G>(
init: B, init: B,
mut f: F, mut f: F,
mut separator: G, mut separator: G,
needs_sep: bool, started: bool,
mut next_item: Option<I::Item>,
) -> B ) -> B
where where
I: Iterator, I: Iterator,
@ -171,12 +223,9 @@ fn intersperse_fold<I, B, F, G>(
{ {
let mut accum = init; let mut accum = init;
if !needs_sep { let first = if started { next_item.take() } else { iter.next() };
if let Some(x) = iter.next() { if let Some(x) = first {
accum = f(accum, x); accum = f(accum, x);
} else {
return accum;
}
} }
iter.fold(accum, |mut accum, x| { iter.fold(accum, |mut accum, x| {