Iterate mutably over two lanes of an array #775

Activity

1. jturner314 commented on Jan 21, 2020

   

   jturner314

   on Jan 21, 2020

   Member

   More actions

   I haven't thought about using axis_iter_mut for this use-case before; that's clever! :)

   With the current version of ndarray, you can use unsafe code like this:

   use ndarray::prelude::*;
   
   fn index_two_lanes_mut<T>(
       mut arr: ArrayViewMut2<'_, T>,
       axis: Axis,
       l1: usize,
       l2: usize,
   ) -> (ArrayViewMut1<'_, T>, ArrayViewMut1<'_, T>) {
       assert_ne!(l1, l2);
   
       let raw_view = arr.raw_view_mut();
       std::mem::forget(arr);
   
       // This is safe because the original data was valid (since it was an
       // `ArrayViewMut`), we've ensured that `l1 != l2`, `arr` no longer exists,
       // and the function signature ensures that the lifetimes of the two lanes
       // are correct.
       unsafe {
           (
               raw_view.clone().index_axis_move(axis, l1).deref_into_view_mut(),
               raw_view.index_axis_move(axis, l2).deref_into_view_mut(),
           )
       }
   }
   
   pub fn apply_two_lanes<T, F>(arr: &mut Array2<T>, axis: Axis, l1: usize, l2: usize, f: F)
   where
       F: Fn(&mut T, &mut T),
   {
       let (lane1, lane2) = index_two_lanes_mut(arr.view_mut(), axis, l1, l2);
       ndarray::Zip::from(lane1).and(lane2).apply(f);
   }
   
   fn main() {
       let mut a = array![[1, 2, 3], [4, 5, 6]];
       apply_two_lanes(&mut a, Axis(0), 0, 1, |x1, x2| *x1 += *x2);
       assert_eq!(a, array![[5, 7, 9], [4, 5, 6]]);
   }

   The next release of ndarray will have multi_slice_mut and multi_slice_move methods (#717), which make it possible to create multiple mutable (disjoint) slices of an array without unsafe more cleanly than other methods. Unfortunately, they don't have an axis parameter like in your example; you'd have to handle the axis == Axis(0) and axis == Axis(1) cases separately.

   We should add multi_slice_axis_mut and multi_slice_axis_move methods, in the same way that slice_axis_mut is analogous to slice_mut.

   🎉React with 🎉1Reacted by Jason Miller
2. Wojtek242 commented on Jan 21, 2020

   

   Wojtek242

   on Jan 21, 2020

   Author

   More actions

   Awesome! Thanks for the advice. Just for my understanding before I close the issue. Unsafe is necessary, because I need to dereference a RawArrayView. In turn, I need a RawArrayView, because I cannot clone an ArrayView. Is that correct?
3. jturner314 commented on Feb 14, 2020

   

   jturner314

   on Feb 14, 2020

   Member

   More actions

   That's mostly correct. The type of raw_view is RawArrayViewMut, not RawArrayView. It's possible to clone an ArrayView, but not an ArrayViewMut. If we allowed cloning ArrayViewMut instances, then it would be possible to get two mutable references to the same data, which would violate Rust's aliasing rules.

   You can think of Array<A, D> as analogous to Vec<A>, ArrayView<'a, A, D> as analogous to &'a [A], and ArrayViewMut<'a, A, D> as analogous to &'a mut [A]. Clone is implemented for Array<A, D>/Vec<A> and ArrayView<'a, A, D>/&'a [A] but not ArrayViewMut<'a, A, D>/&'a mut [A] due to Rust's aliasing rules.

   AFAIK, the only way without unsafe to get mutable references to two elements in a Vec at the same time is to use an iterator (as in your original comment) or to use split_at_mut/split_first_mut/split_last_mut. In ndarray 0.13.0, the same is true of ndarray's types (use an iterator or a splitting method.) Both of those options require a test for which index is less.
4. Wojtek242 commented on Apr 3, 2020

   

   Wojtek242

   on Apr 3, 2020

   Author

   More actions

   Thanks!
5. Wojtek242

   closed this as completedon Apr 3, 2020