Skip to content

Commit

Permalink
Fix zero handling in AxisChunksIter/Mut
Browse files Browse the repository at this point in the history
Now, chunk size of zero and axis length of zero are handled correctly.
  • Loading branch information
jturner314 committed Jul 22, 2019
1 parent a0130ad commit 4f857a3
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 4 deletions.
4 changes: 2 additions & 2 deletions src/impl_methods.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1005,7 +1005,7 @@ where
/// The last view may have less elements if `size` does not divide
/// the axis' dimension.
///
/// **Panics** if `axis` is out of bounds.
/// **Panics** if `axis` is out of bounds or if `size` is zero.
///
/// ```
/// use ndarray::Array;
Expand Down Expand Up @@ -1036,7 +1036,7 @@ where
///
/// Iterator element is `ArrayViewMut<A, D>`
///
/// **Panics** if `axis` is out of bounds.
/// **Panics** if `axis` is out of bounds or if `size` is zero.
pub fn axis_chunks_iter_mut(&mut self, axis: Axis, size: usize) -> AxisChunksIterMut<A, D>
where
S: DataMut,
Expand Down
11 changes: 9 additions & 2 deletions src/iterators/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1206,21 +1206,28 @@ clone_bounds!(
///
/// Returns an axis iterator with the correct stride to move between chunks,
/// the number of chunks, and the shape of the last chunk.
///
/// **Panics** if `size == 0`.
fn chunk_iter_parts<A, D: Dimension>(
v: ArrayView<A, D>,
axis: Axis,
size: usize,
) -> (AxisIterCore<A, D>, usize, D) {
assert_ne!(size, 0, "Chunk size must be nonzero.");
let axis_len = v.len_of(axis);
let size = if size > axis_len { axis_len } else { size };
let n_whole_chunks = axis_len / size;
let chunk_remainder = axis_len % size;
let iter_len = if chunk_remainder == 0 {
n_whole_chunks
} else {
n_whole_chunks + 1
};
let stride = v.stride_of(axis) * size as isize;
let stride = if n_whole_chunks == 0 {
// This case avoids potential overflow when `size > axis_len`.
0
} else {
v.stride_of(axis) * size as isize
};

let axis = axis.index();
let mut inner_dim = v.dim.clone();
Expand Down

0 comments on commit 4f857a3

Please sign in to comment.