Skip to content

Boost iterator intersperse(_with) performance #111379

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.

Sign up for GitHub

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

Merged
merged 5 commits into from
Jan 27, 2024
+98 −49
Merged

Boost iterator intersperse(_with) performance #111379

Changes from 3 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
f9259d1
Boost intersperse(_with) performance
nyurik May 9, 2023
b8d245e
Postpone .next() call until iteration
nyurik May 9, 2023
f1dbc7b
fmt
nyurik May 9, 2023
8cbff0b
Update library/core/src/iter/adapters/intersperse.rs
nyurik Jan 26, 2024
77f31ef
use checked_add for upper bound
nyurik Jan 26, 2024
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
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use super::Peekable;
use core::fmt;
use core::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,
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?)

I::Item: Clone,
{
}

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.map(|hi| {
hi.saturating_sub(!started as usize)
.saturating_add(next_is_some as usize)
.saturating_add(hi)
Copy link
Member

Choose a reason for hiding this comment

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

Why is this saturating instead of and_then-checked_add as before? It should return None when the upper bound is greater than usize::MAX.

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, fixed

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I wonder if there should be more tests around hints - esp when working with iters that have underlying iters... seems like this is something that can relatively easily can be messed up

Copy link
Member

Choose a reason for hiding this comment

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

Maybe so -- we'd certainly like for the standard library to get it right. But Iterator::size_hint has limited utility anyway, given that it can be implemented incorrectly by safe code for custom iterators.

Related discussion: https://internals.rust-lang.org/t/is-size-hint-1-ever-used/8187

}),
)
}

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