Reorder an array according to given indexes

Here is an algorithm that reorders an array arr given an array of indices indices in such a way that every element arr[i] is moved to arr[indices[i]] after the reordering. This is done with only O(1) space.

Input:  arr[]   = [50, 40, 70, 60, 90]
        indices[] = [3,  0,  4,  1,  2]
Output: arr[]   = [40, 60, 90, 50, 70]
        indices[] = [0,  1,  2,  3,   4]

There is a crate implementing this algorithm, which contained some undefined behaviors. I am trying to fix it up but I can't be too sure the resulting code is UB-free.

The Issue

My Implementation:

pub fn reorder_index<A>(arr: &mut [A], indices: &mut [usize]) {
    assert_eq!(arr.len(), indices.len());

    let arr = unsafe { &mut *(arr as *mut [A] as *mut [MaybeUninit<A>]) };
    for i in 0..arr.len() {
        while indices[i] != i {
            let old_target_i = indices[indices[i]];

            //let old_target_e = arr[index[i]];
            let mut old_target_e = std::mem::MaybeUninit::<A>::uninit();
            unsafe {
                old_target_e
                    .as_mut_ptr()
                    .write(arr[indices[i]].as_ptr().read());
            }

            // arr[index[i]] = arr[i];
            // Since the object arr[index[i]] is aliased to old_target_e, we cannot produce
            // mutable reference to it nor assign to it.
            unsafe {
                arr[indices[i]].as_mut_ptr().write(arr[i].as_ptr().read());
            }
            // Now that old_target_e is not longer aliased, we can call assume_init
            let old_target_e = unsafe { old_target_e.assume_init() };

            indices[indices[i]] = indices[i];

            indices[i] = old_target_i;

            //arr[i] = old_target_e;
            // arr[i] is aliased to arr[index[i]
            arr[i].as_mut_ptr().write(old_target_e);
        }
    }
}

I tested somewhat thoroughly locally and the business logic seems to be correct. It would be great if you guys can spot some more memory bugs.

I actually implemented this algorithm a few weeks ago. You can find my implementation here. My implementation uses safe code only, so it is guaranteed to have no UB.

2 Likes

Isn’t this section just reimplementing slice::swap?

I'm pretty sure you are answering a different question: numpy array array indexing,

Input:  arr[]   = [50, 40, 70, 60, 90]
        indices[] = [3,  0,  4,  1,  2]
Correct Output: arr[]   = [40, 60, 90, 50, 70]
        indices[] = [0,  1,  2,  3,   4]
Wrong Output: arr[]   = [60, 50, 90, 40, 70]
        indices[] = [0,  1,  2,  3,   4]

but it just happens I need that algorithm as well, thanks nonetheless.

Yeah.

2 Likes

Right.. Anyway, use slice::swap instead of rewriting it yourself with raw pointers.

1 Like

Thanks, I totally missed that.

For anyone curious, if index is a permutation on n with c cycles, this algorithm performs n - c swaps (|X|-1 swaps for each cycle X; each swap decomposes a cycle into a fixed point and a cycle one shorter).

2 Likes

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.