rust/library/core/tests
Manish Goregaokar 1dd515f273
Rollup merge of #83608 - Kimundi:index_many, r=Mark-Simulacrum
Add slice methods for indexing via an array of indices.

Disclaimer: It's been a while since I contributed to the main Rust repo, apologies in advance if this is large enough already that it should've been an RFC.

---

# Update:

- Based on feedback, removed the `&[T]` variant of this API, and removed the requirements for the indices to be sorted.

# Description

This adds the following slice methods to `core`:

```rust
impl<T> [T] {
    pub unsafe fn get_many_unchecked_mut<const N: usize>(&mut self, indices: [usize; N]) -> [&mut T; N];
    pub fn get_many_mut<const N: usize>(&mut self, indices: [usize; N]) -> Option<[&mut T; N]>;
}
```

This allows creating multiple mutable references to disjunct positions in a slice, which previously required writing some awkward code with `split_at_mut()` or `iter_mut()`. For the bound-checked variant, the indices are checked against each other and against the bounds of the slice, which requires `N * (N + 1) / 2` comparison operations.

This has a proof-of-concept standalone implementation here: https://crates.io/crates/index_many

Care has been taken that the implementation passes miri borrow checks, and generates straight-forward assembly (though this was only checked on x86_64).

# Example

```rust
let v = &mut [1, 2, 3, 4];
let [a, b] = v.get_many_mut([0, 2]).unwrap();
std::mem::swap(a, b);
*v += 100;
assert_eq!(v, &[3, 2, 101, 4]);
```

# Codegen Examples

<details>
  <summary>Click to expand!</summary>

Disclaimer: Taken from local tests with the standalone implementation.

## Unchecked Indexing:

```rust
pub unsafe fn example_unchecked(slice: &mut [usize], indices: [usize; 3]) -> [&mut usize; 3] {
    slice.get_many_unchecked_mut(indices)
}
```

```nasm
example_unchecked:
 mov     rcx, qword, ptr, [r9]
 mov     r8, qword, ptr, [r9, +, 8]
 mov     r9, qword, ptr, [r9, +, 16]
 lea     rcx, [rdx, +, 8*rcx]
 lea     r8, [rdx, +, 8*r8]
 lea     rdx, [rdx, +, 8*r9]
 mov     qword, ptr, [rax], rcx
 mov     qword, ptr, [rax, +, 8], r8
 mov     qword, ptr, [rax, +, 16], rdx
 ret
```

## Checked Indexing (Option):

```rust
pub unsafe fn example_option(slice: &mut [usize], indices: [usize; 3]) -> Option<[&mut usize; 3]> {
    slice.get_many_mut(indices)
}
```

```nasm
 mov     r10, qword, ptr, [r9, +, 8]
 mov     rcx, qword, ptr, [r9, +, 16]
 cmp     rcx, r10
 je      .LBB0_7
 mov     r9, qword, ptr, [r9]
 cmp     rcx, r9
 je      .LBB0_7
 cmp     rcx, r8
 jae     .LBB0_7
 cmp     r10, r9
 je      .LBB0_7
 cmp     r9, r8
 jae     .LBB0_7
 cmp     r10, r8
 jae     .LBB0_7
 lea     r8, [rdx, +, 8*r9]
 lea     r9, [rdx, +, 8*r10]
 lea     rcx, [rdx, +, 8*rcx]
 mov     qword, ptr, [rax], r8
 mov     qword, ptr, [rax, +, 8], r9
 mov     qword, ptr, [rax, +, 16], rcx
 ret
.LBB0_7:
 mov     qword, ptr, [rax], 0
 ret
```

## Checked Indexing (Panic):

```rust
pub fn example_panic(slice: &mut [usize], indices: [usize; 3]) -> [&mut usize; 3] {
    let len = slice.len();
    match slice.get_many_mut(indices) {
        Some(s) => s,
        None => {
            let tmp = indices;
            index_many::sorted_bound_check_failed(&tmp, len)
        }
    }
}
```

```nasm
example_panic:
 sub     rsp, 56
 mov     rax, qword, ptr, [r9]
 mov     r10, qword, ptr, [r9, +, 8]
 mov     r9, qword, ptr, [r9, +, 16]
 cmp     r9, r10
 je      .LBB0_6
 cmp     r9, rax
 je      .LBB0_6
 cmp     r9, r8
 jae     .LBB0_6
 cmp     r10, rax
 je      .LBB0_6
 cmp     rax, r8
 jae     .LBB0_6
 cmp     r10, r8
 jae     .LBB0_6
 lea     rax, [rdx, +, 8*rax]
 lea     r8, [rdx, +, 8*r10]
 lea     rdx, [rdx, +, 8*r9]
 mov     qword, ptr, [rcx], rax
 mov     qword, ptr, [rcx, +, 8], r8
 mov     qword, ptr, [rcx, +, 16], rdx
 mov     rax, rcx
 add     rsp, 56
 ret
.LBB0_6:
 mov     qword, ptr, [rsp, +, 32], rax
 mov     qword, ptr, [rsp, +, 40], r10
 mov     qword, ptr, [rsp, +, 48], r9
 lea     rcx, [rsp, +, 32]
 mov     edx, 3
 call    index_many::bound_check_failed
 ud2
```
</details>

# Extensions

There are multiple optional extensions to this.

## Indexing With Ranges

This could easily be expanded to allow indexing with `[I; N]` where `I: SliceIndex<Self>`.  I wanted to keep the initial implementation simple, so I didn't include it yet.

## Panicking Variant

We could also add this method:

```rust
impl<T> [T] {
    fn index_many_mut<const N: usize>(&mut self, indices: [usize; N]) -> [&mut T; N];
}
```

This would work similar to the regular index operator and panic with out-of-bound indices. The advantage would be that we could more easily ensure good codegen with a useful panic message, which is non-trivial with the `Option` variant.

This is implemented in the standalone implementation, and used as basis for the codegen examples here and there.
2022-11-22 01:26:05 -05:00
..
fmt Add tests for rounding of ties during float formatting 2022-10-20 22:09:24 +02:00
hash Test const Hash, fix nits 2022-11-08 17:39:40 +01:00
iter VecDeque::resize should re-use the buffer in the passed-in element 2022-11-15 00:53:26 -08:00
num Auto merge of #102935 - ajtribick:display-float-0.5-fixed-0, r=scottmcm 2022-11-16 07:20:30 +00:00
ops
panic Fix test (location_const_file) 2022-10-08 11:48:53 +00:00
alloc.rs
any.rs Format dyn Trait better in type_name intrinsic 2022-11-01 20:41:47 +00:00
array.rs
ascii.rs introduce {char, u8}::is_ascii_octdigit 2022-09-27 11:55:13 +05:30
asserting.rs
atomic.rs Make use of [wrapping_]byte_{add,sub} 2022-08-23 19:32:37 +04:00
bool.rs
cell.rs
char.rs
clone.rs
cmp.rs
const_ptr.rs
convert.rs
future.rs
intrinsics.rs
lazy.rs
lib.rs Rollup merge of #83608 - Kimundi:index_many, r=Mark-Simulacrum 2022-11-22 01:26:05 -05:00
macros.rs
manually_drop.rs
mem.rs interpret: fix align_of_val on packed types 2022-10-29 15:58:32 +02:00
nonzero.rs
ops.rs
option.rs cfg-step code 2022-11-06 17:21:21 -05:00
panic.rs Add newlines 2022-09-27 19:23:52 +00:00
pattern.rs
pin_macro.rs
pin.rs
ptr.rs avoid non-strict-provenance casts in libcore tests 2022-11-20 09:58:29 +01:00
result.rs
simd.rs
slice.rs Add get_many_mut methods to slice 2022-11-20 11:19:11 -05:00
str_lossy.rs
str.rs
task.rs Made from_waker, waker, from_raw const 2022-09-14 14:53:16 +02:00
time.rs Fix Duration::{try_,}from_secs_f{32,64}(-0.0) 2022-10-14 16:07:09 +01:00
tuple.rs
unicode.rs
waker.rs