Skip to content

Commit ede4616

Browse files
author
sarah
committed
optimize implementation of Zip::fold and company
1 parent 75b7e52 commit ede4616

File tree

1 file changed

+226
-14
lines changed
  • library/core/src/iter/adapters

1 file changed

+226
-14
lines changed

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

Lines changed: 226 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ use crate::cmp;
22
use crate::fmt::{self, Debug};
33
use crate::iter::{DoubleEndedIterator, ExactSizeIterator, FusedIterator, Iterator};
44
use crate::iter::{InPlaceIterable, SourceIter, TrustedLen};
5+
use crate::ops::{ControlFlow, Try};
56

67
/// An iterator that iterates two other iterators simultaneously.
78
///
@@ -31,6 +32,27 @@ impl<A: Iterator, B: Iterator> Zip<A, B> {
3132
}
3233
None
3334
}
35+
#[inline]
36+
fn adjust_back(&mut self)
37+
where
38+
A: DoubleEndedIterator + ExactSizeIterator,
39+
B: DoubleEndedIterator + ExactSizeIterator,
40+
{
41+
let a_sz = self.a.len();
42+
let b_sz = self.b.len();
43+
if a_sz != b_sz {
44+
// Adjust a, b to equal length
45+
if a_sz > b_sz {
46+
for _ in 0..a_sz - b_sz {
47+
self.a.next_back();
48+
}
49+
} else {
50+
for _ in 0..b_sz - a_sz {
51+
self.b.next_back();
52+
}
53+
}
54+
}
55+
}
3456
}
3557

3658
/// Converts the arguments to iterators and zips them.
@@ -94,6 +116,23 @@ where
94116
ZipImpl::nth(self, n)
95117
}
96118

119+
#[inline]
120+
fn fold<T, F>(self, init: T, f: F) -> T
121+
where
122+
F: FnMut(T, Self::Item) -> T,
123+
{
124+
ZipImpl::fold(self, init, f)
125+
}
126+
127+
#[inline]
128+
fn try_fold<T, F, R>(&mut self, init: T, f: F) -> R
129+
where
130+
F: FnMut(T, Self::Item) -> R,
131+
R: Try<Output = T>,
132+
{
133+
ZipImpl::try_fold(self, init, f)
134+
}
135+
97136
#[inline]
98137
unsafe fn __iterator_get_unchecked(&mut self, idx: usize) -> Self::Item
99138
where
@@ -115,6 +154,23 @@ where
115154
fn next_back(&mut self) -> Option<(A::Item, B::Item)> {
116155
ZipImpl::next_back(self)
117156
}
157+
158+
#[inline]
159+
fn rfold<T, F>(self, init: T, f: F) -> T
160+
where
161+
F: FnMut(T, Self::Item) -> T,
162+
{
163+
ZipImpl::rfold(self, init, f)
164+
}
165+
166+
#[inline]
167+
fn try_rfold<T, F, R>(&mut self, init: T, f: F) -> R
168+
where
169+
F: FnMut(T, Self::Item) -> R,
170+
R: Try<Output = T>,
171+
{
172+
ZipImpl::try_rfold(self, init, f)
173+
}
118174
}
119175

120176
// Zip specialization trait
@@ -129,6 +185,25 @@ trait ZipImpl<A, B> {
129185
where
130186
A: DoubleEndedIterator + ExactSizeIterator,
131187
B: DoubleEndedIterator + ExactSizeIterator;
188+
fn fold<T, F>(self, init: T, f: F) -> T
189+
where
190+
F: FnMut(T, Self::Item) -> T;
191+
fn try_fold<T, F, R>(&mut self, init: T, f: F) -> R
192+
where
193+
F: FnMut(T, Self::Item) -> R,
194+
R: Try<Output = T>;
195+
fn rfold<T, F>(self, init: T, f: F) -> T
196+
where
197+
A: DoubleEndedIterator + ExactSizeIterator,
198+
B: DoubleEndedIterator + ExactSizeIterator,
199+
F: FnMut(T, Self::Item) -> T;
200+
fn try_rfold<T, F, R>(&mut self, init: T, f: F) -> R
201+
where
202+
A: DoubleEndedIterator + ExactSizeIterator,
203+
B: DoubleEndedIterator + ExactSizeIterator,
204+
F: FnMut(T, Self::Item) -> R,
205+
R: Try<Output = T>;
206+
132207
// This has the same safety requirements as `Iterator::__iterator_get_unchecked`
133208
unsafe fn get_unchecked(&mut self, idx: usize) -> <Self as Iterator>::Item
134209
where
@@ -171,26 +246,109 @@ macro_rules! zip_impl_general_defaults {
171246
// and doesn’t call `next_back` too often, so this implementation is safe in
172247
// the `TrustedRandomAccessNoCoerce` specialization
173248

174-
let a_sz = self.a.len();
175-
let b_sz = self.b.len();
176-
if a_sz != b_sz {
177-
// Adjust a, b to equal length
178-
if a_sz > b_sz {
179-
for _ in 0..a_sz - b_sz {
180-
self.a.next_back();
181-
}
182-
} else {
183-
for _ in 0..b_sz - a_sz {
184-
self.b.next_back();
185-
}
186-
}
187-
}
249+
self.adjust_back();
188250
match (self.a.next_back(), self.b.next_back()) {
189251
(Some(x), Some(y)) => Some((x, y)),
190252
(None, None) => None,
191253
_ => unreachable!(),
192254
}
193255
}
256+
257+
#[inline]
258+
default fn fold<T, F>(self, init: T, mut f: F) -> T
259+
where
260+
F: FnMut(T, Self::Item) -> T,
261+
{
262+
let mut a = self.a;
263+
let mut b = self.b;
264+
265+
let acc = a.try_fold(init, move |acc, x| match b.next() {
266+
Some(y) => Ok(f(acc, (x, y))),
267+
None => Err(acc),
268+
});
269+
270+
match acc {
271+
Ok(exhausted_a) => exhausted_a,
272+
Err(exhausted_b) => exhausted_b,
273+
}
274+
}
275+
276+
#[inline]
277+
default fn try_fold<T, F, R>(&mut self, init: T, mut f: F) -> R
278+
where
279+
F: FnMut(T, Self::Item) -> R,
280+
R: Try<Output = T>,
281+
{
282+
let a = &mut self.a;
283+
let b = &mut self.b;
284+
285+
let acc = a.try_fold(init, move |acc, x| match b.next() {
286+
Some(y) => {
287+
let result = f(acc, (x, y));
288+
match result.branch() {
289+
ControlFlow::Continue(continue_) => Ok(continue_),
290+
ControlFlow::Break(break_) => Err(R::from_residual(break_)),
291+
}
292+
}
293+
None => Err(R::from_output(acc)),
294+
});
295+
296+
match acc {
297+
Ok(ok) => R::from_output(ok),
298+
Err(err) => err,
299+
}
300+
}
301+
302+
#[inline]
303+
default fn rfold<T, F>(mut self, init: T, mut f: F) -> T
304+
where
305+
A: DoubleEndedIterator + ExactSizeIterator,
306+
B: DoubleEndedIterator + ExactSizeIterator,
307+
F: FnMut(T, Self::Item) -> T,
308+
{
309+
self.adjust_back();
310+
let mut a = self.a;
311+
let mut b = self.b;
312+
313+
let acc = a.try_rfold(init, move |acc, x| match b.next_back() {
314+
Some(y) => Ok(f(acc, (x, y))),
315+
None => Err(acc),
316+
});
317+
318+
match acc {
319+
Ok(exhausted_a) => exhausted_a,
320+
Err(exhausted_b) => exhausted_b,
321+
}
322+
}
323+
324+
#[inline]
325+
default fn try_rfold<T, F, R>(&mut self, init: T, mut f: F) -> R
326+
where
327+
A: DoubleEndedIterator + ExactSizeIterator,
328+
B: DoubleEndedIterator + ExactSizeIterator,
329+
F: FnMut(T, Self::Item) -> R,
330+
R: Try<Output = T>,
331+
{
332+
self.adjust_back();
333+
let a = &mut self.a;
334+
let b = &mut self.b;
335+
336+
let acc = a.try_rfold(init, move |acc, x| match b.next_back() {
337+
Some(y) => {
338+
let result = f(acc, (x, y));
339+
match result.branch() {
340+
ControlFlow::Continue(c) => ControlFlow::Continue(c),
341+
ControlFlow::Break(b) => ControlFlow::Break(R::from_residual(b)),
342+
}
343+
}
344+
None => ControlFlow::Break(R::from_output(acc)),
345+
});
346+
347+
match acc {
348+
ControlFlow::Continue(c) => R::from_output(c),
349+
ControlFlow::Break(b) => b,
350+
}
351+
}
194352
};
195353
}
196354

@@ -372,6 +530,60 @@ where
372530
None
373531
}
374532
}
533+
534+
#[inline]
535+
fn fold<T, F>(mut self, init: T, mut f: F) -> T
536+
where
537+
F: FnMut(T, Self::Item) -> T,
538+
{
539+
let mut accum = init;
540+
while let Some(x) = <Self as ZipImpl<A, B>>::next(&mut self) {
541+
accum = f(accum, x);
542+
}
543+
accum
544+
}
545+
546+
#[inline]
547+
fn try_fold<T, F, R>(&mut self, init: T, mut f: F) -> R
548+
where
549+
F: FnMut(T, Self::Item) -> R,
550+
R: Try<Output = T>,
551+
{
552+
let mut accum = init;
553+
while let Some(x) = <Self as ZipImpl<A, B>>::next(self) {
554+
accum = f(accum, x)?;
555+
}
556+
try { accum }
557+
}
558+
559+
#[inline]
560+
fn rfold<T, F>(mut self, init: T, mut f: F) -> T
561+
where
562+
A: DoubleEndedIterator + ExactSizeIterator,
563+
B: DoubleEndedIterator + ExactSizeIterator,
564+
F: FnMut(T, Self::Item) -> T,
565+
{
566+
let mut accum = init;
567+
while let Some(x) = <Self as ZipImpl<A, B>>::next_back(&mut self) {
568+
accum = f(accum, x);
569+
}
570+
accum
571+
}
572+
573+
#[inline]
574+
fn try_rfold<T, F, R>(&mut self, init: T, mut f: F) -> R
575+
where
576+
A: DoubleEndedIterator + ExactSizeIterator,
577+
B: DoubleEndedIterator + ExactSizeIterator,
578+
F: FnMut(T, Self::Item) -> R,
579+
R: Try<Output = T>,
580+
{
581+
let mut accum = init;
582+
while let Some(x) = <Self as ZipImpl<A, B>>::next_back(self) {
583+
accum = f(accum, x)?;
584+
}
585+
try { accum }
586+
}
375587
}
376588

377589
#[stable(feature = "rust1", since = "1.0.0")]

0 commit comments

Comments
 (0)