[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.

[leonardo-m]

I think the new get_many_mut method is enough, I think there isn't a need for get_two_mut now.