Repository: bluss/rust-ndarray
Pull request #311 (2017-04-13 21:03:07)

Fix two bugs with `Array::from_shape_fn`
- The indices did not map correctly to element position for f-order shapes
- from_shape_fn did not correctly check the size for overflow


Patch:
diff --git a/src/dimension/dimension_trait.rs b/src/dimension/dimension_trait.rs
index 1aa310c7..92f0e318 100644
--- a/src/dimension/dimension_trait.rs
+++ b/src/dimension/dimension_trait.rs
@@ -191,6 +191,27 @@ pub trait Dimension : Clone + Eq + Debug + Send + Sync + Default +
         }
     }
 
+    #[doc(hidden)]
+    /// Iteration -- Use self as size, and create the next index after `index`
+    /// Return false if iteration is done
+    ///
+    /// Next in f-order
+    #[inline]
+    fn next_for_f(&self, index: &mut Self) -> bool {
+        let mut end_iteration = true;
+        for (&dim, ix) in zip(self.slice(), index.slice_mut()) {
+            *ix += 1;
+            if *ix == dim {
+                *ix = 0;
+            } else {
+                end_iteration = false;
+                break;
+            }
+        }
+        !end_iteration
+    }
+
+
     #[doc(hidden)]
     /// Return stride offset for index.
     fn stride_offset(index: &Self, strides: &Self) -> isize {
diff --git a/src/impl_constructors.rs b/src/impl_constructors.rs
index a4a0c13b..ec5eaa57 100644
--- a/src/impl_constructors.rs
+++ b/src/impl_constructors.rs
@@ -18,6 +18,7 @@ use dimension;
 use linspace;
 use error::{self, ShapeError, ErrorKind};
 use indices;
+use indexes;
 use iterators::{to_vec, to_vec_mapped};
 
 /// # Constructor Methods for Owned Arrays
@@ -192,7 +193,8 @@ impl<S, A, D> ArrayBase<S, D>
 
     /// Create an array with values created by the function `f`.
     ///
-    /// The elements are visited in arbitirary order.
+    /// `f` is called with the index of the element to create; the elements are
+    /// visited in arbitirary order.
     ///
     /// **Panics** if the number of elements in `shape` would overflow usize.
     pub fn from_shape_fn<Sh, F>(shape: Sh, f: F) -> Self
@@ -200,8 +202,15 @@ impl<S, A, D> ArrayBase<S, D>
               F: FnMut(D::Pattern) -> A,
     {
         let shape = shape.into_shape();
-        let v = to_vec_mapped(indices(shape.dim.clone()).into_iter(), f);
-        unsafe { Self::from_shape_vec_unchecked(shape, v) }
+        let _ = size_checked_unwrap!(shape.dim);
+        if shape.is_c {
+            let v = to_vec_mapped(indices(shape.dim.clone()).into_iter(), f);
+            unsafe { Self::from_shape_vec_unchecked(shape, v) }
+        } else {
+            let dim = shape.dim.clone();
+            let v = to_vec_mapped(indexes::indices_iter_f(dim).into_iter(), f);
+            unsafe { Self::from_shape_vec_unchecked(shape, v) }
+        }
     }
 
     /// Create an array with the given shape from a vector. (No cloning of
diff --git a/src/indexes.rs b/src/indexes.rs
index b03e8a90..f2861b48 100644
--- a/src/indexes.rs
+++ b/src/indexes.rs
@@ -196,3 +196,59 @@ impl<D: Dimension + Copy> NdProducer for Indices<D> {
     }
 }
 
+/// An iterator over the indexes of an array shape.
+///
+/// Iterator element type is `D`.
+#[derive(Clone)]
+pub struct IndicesIterF<D> {
+    dim: D,
+    index: D,
+    has_remaining: bool,
+}
+
+pub fn indices_iter_f<E>(shape: E) -> IndicesIterF<E::Dim>
+    where E: IntoDimension,
+{
+    let dim = shape.into_dimension();
+    let zero = dim.zero_index();
+    IndicesIterF {
+        has_remaining: zero != dim,
+        index: zero,
+        dim: dim,
+    }
+}
+
+impl<D> Iterator for IndicesIterF<D>
+    where D: Dimension,
+{
+    type Item = D::Pattern;
+    #[inline]
+    fn next(&mut self) -> Option<Self::Item> {
+        if !self.has_remaining {
+            None
+        } else {
+            let elt = self.index.clone().into_pattern();
+            self.has_remaining = self.dim.next_for_f(&mut self.index);
+            Some(elt)
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let l = match self.index {
+            ref ix => {
+                let gone = self.dim
+                               .default_strides()
+                               .slice()
+                               .iter()
+                               .zip(ix.slice().iter())
+                               .fold(0, |s, (&a, &b)| s + a as usize * b as usize);
+                self.dim.size() - gone
+            }
+        };
+        (l, Some(l))
+    }
+}
+
+impl<D> ExactSizeIterator for IndicesIterF<D>
+    where D: Dimension
+{}
diff --git a/src/iterators/mod.rs b/src/iterators/mod.rs
index 3a407a3b..cd00bfda 100644
--- a/src/iterators/mod.rs
+++ b/src/iterators/mod.rs
@@ -1178,6 +1178,7 @@ pub unsafe trait TrustedIterator { }
 use std;
 use linspace::Linspace;
 use iter::IndicesIter;
+use indexes::IndicesIterF;
 
 unsafe impl<F> TrustedIterator for Linspace<F> { }
 unsafe impl<'a, A, D> TrustedIterator for Iter<'a, A, D> { }
@@ -1185,7 +1186,9 @@ unsafe impl<I, F> TrustedIterator for std::iter::Map<I, F>
     where I: TrustedIterator { }
 unsafe impl<'a, A> TrustedIterator for slice::Iter<'a, A> { }
 unsafe impl TrustedIterator for ::std::ops::Range<usize> { }
+// FIXME: These indices iter are dubious -- size needs to be checked up front.
 unsafe impl<D> TrustedIterator for IndicesIter<D> where D: Dimension { }
+unsafe impl<D> TrustedIterator for IndicesIterF<D> where D: Dimension { }
 
 
 /// Like Iterator::collect, but only for trusted length iterators


Test patch:
diff --git a/tests/array-construct.rs b/tests/array-construct.rs
index 11637fa6..ef690d40 100644
--- a/tests/array-construct.rs
+++ b/tests/array-construct.rs
@@ -43,6 +43,54 @@ fn test_uninit() {
     }
 }
 
+#[test]
+fn test_from_fn_c1() {
+    let a = Array::from_shape_fn(28, |i| i);
+    for (i, elt) in a.indexed_iter() {
+        assert_eq!(i, *elt);
+    }
+}
+
+#[test]
+fn test_from_fn_c() {
+    let a = Array::from_shape_fn((4, 7), |i| i);
+    for (i, elt) in a.indexed_iter() {
+        assert_eq!(i, *elt);
+    }
+}
+
+#[test]
+fn test_from_fn_c3() {
+    let a = Array::from_shape_fn((4, 3, 7), |i| i);
+    for (i, elt) in a.indexed_iter() {
+        assert_eq!(i, *elt);
+    }
+}
+
+#[test]
+fn test_from_fn_f1() {
+    let a = Array::from_shape_fn(28.f(), |i| i);
+    for (i, elt) in a.indexed_iter() {
+        assert_eq!(i, *elt);
+    }
+}
+
+#[test]
+fn test_from_fn_f() {
+    let a = Array::from_shape_fn((4, 7).f(), |i| i);
+    for (i, elt) in a.indexed_iter() {
+        assert_eq!(i, *elt);
+    }
+}
+
+#[test]
+fn test_from_fn_f3() {
+    let a = Array::from_shape_fn((4, 2, 7).f(), |i| i);
+    for (i, elt) in a.indexed_iter() {
+        assert_eq!(i, *elt);
+    }
+}
+
 #[test]
 fn deny_wraparound_from_vec() {
     let five = vec![0; 5];
@@ -73,6 +121,12 @@ fn deny_wraparound_default() {
     let _five_large = Array::<f32, _>::default((3, 7, 29, 36760123, 823996703));
 }
 
+#[should_panic]
+#[test]
+fn deny_wraparound_from_shape_fn() {
+    let _five_large = Array::<f32, _>::from_shape_fn((3, 7, 29, 36760123, 823996703), |_| 0.);
+}
+
 #[should_panic]
 #[test]
 fn deny_wraparound_uninit() {