Skip to content

Commit

Permalink
Rollup merge of rust-lang#37315 - bluss:fold-more, r=alexcrichton
Browse files Browse the repository at this point in the history
Implement Iterator::fold for .chain(), .cloned(), .map() and the VecDeque iterators.

Chain can do something interesting here where it passes on the fold
into its inner iterators.

The lets the underlying iterator's custom fold() be used, and skips the
regular chain logic in next.

Also implement .fold() specifically for .map() and .cloned() so that any
inner fold improvements are available through map and cloned.

The same way, a VecDeque iterator fold can be turned into two slice folds.

These changes lend the power of the slice iterator's loop codegen to
VecDeque, and to chains and flat maps of slice iterators, and so on.
It's an improvement for .sum() and .product(), and other uses of fold.
  • Loading branch information
Manishearth committed Oct 26, 2016
2 parents 1802610 + a16626f commit e2bee0f
Show file tree
Hide file tree
Showing 3 changed files with 98 additions and 20 deletions.
74 changes: 54 additions & 20 deletions src/libcollections/vec_deque.rs
Original file line number Diff line number Diff line change
Expand Up @@ -743,16 +743,8 @@ impl<T> VecDeque<T> {
#[stable(feature = "deque_extras_15", since = "1.5.0")]
pub fn as_slices(&self) -> (&[T], &[T]) {
unsafe {
let contiguous = self.is_contiguous();
let buf = self.buffer_as_slice();
if contiguous {
let (empty, buf) = buf.split_at(0);
(&buf[self.tail..self.head], empty)
} else {
let (mid, right) = buf.split_at(self.tail);
let (left, _) = mid.split_at(self.head);
(right, left)
}
RingSlices::ring_slices(buf, self.head, self.tail)
}
}

Expand Down Expand Up @@ -780,20 +772,10 @@ impl<T> VecDeque<T> {
#[stable(feature = "deque_extras_15", since = "1.5.0")]
pub fn as_mut_slices(&mut self) -> (&mut [T], &mut [T]) {
unsafe {
let contiguous = self.is_contiguous();
let head = self.head;
let tail = self.tail;
let buf = self.buffer_as_mut_slice();

if contiguous {
let (empty, buf) = buf.split_at_mut(0);
(&mut buf[tail..head], empty)
} else {
let (mid, right) = buf.split_at_mut(tail);
let (left, _) = mid.split_at_mut(head);

(right, left)
}
RingSlices::ring_slices(buf, head, tail)
}
}

Expand Down Expand Up @@ -1829,6 +1811,42 @@ fn wrap_index(index: usize, size: usize) -> usize {
index & (size - 1)
}

/// Returns the two slices that cover the VecDeque's valid range
trait RingSlices : Sized {
fn slice(self, from: usize, to: usize) -> Self;
fn split_at(self, i: usize) -> (Self, Self);

fn ring_slices(buf: Self, head: usize, tail: usize) -> (Self, Self) {
let contiguous = tail <= head;
if contiguous {
let (empty, buf) = buf.split_at(0);
(buf.slice(tail, head), empty)
} else {
let (mid, right) = buf.split_at(tail);
let (left, _) = mid.split_at(head);
(right, left)
}
}
}

impl<'a, T> RingSlices for &'a [T] {
fn slice(self, from: usize, to: usize) -> Self {
&self[from..to]
}
fn split_at(self, i: usize) -> (Self, Self) {
(*self).split_at(i)
}
}

impl<'a, T> RingSlices for &'a mut [T] {
fn slice(self, from: usize, to: usize) -> Self {
&mut self[from..to]
}
fn split_at(self, i: usize) -> (Self, Self) {
(*self).split_at_mut(i)
}
}

/// Calculate the number of elements left to be read in the buffer
#[inline]
fn count(tail: usize, head: usize, size: usize) -> usize {
Expand Down Expand Up @@ -1875,6 +1893,14 @@ impl<'a, T> Iterator for Iter<'a, T> {
let len = count(self.tail, self.head, self.ring.len());
(len, Some(len))
}

fn fold<Acc, F>(self, mut accum: Acc, mut f: F) -> Acc
where F: FnMut(Acc, Self::Item) -> Acc,
{
let (front, back) = RingSlices::ring_slices(self.ring, self.head, self.tail);
accum = front.iter().fold(accum, &mut f);
back.iter().fold(accum, &mut f)
}
}

#[stable(feature = "rust1", since = "1.0.0")]
Expand Down Expand Up @@ -1927,6 +1953,14 @@ impl<'a, T> Iterator for IterMut<'a, T> {
let len = count(self.tail, self.head, self.ring.len());
(len, Some(len))
}

fn fold<Acc, F>(self, mut accum: Acc, mut f: F) -> Acc
where F: FnMut(Acc, Self::Item) -> Acc,
{
let (front, back) = RingSlices::ring_slices(self.ring, self.head, self.tail);
accum = front.iter_mut().fold(accum, &mut f);
back.iter_mut().fold(accum, &mut f)
}
}

#[stable(feature = "rust1", since = "1.0.0")]
Expand Down
32 changes: 32 additions & 0 deletions src/libcore/iter/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -399,6 +399,12 @@ impl<'a, I, T: 'a> Iterator for Cloned<I>
fn size_hint(&self) -> (usize, Option<usize>) {
self.it.size_hint()
}

fn fold<Acc, F>(self, init: Acc, mut f: F) -> Acc
where F: FnMut(Acc, Self::Item) -> Acc,
{
self.it.fold(init, move |acc, elt| f(acc, elt.clone()))
}
}

#[stable(feature = "iter_cloned", since = "1.1.0")]
Expand Down Expand Up @@ -544,6 +550,25 @@ impl<A, B> Iterator for Chain<A, B> where
}
}

fn fold<Acc, F>(self, init: Acc, mut f: F) -> Acc
where F: FnMut(Acc, Self::Item) -> Acc,
{
let mut accum = init;
match self.state {
ChainState::Both | ChainState::Front => {
accum = self.a.fold(accum, &mut f);
}
_ => { }
}
match self.state {
ChainState::Both | ChainState::Back => {
accum = self.b.fold(accum, &mut f);
}
_ => { }
}
accum
}

#[inline]
fn nth(&mut self, mut n: usize) -> Option<A::Item> {
match self.state {
Expand Down Expand Up @@ -939,6 +964,13 @@ impl<B, I: Iterator, F> Iterator for Map<I, F> where F: FnMut(I::Item) -> B {
fn size_hint(&self) -> (usize, Option<usize>) {
self.iter.size_hint()
}

fn fold<Acc, G>(self, init: Acc, mut g: G) -> Acc
where G: FnMut(Acc, Self::Item) -> Acc,
{
let mut f = self.f;
self.iter.fold(init, move |acc, elt| g(acc, f(elt)))
}
}

#[stable(feature = "rust1", since = "1.0.0")]
Expand Down
12 changes: 12 additions & 0 deletions src/libcoretest/iter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -985,6 +985,18 @@ fn test_empty() {
assert_eq!(it.next(), None);
}

#[test]
fn test_chain_fold() {
let xs = [1, 2, 3];
let ys = [1, 2, 0];

let mut iter = xs.iter().chain(&ys);
iter.next();
let mut result = Vec::new();
iter.fold((), |(), &elt| result.push(elt));
assert_eq!(&[2, 3, 1, 2, 0], &result[..]);
}

#[bench]
fn bench_rposition(b: &mut Bencher) {
let it: Vec<usize> = (0..300).collect();
Expand Down

0 comments on commit e2bee0f

Please sign in to comment.