Skip to content

Commit f9259d1

Browse files
committed
Boost intersperse(_with) performance
I did some benchmark digging into the `intersperse` and `intersperse_with` code as part of the https://internals.rust-lang.org/t/add-iterate-with-separators-iterator-function/18781/13 discussion, and as a result I optimized them a bit, without relying on the peekable iterator.
1 parent 5bd5d21 commit f9259d1

File tree

1 file changed

+74
-52
lines changed

1 file changed

+74
-52
lines changed

library/core/src/iter/adapters/intersperse.rs

+74-52
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
1-
use super::Peekable;
1+
use core::fmt;
2+
use core::iter::{Fuse, FusedIterator};
23

34
/// An iterator adapter that places a separator between all elements.
45
///
@@ -11,16 +12,25 @@ where
1112
I::Item: Clone,
1213
{
1314
separator: I::Item,
14-
iter: Peekable<I>,
15-
needs_sep: bool,
15+
next_item: Option<I::Item>,
16+
iter: Fuse<I>,
17+
}
18+
19+
#[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")]
20+
impl<I> FusedIterator for Intersperse<I>
21+
where
22+
I: FusedIterator,
23+
I::Item: Clone,
24+
{
1625
}
1726

1827
impl<I: Iterator> Intersperse<I>
1928
where
2029
I::Item: Clone,
2130
{
2231
pub(in crate::iter) fn new(iter: I, separator: I::Item) -> Self {
23-
Self { iter: iter.peekable(), separator, needs_sep: false }
32+
let mut iter = iter.fuse();
33+
Self { separator, next_item: iter.next(), iter }
2434
}
2535
}
2636

@@ -33,27 +43,31 @@ where
3343
type Item = I::Item;
3444

3545
#[inline]
36-
fn next(&mut self) -> Option<I::Item> {
37-
if self.needs_sep && self.iter.peek().is_some() {
38-
self.needs_sep = false;
39-
Some(self.separator.clone())
46+
fn next(&mut self) -> Option<Self::Item> {
47+
if let Some(v) = self.next_item.take() {
48+
Some(v)
4049
} else {
41-
self.needs_sep = true;
42-
self.iter.next()
50+
let next_item = self.iter.next();
51+
if next_item.is_some() {
52+
self.next_item = next_item;
53+
Some(self.separator.clone())
54+
} else {
55+
None
56+
}
4357
}
4458
}
4559

60+
fn size_hint(&self) -> (usize, Option<usize>) {
61+
intersperse_size_hint(&self.iter, self.next_item.is_some())
62+
}
63+
4664
fn fold<B, F>(self, init: B, f: F) -> B
4765
where
4866
Self: Sized,
4967
F: FnMut(B, Self::Item) -> B,
5068
{
5169
let separator = self.separator;
52-
intersperse_fold(self.iter, init, f, move || separator.clone(), self.needs_sep)
53-
}
54-
55-
fn size_hint(&self) -> (usize, Option<usize>) {
56-
intersperse_size_hint(&self.iter, self.needs_sep)
70+
intersperse_fold(self.iter, init, f, move || separator.clone(), self.next_item)
5771
}
5872
}
5973

@@ -67,38 +81,46 @@ where
6781
I: Iterator,
6882
{
6983
separator: G,
70-
iter: Peekable<I>,
71-
needs_sep: bool,
84+
next_item: Option<I::Item>,
85+
iter: Fuse<I>,
86+
}
87+
88+
#[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")]
89+
impl<I, G> FusedIterator for IntersperseWith<I, G>
90+
where
91+
I: FusedIterator,
92+
G: FnMut() -> I::Item,
93+
{
7294
}
7395

7496
#[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")]
75-
impl<I, G> crate::fmt::Debug for IntersperseWith<I, G>
97+
impl<I, G> fmt::Debug for IntersperseWith<I, G>
7698
where
77-
I: Iterator + crate::fmt::Debug,
78-
I::Item: crate::fmt::Debug,
79-
G: crate::fmt::Debug,
99+
I: Iterator + fmt::Debug,
100+
I::Item: fmt::Debug,
101+
G: fmt::Debug,
80102
{
81-
fn fmt(&self, f: &mut crate::fmt::Formatter<'_>) -> crate::fmt::Result {
103+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
82104
f.debug_struct("IntersperseWith")
83105
.field("separator", &self.separator)
84106
.field("iter", &self.iter)
85-
.field("needs_sep", &self.needs_sep)
107+
.field("next_item", &self.next_item)
86108
.finish()
87109
}
88110
}
89111

90112
#[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")]
91-
impl<I, G> crate::clone::Clone for IntersperseWith<I, G>
113+
impl<I, G> Clone for IntersperseWith<I, G>
92114
where
93-
I: Iterator + crate::clone::Clone,
94-
I::Item: crate::clone::Clone,
115+
I: Iterator + Clone,
116+
I::Item: Clone,
95117
G: Clone,
96118
{
97119
fn clone(&self) -> Self {
98-
IntersperseWith {
120+
Self {
99121
separator: self.separator.clone(),
100122
iter: self.iter.clone(),
101-
needs_sep: self.needs_sep.clone(),
123+
next_item: self.next_item.clone(),
102124
}
103125
}
104126
}
@@ -109,7 +131,8 @@ where
109131
G: FnMut() -> I::Item,
110132
{
111133
pub(in crate::iter) fn new(iter: I, separator: G) -> Self {
112-
Self { iter: iter.peekable(), separator, needs_sep: false }
134+
let mut iter = iter.fuse();
135+
Self { separator, next_item: iter.next(), iter }
113136
}
114137
}
115138

@@ -122,47 +145,50 @@ where
122145
type Item = I::Item;
123146

124147
#[inline]
125-
fn next(&mut self) -> Option<I::Item> {
126-
if self.needs_sep && self.iter.peek().is_some() {
127-
self.needs_sep = false;
128-
Some((self.separator)())
148+
fn next(&mut self) -> Option<Self::Item> {
149+
if let Some(v) = self.next_item.take() {
150+
Some(v)
129151
} else {
130-
self.needs_sep = true;
131-
self.iter.next()
152+
let next_item = self.iter.next();
153+
if next_item.is_some() {
154+
self.next_item = next_item;
155+
Some((self.separator)())
156+
} else {
157+
None
158+
}
132159
}
133160
}
134161

162+
fn size_hint(&self) -> (usize, Option<usize>) {
163+
intersperse_size_hint(&self.iter, self.next_item.is_some())
164+
}
165+
135166
fn fold<B, F>(self, init: B, f: F) -> B
136167
where
137168
Self: Sized,
138169
F: FnMut(B, Self::Item) -> B,
139170
{
140-
intersperse_fold(self.iter, init, f, self.separator, self.needs_sep)
141-
}
142-
143-
fn size_hint(&self) -> (usize, Option<usize>) {
144-
intersperse_size_hint(&self.iter, self.needs_sep)
171+
intersperse_fold(self.iter, init, f, self.separator, self.next_item)
145172
}
146173
}
147174

148-
fn intersperse_size_hint<I>(iter: &I, needs_sep: bool) -> (usize, Option<usize>)
175+
fn intersperse_size_hint<I>(iter: &I, next_is_elem: bool) -> (usize, Option<usize>)
149176
where
150177
I: Iterator,
151178
{
152179
let (lo, hi) = iter.size_hint();
153-
let next_is_elem = !needs_sep;
154180
(
155-
lo.saturating_sub(next_is_elem as usize).saturating_add(lo),
156-
hi.and_then(|hi| hi.saturating_sub(next_is_elem as usize).checked_add(hi)),
181+
lo.saturating_add(next_is_elem as usize).saturating_add(lo),
182+
hi.map(|hi| hi.saturating_add(next_is_elem as usize).saturating_add(hi)),
157183
)
158184
}
159185

160186
fn intersperse_fold<I, B, F, G>(
161-
mut iter: I,
187+
iter: I,
162188
init: B,
163189
mut f: F,
164190
mut separator: G,
165-
needs_sep: bool,
191+
mut next_item: Option<I::Item>,
166192
) -> B
167193
where
168194
I: Iterator,
@@ -171,12 +197,8 @@ where
171197
{
172198
let mut accum = init;
173199

174-
if !needs_sep {
175-
if let Some(x) = iter.next() {
176-
accum = f(accum, x);
177-
} else {
178-
return accum;
179-
}
200+
if let Some(x) = next_item.take() {
201+
accum = f(accum, x);
180202
}
181203

182204
iter.fold(accum, |mut accum, x| {

0 commit comments

Comments
 (0)