From 3fe37b8c6e02daa6a7c33a71bb7f44682e43d08d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Marvin=20L=C3=B6bel?= Date: Sun, 28 Mar 2021 14:24:49 +0200 Subject: [PATCH] Add get_many_mut methods to slice --- library/core/src/error.rs | 3 + library/core/src/slice/mod.rs | 136 ++++++++++++++++++++++++++++++++++ library/core/tests/lib.rs | 1 + library/core/tests/slice.rs | 60 +++++++++++++++ library/std/src/lib.rs | 1 + 5 files changed, 201 insertions(+) diff --git a/library/core/src/error.rs b/library/core/src/error.rs index c9c7cdf4def..7152300abcb 100644 --- a/library/core/src/error.rs +++ b/library/core/src/error.rs @@ -506,3 +506,6 @@ fn description(&self) -> &str { #[unstable(feature = "cstr_from_bytes_until_nul", issue = "95027")] impl Error for crate::ffi::FromBytesUntilNulError {} + +#[unstable(feature = "get_many_mut", issue = "104642")] +impl Error for crate::slice::GetManyMutError {} diff --git a/library/core/src/slice/mod.rs b/library/core/src/slice/mod.rs index ad90e00b8a0..04486ed2d14 100644 --- a/library/core/src/slice/mod.rs +++ b/library/core/src/slice/mod.rs @@ -7,6 +7,7 @@ #![stable(feature = "rust1", since = "1.0.0")] use crate::cmp::Ordering::{self, Greater, Less}; +use crate::fmt; use crate::intrinsics::{assert_unsafe_precondition, exact_div}; use crate::marker::Copy; use crate::mem::{self, SizedTypeProperties}; @@ -4082,6 +4083,88 @@ pub fn take_last_mut<'a>(self: &mut &'a mut Self) -> Option<&'a mut T> { *self = rem; Some(last) } + + /// Returns mutable references to many indices at once, without doing any checks. + /// + /// For a safe alternative see [`get_many_mut`]. + /// + /// # Safety + /// + /// Calling this method with overlapping or out-of-bounds indices is *[undefined behavior]* + /// even if the resulting references are not used. + /// + /// # Examples + /// + /// ``` + /// #![feature(get_many_mut)] + /// + /// let x = &mut [1, 2, 4]; + /// + /// unsafe { + /// let [a, b] = x.get_many_unchecked_mut([0, 2]); + /// *a *= 10; + /// *b *= 100; + /// } + /// assert_eq!(x, &[10, 2, 400]); + /// ``` + /// + /// [`get_many_mut`]: slice::get_many_mut + /// [undefined behavior]: https://doc.rust-lang.org/reference/behavior-considered-undefined.html + #[unstable(feature = "get_many_mut", issue = "104642")] + #[inline] + pub unsafe fn get_many_unchecked_mut( + &mut self, + indices: [usize; N], + ) -> [&mut T; N] { + // NB: This implementation is written as it is because any variation of + // `indices.map(|i| self.get_unchecked_mut(i))` would make miri unhappy, + // or generate worse code otherwise. This is also why we need to go + // through a raw pointer here. + let slice: *mut [T] = self; + let mut arr: mem::MaybeUninit<[&mut T; N]> = mem::MaybeUninit::uninit(); + let arr_ptr = arr.as_mut_ptr(); + + // SAFETY: We expect `indices` to contain disjunct values that are + // in bounds of `self`. + unsafe { + for i in 0..N { + let idx = *indices.get_unchecked(i); + *(*arr_ptr).get_unchecked_mut(i) = &mut *slice.get_unchecked_mut(idx); + } + arr.assume_init() + } + } + + /// Returns mutable references to many indices at once. + /// + /// Returns an error if any index is out-of-bounds, or if the same index was + /// passed more than once. + /// + /// # Examples + /// + /// ``` + /// #![feature(get_many_mut)] + /// + /// let v = &mut [1, 2, 3]; + /// if let Ok([a, b]) = v.get_many_mut([0, 2]) { + /// *a = 413; + /// *b = 612; + /// } + /// assert_eq!(v, &[413, 2, 612]); + /// ``` + #[unstable(feature = "get_many_mut", issue = "104642")] + #[inline] + pub fn get_many_mut( + &mut self, + indices: [usize; N], + ) -> Result<[&mut T; N], GetManyMutError> { + if !get_many_check_valid(&indices, self.len()) { + return Err(GetManyMutError { _private: () }); + } + // SAFETY: The `get_many_check_valid()` call checked that all indices + // are disjunct and in bounds. + unsafe { Ok(self.get_many_unchecked_mut(indices)) } + } } impl [[T; N]] { @@ -4304,3 +4387,56 @@ fn as_slice(&self) -> &[Self::Item] { self } } + +/// This checks every index against each other, and against `len`. +/// +/// This will do `binomial(N + 1, 2) = N * (N + 1) / 2 = 0, 1, 3, 6, 10, ..` +/// comparison operations. +fn get_many_check_valid(indices: &[usize; N], len: usize) -> bool { + // NB: The optimzer should inline the loops into a sequence + // of instructions without additional branching. + let mut valid = true; + for (i, &idx) in indices.iter().enumerate() { + valid &= idx < len; + for &idx2 in &indices[..i] { + valid &= idx != idx2; + } + } + valid +} + +/// The error type returned by [`get_many_mut`][`slice::get_many_mut`]. +/// +/// It indicates one of two possible errors: +/// - An index is out-of-bounds. +/// - The same index appeared multiple times in the array. +/// +/// # Examples +/// +/// ``` +/// #![feature(get_many_mut)] +/// +/// let v = &mut [1, 2, 3]; +/// assert!(v.get_many_mut([0, 999]).is_err()); +/// assert!(v.get_many_mut([1, 1]).is_err()); +/// ``` +#[unstable(feature = "get_many_mut", issue = "104642")] +// NB: The N here is there to be forward-compatible with adding more details +// to the error type at a later point +pub struct GetManyMutError { + _private: (), +} + +#[unstable(feature = "get_many_mut", issue = "104642")] +impl fmt::Debug for GetManyMutError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("GetManyMutError").finish_non_exhaustive() + } +} + +#[unstable(feature = "get_many_mut", issue = "104642")] +impl fmt::Display for GetManyMutError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Display::fmt("an index is out of bounds or appeared multiple times in the array", f) + } +} diff --git a/library/core/tests/lib.rs b/library/core/tests/lib.rs index 4c0c0da652f..abbad96c4da 100644 --- a/library/core/tests/lib.rs +++ b/library/core/tests/lib.rs @@ -108,6 +108,7 @@ #![feature(provide_any)] #![feature(utf8_chunks)] #![feature(is_ascii_octdigit)] +#![feature(get_many_mut)] #![deny(unsafe_op_in_unsafe_fn)] extern crate test; diff --git a/library/core/tests/slice.rs b/library/core/tests/slice.rs index 9e1fbea7914..4e06e0f4398 100644 --- a/library/core/tests/slice.rs +++ b/library/core/tests/slice.rs @@ -2595,3 +2595,63 @@ fn test_flatten_mut_size_overflow() { let x = &mut [[(); usize::MAX]; 2][..]; let _ = x.flatten_mut(); } + +#[test] +fn test_get_many_mut_normal_2() { + let mut v = vec![1, 2, 3, 4, 5]; + let [a, b] = v.get_many_mut([3, 0]).unwrap(); + *a += 10; + *b += 100; + assert_eq!(v, vec![101, 2, 3, 14, 5]); +} + +#[test] +fn test_get_many_mut_normal_3() { + let mut v = vec![1, 2, 3, 4, 5]; + let [a, b, c] = v.get_many_mut([0, 4, 2]).unwrap(); + *a += 10; + *b += 100; + *c += 1000; + assert_eq!(v, vec![11, 2, 1003, 4, 105]); +} + +#[test] +fn test_get_many_mut_empty() { + let mut v = vec![1, 2, 3, 4, 5]; + let [] = v.get_many_mut([]).unwrap(); + assert_eq!(v, vec![1, 2, 3, 4, 5]); +} + +#[test] +fn test_get_many_mut_single_first() { + let mut v = vec![1, 2, 3, 4, 5]; + let [a] = v.get_many_mut([0]).unwrap(); + *a += 10; + assert_eq!(v, vec![11, 2, 3, 4, 5]); +} + +#[test] +fn test_get_many_mut_single_last() { + let mut v = vec![1, 2, 3, 4, 5]; + let [a] = v.get_many_mut([4]).unwrap(); + *a += 10; + assert_eq!(v, vec![1, 2, 3, 4, 15]); +} + +#[test] +fn test_get_many_mut_oob_nonempty() { + let mut v = vec![1, 2, 3, 4, 5]; + assert!(v.get_many_mut([5]).is_err()); +} + +#[test] +fn test_get_many_mut_oob_empty() { + let mut v: Vec = vec![]; + assert!(v.get_many_mut([0]).is_err()); +} + +#[test] +fn test_get_many_mut_duplicate() { + let mut v = vec![1, 2, 3, 4, 5]; + assert!(v.get_many_mut([1, 3, 3, 4]).is_err()); +} diff --git a/library/std/src/lib.rs b/library/std/src/lib.rs index 9334c833bb6..63ee6c521d7 100644 --- a/library/std/src/lib.rs +++ b/library/std/src/lib.rs @@ -347,6 +347,7 @@ #![feature(stdsimd)] #![feature(test)] #![feature(trace_macros)] +#![feature(get_many_mut)] // // Only used in tests/benchmarks: //