Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Boost iterator intersperse(_with) performance #111379

Merged
merged 5 commits into from
Jan 27, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
147 changes: 98 additions & 49 deletions library/core/src/iter/adapters/intersperse.rs
Original file line number Diff line number Diff line change
@@ -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.
///
Expand All @@ -10,17 +11,26 @@ pub struct Intersperse<I: Iterator>
where
I::Item: Clone,
{
started: bool,
separator: I::Item,
iter: Peekable<I>,
needs_sep: bool,
next_item: Option<I::Item>,
iter: Fuse<I>,
}

#[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")]
impl<I> FusedIterator for Intersperse<I>
where
I: FusedIterator,
I::Item: Clone,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just to note -- with iter: Fuse<I>, we shouldn't need this constraint, but it's conservative to keep it.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thx, so i left it as is (unless it gets in the way of anything?)

{
}

impl<I: Iterator> Intersperse<I>
where
I::Item: Clone,
{
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() }
}
}

Expand All @@ -33,27 +43,43 @@ where
type Item = I::Item;

#[inline]
fn next(&mut self) -> Option<I::Item> {
if self.needs_sep && self.iter.peek().is_some() {
self.needs_sep = false;
Some(self.separator.clone())
fn next(&mut self) -> Option<Self::Item> {
if self.started {
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())
} else {
None
}
}
} else {
self.needs_sep = true;
self.started = true;
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
where
Self: Sized,
F: FnMut(B, Self::Item) -> B,
{
let separator = self.separator;
intersperse_fold(self.iter, init, f, move || separator.clone(), self.needs_sep)
}

fn size_hint(&self) -> (usize, Option<usize>) {
intersperse_size_hint(&self.iter, self.needs_sep)
intersperse_fold(
self.iter,
init,
f,
move || separator.clone(),
self.started,
self.next_item,
)
}
}

Expand All @@ -66,39 +92,50 @@ pub struct IntersperseWith<I, G>
where
I: Iterator,
{
started: bool,
separator: G,
iter: Peekable<I>,
needs_sep: bool,
next_item: Option<I::Item>,
iter: Fuse<I>,
}

#[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")]
impl<I, G> FusedIterator for IntersperseWith<I, G>
where
I: FusedIterator,
G: FnMut() -> I::Item,
{
}

#[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")]
impl<I, G> crate::fmt::Debug for IntersperseWith<I, G>
impl<I, G> fmt::Debug for IntersperseWith<I, G>
where
I: Iterator + crate::fmt::Debug,
I::Item: crate::fmt::Debug,
G: crate::fmt::Debug,
I: Iterator + fmt::Debug,
I::Item: fmt::Debug,
G: fmt::Debug,
{
fn fmt(&self, f: &mut crate::fmt::Formatter<'_>) -> crate::fmt::Result {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("IntersperseWith")
.field("started", &self.started)
.field("separator", &self.separator)
.field("iter", &self.iter)
.field("needs_sep", &self.needs_sep)
.field("next_item", &self.next_item)
.finish()
}
}

#[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
I: Iterator + crate::clone::Clone,
I::Item: crate::clone::Clone,
I: Iterator + Clone,
I::Item: Clone,
G: Clone,
{
fn clone(&self) -> Self {
IntersperseWith {
Self {
started: self.started,
separator: self.separator.clone(),
iter: self.iter.clone(),
needs_sep: self.needs_sep.clone(),
next_item: self.next_item.clone(),
}
}
}
Expand All @@ -109,7 +146,7 @@ where
G: FnMut() -> I::Item,
{
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() }
}
}

Expand All @@ -122,38 +159,52 @@ where
type Item = I::Item;

#[inline]
fn next(&mut self) -> Option<I::Item> {
if self.needs_sep && self.iter.peek().is_some() {
self.needs_sep = false;
Some((self.separator)())
fn next(&mut self) -> Option<Self::Item> {
if self.started {
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)())
} else {
None
}
}
} else {
self.needs_sep = true;
self.started = true;
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
where
Self: Sized,
F: FnMut(B, Self::Item) -> B,
{
intersperse_fold(self.iter, init, f, self.separator, self.needs_sep)
}

fn size_hint(&self) -> (usize, Option<usize>) {
intersperse_size_hint(&self.iter, self.needs_sep)
intersperse_fold(self.iter, init, f, self.separator, self.started, self.next_item)
}
}

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
I: Iterator,
{
let (lo, hi) = iter.size_hint();
let next_is_elem = !needs_sep;
(
lo.saturating_sub(next_is_elem as usize).saturating_add(lo),
hi.and_then(|hi| hi.saturating_sub(next_is_elem as usize).checked_add(hi)),
lo.saturating_sub(!started as usize)
.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)
}),
)
}

Expand All @@ -162,7 +213,8 @@ fn intersperse_fold<I, B, F, G>(
init: B,
mut f: F,
mut separator: G,
needs_sep: bool,
started: bool,
mut next_item: Option<I::Item>,
) -> B
where
I: Iterator,
Expand All @@ -171,12 +223,9 @@ where
{
let mut accum = init;

if !needs_sep {
if let Some(x) = iter.next() {
accum = f(accum, x);
} else {
return accum;
}
let first = if started { next_item.take() } else { iter.next() };
if let Some(x) = first {
accum = f(accum, x);
}

iter.fold(accum, |mut accum, x| {
Expand Down
Loading