Skip to content

Commit 08964fd

Browse files
committed
Implement fold for IndicesIter
This improves performance somewhat.
1 parent 277e49b commit 08964fd

File tree

1 file changed

+34
-0
lines changed

1 file changed

+34
-0
lines changed

src/indexes.rs

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,25 @@ where
8080
};
8181
(l, Some(l))
8282
}
83+
84+
fn fold<B, F>(self, init: B, mut f: F) -> B
85+
where
86+
F: FnMut(B, D::Pattern) -> B,
87+
{
88+
let IndicesIter { mut index, dim } = self;
89+
let inner_axis = dim.ndim() - 1;
90+
let inner_len = dim[inner_axis];
91+
let mut acc = init;
92+
while let Some(mut ix) = index {
93+
// unroll innermost axis
94+
while ix[inner_axis] < inner_len {
95+
acc = f(acc, ix.clone().into_pattern());
96+
ix[inner_axis] += 1;
97+
}
98+
index = dim.next_for(ix);
99+
}
100+
acc
101+
}
83102
}
84103

85104
impl<D> ExactSizeIterator for IndicesIter<D> where D: Dimension {}
@@ -283,6 +302,21 @@ mod tests {
283302
assert_eq!(len, 0);
284303
}
285304

305+
#[test]
306+
fn test_indices_iter_c_fold() {
307+
let dim = (3, 4);
308+
let mut it = indices(dim).into_iter();
309+
it.next();
310+
let clone = it.clone();
311+
let len = it.len();
312+
let acc = clone.fold(0, |acc, ix| {
313+
assert_eq!(ix, it.next().unwrap());
314+
acc + 1
315+
});
316+
assert_eq!(acc, len);
317+
assert!(it.next().is_none());
318+
}
319+
286320
#[test]
287321
fn test_indices_iter_f_size_hint() {
288322
let dim = (3, 4);

0 commit comments

Comments
 (0)