Skip to content

Commit e6d6e5c

Browse files
committed
Reimplement TrustedRandomAccess specializations for Zip::fold and Zip::try_fold
Previously these optimizations automatically fell out of the default implementation for `fold` and `try_fold` because `next()` was specialized. Since that is gone we now need to implement it explicitly.
1 parent a56524a commit e6d6e5c

File tree

2 files changed

+98
-13
lines changed
  • library/core

2 files changed

+98
-13
lines changed

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

Lines changed: 98 additions & 0 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::Try;
56

67
/// An iterator that iterates two other iterators simultaneously.
78
///
@@ -95,6 +96,14 @@ where
9596
(lower, upper)
9697
}
9798

99+
fn fold<T, F>(self, init: T, f: F) -> T
100+
where
101+
Self: Sized,
102+
F: FnMut(T, Self::Item) -> T,
103+
{
104+
self.spec_fold(init, f)
105+
}
106+
98107
#[inline]
99108
#[doc(hidden)]
100109
unsafe fn __iterator_get_unchecked(&mut self, idx: usize) -> Self::Item
@@ -137,6 +146,95 @@ where
137146
}
138147
}
139148

149+
trait ZipSpec: Iterator {
150+
fn spec_fold<T, F>(self, init: T, f: F) -> T
151+
where
152+
Self: Sized,
153+
F: FnMut(T, Self::Item) -> T;
154+
155+
fn spec_try_fold<B, F, R>(&mut self, init: B, f: F) -> R
156+
where
157+
Self: Sized,
158+
F: FnMut(B, Self::Item) -> R,
159+
R: Try<Output = B>;
160+
161+
}
162+
163+
impl<A, B> ZipSpec for Zip<A, B>
164+
where
165+
A: Iterator,
166+
B: Iterator,
167+
{
168+
default fn spec_fold<T, F>(mut self, init: T, mut f: F) -> T
169+
where
170+
Self: Sized,
171+
F: FnMut(T, Self::Item) -> T,
172+
{
173+
let mut accum = init;
174+
while let Some(x) = self.next() {
175+
accum = f(accum, x);
176+
}
177+
accum
178+
}
179+
180+
default fn spec_try_fold<T, F, R>(&mut self, init: T, mut f: F) -> R
181+
where
182+
Self: Sized,
183+
F: FnMut(T, Self::Item) -> R,
184+
R: Try<Output = T>,
185+
{
186+
let mut accum = init;
187+
while let Some(x) = self.next() {
188+
accum = f(accum, x)?;
189+
}
190+
try { accum }
191+
}
192+
}
193+
194+
impl<A, B> ZipSpec for Zip<A, B>
195+
where
196+
A: Iterator,
197+
B: Iterator,
198+
Self: TrustedRandomAccess,
199+
{
200+
fn spec_fold<T, F>(mut self, init: T, mut f: F) -> T
201+
where
202+
Self: Sized,
203+
F: FnMut(T, Self::Item) -> T,
204+
{
205+
let _ = self.advance_by(0);
206+
let len = self.size();
207+
let mut accum = init;
208+
for i in 0..len {
209+
// SAFETY: each item is only accessed once and we run the cleanup function afterwards
210+
let x = unsafe { self.__iterator_get_unchecked(i) };
211+
accum = f(accum, x);
212+
}
213+
// FIXME drop-guard or use ForLoopDesugar
214+
self.cleanup(len, true);
215+
accum
216+
}
217+
218+
fn spec_try_fold<T, F, R>(&mut self, init: T, mut f: F) -> R
219+
where
220+
Self: Sized,
221+
F: FnMut(T, Self::Item) -> R,
222+
R: Try<Output = T>,
223+
{
224+
let _ = self.advance_by(0);
225+
let len = self.size();
226+
let mut accum = init;
227+
for i in 0..len {
228+
// SAFETY: each item is only accessed once and we run the cleanup function afterwards
229+
let x = unsafe { self.__iterator_get_unchecked(i) };
230+
accum = f(accum, x)?;
231+
}
232+
// FIXME drop-guard or use ForLoopDesugar
233+
self.cleanup(len, true);
234+
try { accum }
235+
}
236+
}
237+
140238
#[stable(feature = "rust1", since = "1.0.0")]
141239
impl<A, B> ExactSizeIterator for Zip<A, B>
142240
where

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

Lines changed: 0 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -174,19 +174,6 @@ fn test_zip_map_rev_sideffectful() {
174174
assert_eq!(&ys, &[1, 1, 1, 1]);
175175
}
176176

177-
#[test]
178-
fn test_zip_nested_sideffectful() {
179-
let mut xs = [0; 6];
180-
let ys = [0; 4];
181-
182-
{
183-
// test that it has the side effect nested inside enumerate
184-
let it = xs.iter_mut().map(|x| *x = 1).enumerate().zip(&ys);
185-
it.count();
186-
}
187-
assert_eq!(&xs, &[1, 1, 1, 1, 1, 0]);
188-
}
189-
190177
#[test]
191178
fn test_zip_nth_back_side_effects_exhausted() {
192179
let mut a = Vec::new();

0 commit comments

Comments
 (0)