From ed5015939f02ce340bf3581c866f7b6ddabb6daf Mon Sep 17 00:00:00 2001 From: Ulrik Sverdrup Date: Mon, 17 Oct 2016 10:58:21 +0200 Subject: [PATCH] Expand .zip() specialization to .map() and .cloned() Implement .zip() specialization for Map and Cloned. The crucial thing for transparent specialization is that we want to preserve the potential side effects. The simplest example is that in this code snippet: `(0..6).map(f).zip((0..4).map(g)).count()` `f` will be called five times, and `g` four times. The last time for `f` is when the other iterator is at its end, so this element is unused. This side effect can be preserved without disturbing code generation for simple uses of `.map()`. The `Zip::next_back()` case is even more complicated, unfortunately. --- src/libcore/iter/mod.rs | 59 ++++++++++++++++++ src/libcore/iter_private.rs | 5 +- src/libcore/slice.rs | 2 + src/test/codegen/zip.rs | 9 +++ src/test/run-pass/iter-zip.rs | 112 ++++++++++++++++++++++++++++++++++ 5 files changed, 186 insertions(+), 1 deletion(-) create mode 100644 src/test/run-pass/iter-zip.rs diff --git a/src/libcore/iter/mod.rs b/src/libcore/iter/mod.rs index 035418570a019..9eeb2608071c2 100644 --- a/src/libcore/iter/mod.rs +++ b/src/libcore/iter/mod.rs @@ -420,6 +420,18 @@ impl<'a, I, T: 'a> FusedIterator for Cloned where I: FusedIterator, T: Clone {} +#[doc(hidden)] +unsafe impl<'a, I, T: 'a> TrustedRandomAccess for Cloned + where I: TrustedRandomAccess, T: Clone +{ + unsafe fn get_unchecked(&mut self, i: usize) -> Self::Item { + self.it.get_unchecked(i).clone() + } + + #[inline] + fn may_have_side_effect() -> bool { true } +} + /// An iterator that repeats endlessly. /// /// This `struct` is created by the [`cycle()`] method on [`Iterator`]. See its @@ -773,6 +785,13 @@ impl ZipImpl for Zip unsafe { Some((self.a.get_unchecked(i), self.b.get_unchecked(i))) } + } else if A::may_have_side_effect() && self.index < self.a.len() { + // match the base implementation's potential side effects + unsafe { + self.a.get_unchecked(self.index); + } + self.index += 1; + None } else { None } @@ -789,6 +808,23 @@ impl ZipImpl for Zip where A: DoubleEndedIterator + ExactSizeIterator, B: DoubleEndedIterator + ExactSizeIterator { + // Adjust a, b to equal length + if A::may_have_side_effect() { + let sz = self.a.len(); + if sz > self.len { + for _ in 0..sz - cmp::max(self.len, self.index) { + self.a.next_back(); + } + } + } + if B::may_have_side_effect() { + let sz = self.b.len(); + if sz > self.len { + for _ in 0..sz - self.len { + self.b.next_back(); + } + } + } if self.index < self.len { self.len -= 1; let i = self.len; @@ -814,6 +850,9 @@ unsafe impl TrustedRandomAccess for Zip (self.a.get_unchecked(i), self.b.get_unchecked(i)) } + fn may_have_side_effect() -> bool { + A::may_have_side_effect() || B::may_have_side_effect() + } } #[unstable(feature = "fused", issue = "35602")] @@ -920,6 +959,18 @@ impl ExactSizeIterator for Map impl FusedIterator for Map where F: FnMut(I::Item) -> B {} +#[doc(hidden)] +unsafe impl TrustedRandomAccess for Map + where I: TrustedRandomAccess, + F: FnMut(I::Item) -> B, +{ + unsafe fn get_unchecked(&mut self, i: usize) -> Self::Item { + (self.f)(self.iter.get_unchecked(i)) + } + #[inline] + fn may_have_side_effect() -> bool { true } +} + /// An iterator that filters the elements of `iter` with `predicate`. /// /// This `struct` is created by the [`filter()`] method on [`Iterator`]. See its @@ -1135,6 +1186,10 @@ unsafe impl TrustedRandomAccess for Enumerate unsafe fn get_unchecked(&mut self, i: usize) -> (usize, I::Item) { (self.count + i, self.iter.get_unchecked(i)) } + + fn may_have_side_effect() -> bool { + I::may_have_side_effect() + } } #[unstable(feature = "fused", issue = "35602")] @@ -1764,6 +1819,10 @@ unsafe impl TrustedRandomAccess for Fuse unsafe fn get_unchecked(&mut self, i: usize) -> I::Item { self.iter.get_unchecked(i) } + + fn may_have_side_effect() -> bool { + I::may_have_side_effect() + } } #[unstable(feature = "fused", issue = "35602")] diff --git a/src/libcore/iter_private.rs b/src/libcore/iter_private.rs index 83eeef31ab054..bc1aaa09f3dbd 100644 --- a/src/libcore/iter_private.rs +++ b/src/libcore/iter_private.rs @@ -14,6 +14,7 @@ /// # Safety /// /// The iterator's .len() and size_hint() must be exact. +/// `.len()` must be cheap to call. /// /// .get_unchecked() must return distinct mutable references for distinct /// indices (if applicable), and must return a valid reference if index is in @@ -21,5 +22,7 @@ #[doc(hidden)] pub unsafe trait TrustedRandomAccess : ExactSizeIterator { unsafe fn get_unchecked(&mut self, i: usize) -> Self::Item; + /// Return `true` if getting an iterator element may have + /// side effects. Remember to take inner iterators into account. + fn may_have_side_effect() -> bool; } - diff --git a/src/libcore/slice.rs b/src/libcore/slice.rs index d1df56905df24..31be404ba905a 100644 --- a/src/libcore/slice.rs +++ b/src/libcore/slice.rs @@ -1968,6 +1968,7 @@ unsafe impl<'a, T> TrustedRandomAccess for Iter<'a, T> { unsafe fn get_unchecked(&mut self, i: usize) -> &'a T { &*self.ptr.offset(i as isize) } + fn may_have_side_effect() -> bool { false } } #[doc(hidden)] @@ -1975,4 +1976,5 @@ unsafe impl<'a, T> TrustedRandomAccess for IterMut<'a, T> { unsafe fn get_unchecked(&mut self, i: usize) -> &'a mut T { &mut *self.ptr.offset(i as isize) } + fn may_have_side_effect() -> bool { false } } diff --git a/src/test/codegen/zip.rs b/src/test/codegen/zip.rs index 6c956364bf80f..d0051c5165fe1 100644 --- a/src/test/codegen/zip.rs +++ b/src/test/codegen/zip.rs @@ -20,3 +20,12 @@ pub fn zip_copy(xs: &[u8], ys: &mut [u8]) { *y = *x; } } + +// CHECK-LABEL: @zip_copy_mapped +#[no_mangle] +pub fn zip_copy_mapped(xs: &[u8], ys: &mut [u8]) { +// CHECK: memcpy + for (x, y) in xs.iter().map(|&x| x).zip(ys) { + *y = x; + } +} diff --git a/src/test/run-pass/iter-zip.rs b/src/test/run-pass/iter-zip.rs new file mode 100644 index 0000000000000..b0503bc2048e7 --- /dev/null +++ b/src/test/run-pass/iter-zip.rs @@ -0,0 +1,112 @@ +// Copyright 2016 The Rust Project Developers. See the COPYRIGHT +// file at the top-level directory of this distribution and at +// http://rust-lang.org/COPYRIGHT. +// +// Licensed under the Apache License, Version 2.0 or the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + +// Test that .zip() specialization preserves side effects +// in sideeffectful iterator adaptors. + +use std::cell::Cell; + +#[derive(Debug)] +struct CountClone(Cell); + +fn count_clone() -> CountClone { CountClone(Cell::new(0)) } + +impl PartialEq for CountClone { + fn eq(&self, rhs: &i32) -> bool { + self.0.get() == *rhs + } +} + +impl Clone for CountClone { + fn clone(&self) -> Self { + let ret = CountClone(self.0.clone()); + let n = self.0.get(); + self.0.set(n + 1); + ret + } +} + +fn test_zip_cloned_sideffectful() { + let xs = [count_clone(), count_clone(), count_clone(), count_clone()]; + let ys = [count_clone(), count_clone()]; + + for _ in xs.iter().cloned().zip(ys.iter().cloned()) { } + + assert_eq!(&xs, &[1, 1, 1, 0][..]); + assert_eq!(&ys, &[1, 1][..]); + + let xs = [count_clone(), count_clone()]; + let ys = [count_clone(), count_clone(), count_clone(), count_clone()]; + + for _ in xs.iter().cloned().zip(ys.iter().cloned()) { } + + assert_eq!(&xs, &[1, 1][..]); + assert_eq!(&ys, &[1, 1, 0, 0][..]); +} + +fn test_zip_map_sideffectful() { + let mut xs = [0; 6]; + let mut ys = [0; 4]; + + for _ in xs.iter_mut().map(|x| *x += 1).zip(ys.iter_mut().map(|y| *y += 1)) { } + + assert_eq!(&xs, &[1, 1, 1, 1, 1, 0]); + assert_eq!(&ys, &[1, 1, 1, 1]); + + let mut xs = [0; 4]; + let mut ys = [0; 6]; + + for _ in xs.iter_mut().map(|x| *x += 1).zip(ys.iter_mut().map(|y| *y += 1)) { } + + assert_eq!(&xs, &[1, 1, 1, 1]); + assert_eq!(&ys, &[1, 1, 1, 1, 0, 0]); +} + +fn test_zip_map_rev_sideffectful() { + let mut xs = [0; 6]; + let mut ys = [0; 4]; + + { + let mut it = xs.iter_mut().map(|x| *x += 1).zip(ys.iter_mut().map(|y| *y += 1)); + it.next_back(); + } + assert_eq!(&xs, &[0, 0, 0, 1, 1, 1]); + assert_eq!(&ys, &[0, 0, 0, 1]); + + let mut xs = [0; 6]; + let mut ys = [0; 4]; + + { + let mut it = xs.iter_mut().map(|x| *x += 1).zip(ys.iter_mut().map(|y| *y += 1)); + (&mut it).take(5).count(); + it.next_back(); + } + assert_eq!(&xs, &[1, 1, 1, 1, 1, 1]); + assert_eq!(&ys, &[1, 1, 1, 1]); +} + +fn test_zip_nested_sideffectful() { + let mut xs = [0; 6]; + let ys = [0; 4]; + + { + // test that it has the side effect nested inside enumerate + let it = xs.iter_mut().map(|x| *x = 1).enumerate().zip(&ys); + it.count(); + } + assert_eq!(&xs, &[1, 1, 1, 1, 1, 0]); +} + +fn main() { + test_zip_cloned_sideffectful(); + test_zip_map_sideffectful(); + test_zip_map_rev_sideffectful(); + test_zip_nested_sideffectful(); +}