Merge pull request #547 from rust-ndarray/iter-rfold

Implement Iterator::rfold where we can

Author: bluss

benches/iter.rs

@@ -10,6 +10,7 @@ use rawpointer::PointerExt;
 extern crate ndarray;
 use ndarray::prelude::*;
 use ndarray::{Zip, FoldWhile};
+use ndarray::Slice;
 
 #[bench]
 fn iter_sum_2d_regular(bench: &mut Bencher)
@@ -342,3 +343,24 @@ fn indexed_iter_3d_dyn(bench: &mut Bencher) {
         }
     })
 }
+
+#[bench]
+fn iter_sum_1d_strided_fold(bench: &mut Bencher)
+{
+    let mut a = Array::<u64, _>::ones(10240);
+    a.slice_axis_inplace(Axis(0), Slice::new(0, None, 2));
+    bench.iter(|| {
+        a.iter().fold(0, |acc, &x| acc + x)
+    });
+}
+
+#[bench]
+fn iter_sum_1d_strided_rfold(bench: &mut Bencher)
+{
+    let mut a = Array::<u64, _>::ones(10240);
+    a.slice_axis_inplace(Axis(0), Slice::new(0, None, 2));
+    bench.iter(|| {
+        a.iter().rfold(0, |acc, &x| acc + x)
+    });
+}
+

src/iterators/mod.rs

@@ -148,6 +148,23 @@ impl<A> DoubleEndedIterator for Baseiter<A, Ix1> {
 
         unsafe { Some(self.ptr.offset(offset)) }
     }
+
+    fn rfold<Acc, G>(mut self, init: Acc, mut g: G) -> Acc
+        where G: FnMut(Acc, *mut A) -> Acc,
+    {
+        let mut accum = init;
+        if let Some(index) = self.index {
+            let elem_index = index[0];
+            unsafe {
+                // self.dim[0] is the current length
+                while self.dim[0] > elem_index {
+                    self.dim[0] -= 1;
+                    accum = g(accum, self.ptr.offset(Ix1::stride_offset(&self.dim, &self.strides)));
+                }
+            }
+        }
+        accum
+    }
 }
 
 clone_bounds!(
@@ -206,6 +223,14 @@ impl<'a, A> DoubleEndedIterator for ElementsBase<'a, A, Ix1> {
     fn next_back(&mut self) -> Option<&'a A> {
         self.inner.next_back().map(|p| unsafe { &*p })
     }
+
+    fn rfold<Acc, G>(self, init: Acc, mut g: G) -> Acc
+        where G: FnMut(Acc, Self::Item) -> Acc,
+    {
+        unsafe {
+            self.inner.rfold(init, move |acc, ptr| g(acc, &*ptr))
+        }
+    }
 }
 
 impl<'a, A, D> ExactSizeIterator for ElementsBase<'a, A, D>
@@ -370,6 +395,12 @@ impl<'a, A> DoubleEndedIterator for Iter<'a, A, Ix1> {
     fn next_back(&mut self) -> Option<&'a A> {
         either_mut!(self.inner, iter => iter.next_back())
     }
+
+    fn rfold<Acc, G>(self, init: Acc, g: G) -> Acc
+        where G: FnMut(Acc, Self::Item) -> Acc
+    {
+        either!(self.inner, iter => iter.rfold(init, g))
+    }
 }
 
 impl<'a, A, D> ExactSizeIterator for Iter<'a, A, D>
@@ -431,6 +462,12 @@ impl<'a, A> DoubleEndedIterator for IterMut<'a, A, Ix1> {
     fn next_back(&mut self) -> Option<&'a mut A> {
         either_mut!(self.inner, iter => iter.next_back())
     }
+
+    fn rfold<Acc, G>(self, init: Acc, g: G) -> Acc
+        where G: FnMut(Acc, Self::Item) -> Acc
+    {
+        either!(self.inner, iter => iter.rfold(init, g))
+    }
 }
 
 impl<'a, A, D> ExactSizeIterator for IterMut<'a, A, D>
@@ -466,6 +503,14 @@ impl<'a, A> DoubleEndedIterator for ElementsBaseMut<'a, A, Ix1> {
     fn next_back(&mut self) -> Option<&'a mut A> {
         self.inner.next_back().map(|p| unsafe { &mut *p })
     }
+
+    fn rfold<Acc, G>(self, init: Acc, mut g: G) -> Acc
+        where G: FnMut(Acc, Self::Item) -> Acc
+    {
+        unsafe {
+            self.inner.rfold(init, move |acc, ptr| g(acc, &mut *ptr))
+        }
+    }
 }
 
 impl<'a, A, D> ExactSizeIterator for ElementsBaseMut<'a, A, D>

tests/iterators.rs

@@ -2,18 +2,17 @@
 extern crate ndarray;
 extern crate itertools;
 
-use ndarray::{Array0, Array2};
-use ndarray::RcArray;
+use ndarray::prelude::*;
 use ndarray::Ix;
 use ndarray::{
-    ArrayBase,
     Data,
     Dimension,
     aview1,
     arr2,
     arr3,
     Axis,
     indices,
+    Slice,
 };
 
 use itertools::assert_equal;
@@ -498,3 +497,44 @@ fn test_fold() {
     a += 1;
     assert_eq!(a.iter().fold(0, |acc, &x| acc + x), 1);
 }
+
+#[test]
+fn test_rfold() {
+    {
+        let mut a = Array1::<i32>::default(256);
+        a += 1;
+        let mut iter = a.iter();
+        iter.next();
+        assert_eq!(iter.rfold(0, |acc, &x| acc + x), a.sum() - 1);
+    }
+
+    // Test strided arrays
+    {
+        let mut a = Array1::<i32>::default(256);
+        a.slice_axis_inplace(Axis(0), Slice::new(0, None, 2));
+        a += 1;
+        let mut iter = a.iter();
+        iter.next();
+        assert_eq!(iter.rfold(0, |acc, &x| acc + x), a.sum() - 1);
+    }
+
+    {
+        let mut a = Array1::<i32>::default(256);
+        a.slice_axis_inplace(Axis(0), Slice::new(0, None, -2));
+        a += 1;
+        let mut iter = a.iter();
+        iter.next();
+        assert_eq!(iter.rfold(0, |acc, &x| acc + x), a.sum() - 1);
+    }
+
+    // Test order
+    {
+        let mut a = Array1::from_iter(0..20);
+        a.slice_axis_inplace(Axis(0), Slice::new(0, None, 2));
+        let mut iter = a.iter();
+        iter.next();
+        let output = iter.rfold(Vec::new(),
+            |mut acc, elt| { acc.push(*elt); acc });
+        assert_eq!(Array1::from_vec(output), Array::from_iter((1..10).rev().map(|i| i * 2)));
+    }
+}