Merge pull request #1138 from jturner314/fix-map_axis-miri

Fix Miri failure with -Zmiri-tag-raw-pointers
rust-ndarray · Mar 9, 2024 · 3a01497 · 3a01497
2 parents 0598508 + a70bf19
commit 3a01497
Showing 1 changed file with 4 additions and 20 deletions.
diff --git a/src/impl_methods.rs b/src/impl_methods.rs
@@ -2752,17 +2752,11 @@ where
         A: 'a,
         S: Data,
     {
-        let view_len = self.len_of(axis);
-        let view_stride = self.strides.axis(axis);
-        if view_len == 0 {
+        if self.len_of(axis) == 0 {
             let new_dim = self.dim.remove_axis(axis);
             Array::from_shape_simple_fn(new_dim, move || mapping(ArrayView::from(&[])))
         } else {
-            // use the 0th subview as a map to each 1d array view extended from
-            // the 0th element.
-            self.index_axis(axis, 0).map(|first_elt| unsafe {
-                mapping(ArrayView::new_(first_elt, Ix1(view_len), Ix1(view_stride)))
-            })
+            Zip::from(self.lanes(axis)).map_collect(mapping)
         }
     }
 
@@ -2783,21 +2777,11 @@ where
         A: 'a,
         S: DataMut,
     {
-        let view_len = self.len_of(axis);
-        let view_stride = self.strides.axis(axis);
-        if view_len == 0 {
+        if self.len_of(axis) == 0 {
             let new_dim = self.dim.remove_axis(axis);
             Array::from_shape_simple_fn(new_dim, move || mapping(ArrayViewMut::from(&mut [])))
         } else {
-            // use the 0th subview as a map to each 1d array view extended from
-            // the 0th element.
-            self.index_axis_mut(axis, 0).map_mut(|first_elt| unsafe {
-                mapping(ArrayViewMut::new_(
-                    first_elt,
-                    Ix1(view_len),
-                    Ix1(view_stride),
-                ))
-            })
+            Zip::from(self.lanes_mut(axis)).map_collect(mapping)
         }
     }