Rust: [ER] get_two_mut

[leonardo-m]

In various cases I'd like a slice method similar to slice::split_at_mut but simpler, that returns just two items at different indexes:

fn get_two_mut(&mut self, i: usize, j: usize) -> (&mut T, &mut T)

That panics if i == j or i >= self.len() or i >= self.len().

[jyn514]

You can write this without needing to go through the standard library: https://play.rust-lang.org/?version=stable&mode=debug&edition=2018&gist=98fbe51263eeccb354ee3abef6e3e8eb

[leonardo-m]

Yeah, I can also write this, but I think it's worth adding to stdlib:

fn get_two_mut<T, const N: usize>(this: &mut [T; N], i: usize, j: usize) -> (&mut T, &mut T) {
    assert!(i < N && j < N && i != j);
    let ptr = this.as_mut_ptr();
    unsafe { (transmute(ptr.add(i)), transmute(ptr.add(j))) }
}

(Here I've written it for an array, but for a slice it's very similar).

[jyn514]

I'm not sure your example is sound ... certainly not according to stacked borrows. You're taking a pointer with provenance for the whole array and making two &mut T out of it, which means it aliases.

[HeroicKatora]

Miri says it's fine and so would my intuition. Raw pointers, unlike references, do not assert their tags upon using their value but only when actually dereferenced and only for the place they are dereferenced to. But that only occurs as part of producing the &mut _ vaules from the transmutes and those do _not_ alias. If we insert some intermediate

let a = &mut *(ptr as *mut [T; N]);

after the construction of at least one of the return values _then_ it rightfully complains.

That said, this is unecessary¹. This method can be written with entirely safe Rust. (I faintly remember already seeing this question/writing it out somewhere. It might be interesting to document the pattern somewhere).

fn get_two_mut<T, const N: usize>(this: &mut [T; N], i: usize, j: usize) -> (&mut T, &mut T) {
    let (i, j) = if i < j { (i,j) } else { (j,i) };
    assert!(i < j && j < N);

    let s0 = i;
    // The position of the second split relative to the first.
    let s1 = j - i;

    let (_, split) = this.split_at_mut(s0);
    let (first, second) = split.split_at_mut(s1);
    (&mut first[0], &mut second[0])
}

¹Sadly llvm optimization is a lot worse than anticipated and this does no less than _three_ panic bounds checks.

For T = u32 and N = 64

    push    rax
    cmp rsi, rdx
    mov rax, rdx
    cmovb   rax, rsi
    cmovb   rsi, rdx
    mov rcx, rsi
    sub rcx, rax
    jbe .LBB7_6 ;; The manual panic
# %bb.1:
    cmp rsi, 64
    jae .LBB7_6 ;; The manual panic
# %bb.2:
    cmp rax, 65
    jae .LBB7_7 ;; "assertion failed: mid <= self.len()"
# %bb.3:
    mov edx, 64
    sub rdx, rax
    cmp rdx, rcx
    jb  .LBB7_7 ;; "assertion failed: mid <= self.len()"
# %bb.4:
    test    rcx, rcx
    je  .LBB7_8 ;; to core::panicking::panic_bounds_check@GOTPCREL
# %bb.5:
    lea rdx, [rdi + 4*rsi]
    lea rax, [rdi + 4*rax]
    pop rcx
    ret

[leonardo-m]

The asm for my version (on slices):

use std::mem::transmute;

type T = u32;

pub fn get_two_mut(this: &mut [T], i: usize, j: usize) -> (&mut T, &mut T) {
    assert!(i < this.len() && j < this.len() && i != j);
    let ptr = this.as_mut_ptr();
    unsafe { (transmute(ptr.add(i)), transmute(ptr.add(j))) }
}

Asm (optimized build):

get_two_mut:
    cmp     rdx, rcx
    je      .LBB7_4
    cmp     rdx, rsi
    jae     .LBB7_4
    cmp     rcx, rsi
    jae     .LBB7_4
    lea     rax, [rdi + 4*rdx]
    lea     rdx, [rdi + 4*rcx]
    ret
.LBB7_4:
    push    rax
    call    std::panicking::begin_panic
    ud2

[SkiFire13]

@HeroicKatora your implementation is wrong for inputs where i > j.

This is my take:

pub fn get_two_mut_safe(this: &mut [T], i: usize, j: usize) -> (&mut T, &mut T) {
    assert!(i < this.len() && j < this.len() && i != j);
    if i > j {
        let (first, second) = this.split_at_mut(i);
        (&mut second[0], &mut first[j])
    } else {
        let (first, second) = this.split_at_mut(j);
        (&mut first[i], &mut second[0])
    }
}

This safe implementation is compiled to the same assembly as @leonardo-m's unsafe version. https://rust.godbolt.org/z/36c4sx

Curiously swapping the two if branches (checking for i<j first) adds a mov to the assembly.

[HeroicKatora]

your implementation is wrong for inputs where i > j.

@SkiFire13 Could you explain how? Although it's not incredibly important because your implementation looks incredibly pleasing and I'm going to use it going forward :) Note the normalization in the first line that ensures that always i < j from that point onwards.

[SkiFire13]

I would expect this function to always return the equivalent of (&mut this[i], &mut this[j]) if i != j and both are < this.len(), however if i > j then your code returns (&mut this[j], &mut this[i]) which is swapped compared to what I expected.

[RalfJung]

[ER] get_two_mut

Out of curiosity, what is "ER" for?

[leonardo-m]

Out of curiosity, what is "ER" for?

Enhancement Request.