Fix partially-consumed AxisIter/Mut

This fixes the behavior of `.split_at()` and the `NdProducer`
implementations for partially-consumed `AxisIter` or `AxisIterMut`
instances.

Author: jturner314

src/iterators/mod.rs

@@ -738,11 +738,19 @@ where
 
 #[derive(Debug)]
 pub struct AxisIterCore<A, D> {
+    /// Index along the axis of the value of `.next()`, relative to the start
+    /// of the axis.
     index: Ix,
-    len: Ix,
+    /// (Exclusive) upper bound on `index`. Initially, this is equal to the
+    /// length of the axis.
+    end: Ix,
+    /// Stride along the axis (offset between consecutive pointers).
     stride: Ixs,
+    /// Shape of the iterator's items.
     inner_dim: D,
+    /// Strides of the iterator's items.
     inner_strides: D,
+    /// Pointer corresponding to `index == 0`.
     ptr: *mut A,
 }
 
@@ -751,7 +759,7 @@ clone_bounds!(
     AxisIterCore[A, D] {
         @copy {
             index,
-            len,
+            end,
             stride,
             ptr,
         }
@@ -767,54 +775,53 @@ impl<A, D: Dimension> AxisIterCore<A, D> {
         Di: RemoveAxis<Smaller = D>,
         S: Data<Elem = A>,
     {
-        let shape = v.shape()[axis.index()];
-        let stride = v.strides()[axis.index()];
         AxisIterCore {
             index: 0,
-            len: shape,
-            stride,
+            end: v.len_of(axis),
+            stride: v.stride_of(axis),
             inner_dim: v.dim.remove_axis(axis),
             inner_strides: v.strides.remove_axis(axis),
             ptr: v.ptr,
         }
     }
 
+    #[inline]
     unsafe fn offset(&self, index: usize) -> *mut A {
         debug_assert!(
-            index <= self.len,
-            "index={}, len={}, stride={}",
+            index <= self.end,
+            "index={}, end={}, stride={}",
             index,
-            self.len,
+            self.end,
             self.stride
         );
         self.ptr.offset(index as isize * self.stride)
     }
 
-    /// Split the iterator at index, yielding two disjoint iterators.
+    /// Splits the iterator at `index`, yielding two disjoint iterators.
+    ///
+    /// `index` is relative to the current state of the iterator (which is not
+    /// necessarily the start of the axis).
     ///
-    /// **Panics** if `index` is strictly greater than the iterator's length
+    /// **Panics** if `index` is strictly greater than the iterator's remaining
+    /// length.
     fn split_at(self, index: usize) -> (Self, Self) {
-        assert!(index <= self.len);
-        let right_ptr = if index != self.len {
-            unsafe { self.offset(index) }
-        } else {
-            self.ptr
-        };
+        assert!(index <= self.len());
+        let mid = self.index + index;
         let left = AxisIterCore {
-            index: 0,
-            len: index,
+            index: self.index,
+            end: mid,
             stride: self.stride,
             inner_dim: self.inner_dim.clone(),
             inner_strides: self.inner_strides.clone(),
             ptr: self.ptr,
         };
         let right = AxisIterCore {
-            index: 0,
-            len: self.len - index,
+            index: mid,
+            end: self.end,
             stride: self.stride,
             inner_dim: self.inner_dim,
             inner_strides: self.inner_strides,
-            ptr: right_ptr,
+            ptr: self.ptr,
         };
         (left, right)
     }
@@ -827,7 +834,7 @@ where
     type Item = *mut A;
 
     fn next(&mut self) -> Option<Self::Item> {
-        if self.index >= self.len {
+        if self.index >= self.end {
             None
         } else {
             let ptr = unsafe { self.offset(self.index) };
@@ -837,7 +844,7 @@ where
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
-        let len = self.len - self.index;
+        let len = self.len();
         (len, Some(len))
     }
 }
@@ -847,16 +854,25 @@ where
     D: Dimension,
 {
     fn next_back(&mut self) -> Option<Self::Item> {
-        if self.index >= self.len {
+        if self.index >= self.end {
             None
         } else {
-            self.len -= 1;
-            let ptr = unsafe { self.offset(self.len) };
+            self.end -= 1;
+            let ptr = unsafe { self.offset(self.end) };
             Some(ptr)
         }
     }
 }
 
+impl<A, D> ExactSizeIterator for AxisIterCore<A, D>
+where
+    D: Dimension,
+{
+    fn len(&self) -> usize {
+        self.end - self.index
+    }
+}
+
 /// An iterator that traverses over an axis and
 /// and yields each subview.
 ///
@@ -899,9 +915,13 @@ impl<'a, A, D: Dimension> AxisIter<'a, A, D> {
         }
     }
 
-    /// Split the iterator at index, yielding two disjoint iterators.
+    /// Splits the iterator at `index`, yielding two disjoint iterators.
+    ///
+    /// `index` is relative to the current state of the iterator (which is not
+    /// necessarily the start of the axis).
     ///
-    /// **Panics** if `index` is strictly greater than the iterator's length
+    /// **Panics** if `index` is strictly greater than the iterator's remaining
+    /// length.
     pub fn split_at(self, index: usize) -> (Self, Self) {
         let (left, right) = self.iter.split_at(index);
         (
@@ -946,7 +966,7 @@ where
     D: Dimension,
 {
     fn len(&self) -> usize {
-        self.size_hint().0
+        self.iter.len()
     }
 }
 
@@ -981,9 +1001,13 @@ impl<'a, A, D: Dimension> AxisIterMut<'a, A, D> {
         }
     }
 
-    /// Split the iterator at index, yielding two disjoint iterators.
+    /// Splits the iterator at `index`, yielding two disjoint iterators.
     ///
-    /// **Panics** if `index` is strictly greater than the iterator's length
+    /// `index` is relative to the current state of the iterator (which is not
+    /// necessarily the start of the axis).
+    ///
+    /// **Panics** if `index` is strictly greater than the iterator's remaining
+    /// length.
     pub fn split_at(self, index: usize) -> (Self, Self) {
         let (left, right) = self.iter.split_at(index);
         (
@@ -1028,7 +1052,7 @@ where
     D: Dimension,
 {
     fn len(&self) -> usize {
-        self.size_hint().0
+        self.iter.len()
     }
 }
 
@@ -1048,7 +1072,16 @@ impl<'a, A, D: Dimension> NdProducer for AxisIter<'a, A, D> {
     }
     #[doc(hidden)]
     fn as_ptr(&self) -> Self::Ptr {
-        self.iter.ptr
+        if self.len() > 0 {
+            // `self.iter.index` is guaranteed to be in-bounds if any of the
+            // iterator remains (i.e. if `self.len() > 0`).
+            unsafe { self.iter.offset(self.iter.index) }
+        } else {
+            // In this case, `self.iter.index` may be past the end, so we must
+            // not call `.offset()`. It's okay to return a dangling pointer
+            // because it will never be used in the length 0 case.
+            std::ptr::NonNull::dangling().as_ptr()
+        }
     }
 
     fn contiguous_stride(&self) -> isize {
@@ -1065,7 +1098,7 @@ impl<'a, A, D: Dimension> NdProducer for AxisIter<'a, A, D> {
     }
     #[doc(hidden)]
     unsafe fn uget_ptr(&self, i: &Self::Dim) -> Self::Ptr {
-        self.iter.ptr.offset(self.iter.stride * i[0] as isize)
+        self.iter.offset(self.iter.index + i[0])
     }
 
     #[doc(hidden)]
@@ -1096,7 +1129,16 @@ impl<'a, A, D: Dimension> NdProducer for AxisIterMut<'a, A, D> {
     }
     #[doc(hidden)]
     fn as_ptr(&self) -> Self::Ptr {
-        self.iter.ptr
+        if self.len() > 0 {
+            // `self.iter.index` is guaranteed to be in-bounds if any of the
+            // iterator remains (i.e. if `self.len() > 0`).
+            unsafe { self.iter.offset(self.iter.index) }
+        } else {
+            // In this case, `self.iter.index` may be past the end, so we must
+            // not call `.offset()`. It's okay to return a dangling pointer
+            // because it will never be used in the length 0 case.
+            std::ptr::NonNull::dangling().as_ptr()
+        }
     }
 
     fn contiguous_stride(&self) -> isize {
@@ -1113,7 +1155,7 @@ impl<'a, A, D: Dimension> NdProducer for AxisIterMut<'a, A, D> {
     }
     #[doc(hidden)]
     unsafe fn uget_ptr(&self, i: &Self::Dim) -> Self::Ptr {
-        self.iter.ptr.offset(self.iter.stride * i[0] as isize)
+        self.iter.offset(self.iter.index + i[0])
     }
 
     #[doc(hidden)]
@@ -1193,7 +1235,7 @@ fn chunk_iter_parts<A, D: Dimension>(
 
     let iter = AxisIterCore {
         index: 0,
-        len: iter_len,
+        end: iter_len,
         stride,
         inner_dim,
         inner_strides: v.strides,
@@ -1270,7 +1312,7 @@ macro_rules! chunk_iter_impl {
             D: Dimension,
         {
             fn next_back(&mut self) -> Option<Self::Item> {
-                let is_uneven = self.iter.len > self.n_whole_chunks;
+                let is_uneven = self.iter.end > self.n_whole_chunks;
                 let res = self.iter.next_back();
                 self.get_subview(res, is_uneven)
             }