Fix bugs and improve test for IndicesIter::fold
@@ -86,14 +86,21 @@ where
F: FnMut(B, D::Pattern) -> B,
{
let IndicesIter { mut index, dim } = self;
- let inner_axis = dim.ndim() - 1;
+ let ndim = dim.ndim();
+ if ndim == 0 {
+ return match index {
+ Some(ix) => f(init, ix.into_pattern()),
+ None => init,
+ };
+ }
+ let inner_axis = ndim - 1;
let inner_len = dim[inner_axis];
let mut acc = init;
while let Some(mut ix) = index {
// unroll innermost axis
- while ix[inner_axis] < inner_len {
+ for i in ix[inner_axis]..inner_len {
+ ix[inner_axis] = i;
acc = f(acc, ix.clone().into_pattern());
- ix[inner_axis] += 1;
}
index = dim.next_for(ix);
}
@@ -304,9 +311,13 @@ mod tests {
#[test]
fn test_indices_iter_c_fold() {
- let dim = (3, 4);
- let mut it = indices(dim).into_iter();
+ macro_rules! run_test {
+ ($dim:expr) => {
+ for num_consume in 0..3 {
+ let mut it = indices($dim).into_iter();
+ for _ in 0..num_consume {
it.next();
+ }
let clone = it.clone();
let len = it.len();
let acc = clone.fold(0, |acc, ix| {
@@ -316,6 +327,15 @@ mod tests {
assert_eq!(acc, len);
assert!(it.next().is_none());
}
+ };
+ }
+ run_test!(());
+ run_test!((2,));
+ run_test!((2, 3));
+ run_test!((2, 0, 3));
+ run_test!((2, 3, 4));
+ run_test!((2, 3, 4, 2));
+ }
#[test]
fn test_indices_iter_f_size_hint() {
