Iterate mutably over two lanes of an array

[Wojtek242]

I asked this question on reddit, but I didn't get any response over there. I'm currently using ndarray for a project of mine and one thing I need to quite frequently is to mutate the contents of two or more columns/rows based on each others contents. It isn't always the case that both lanes need to be mutable, but they usually need to be. I've figured out how to do it with ndarray::Zip and iter_axis_mut as follows

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),
{
    assert_ne!(l1, l2);

    let mut lanes = arr.axis_iter_mut(axis);
    let (lane1, lane2) = if l1 < l2 {
        (lanes.nth(l1).unwrap(), lanes.nth(l2 - l1 - 1).unwrap())
    } else {
        let (lane2, lane1) =
            (lanes.nth(l2).unwrap(), lanes.nth(l1 - l2 - 1).unwrap());
        (lane1, lane2)
    };
    ndarray::Zip::from(lane1).and(lane2).apply(f);
}

However, as you can see, because I don't know if l1 or l2 comes first, I need the additional if test. I haven't managed to figure out how to do the above without some manner of checking of the ordering of l1 and l2. Is there a way to accomplish the above without having to check which index is first?

[jturner314]

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.

[Wojtek242]

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?

[jturner314]

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.

[Wojtek242]

Thanks!