Add `Option::as_slice`(`_mut`)

This adds the following functions:

* `Option<T>::as_slice(&self) -> &[T]`
* `Option<T>::as_slice_mut(&mut self) -> &[T]`

The `as_slice` and `as_slice_mut` functions benefit from an
optimization that makes them completely branch-free.

Note that the optimization's soundness hinges on the fact that either
the niche optimization makes the offset of the `Some(_)` contents zero
or the mempory layout of `Option<T>` is equal to that of
`Option<MaybeUninit<T>>`.
This commit is contained in:
Andre Bogus 2022-12-18 15:40:46 +01:00
parent bd4a96a12d
commit 41da875fae
3 changed files with 148 additions and 0 deletions

View File

@ -134,6 +134,7 @@
#![feature(const_option)]
#![feature(const_option_ext)]
#![feature(const_pin)]
#![feature(const_pointer_byte_offsets)]
#![feature(const_pointer_is_aligned)]
#![feature(const_ptr_sub_ptr)]
#![feature(const_replace)]

View File

@ -553,6 +553,7 @@ use crate::pin::Pin;
use crate::{
cmp, convert, hint, mem,
ops::{self, ControlFlow, Deref, DerefMut},
slice,
};
/// The `Option` type. See [the module level documentation](self) for more.
@ -734,6 +735,124 @@ impl<T> Option<T> {
}
}
const fn get_some_offset() -> isize {
if mem::size_of::<Option<T>>() == mem::size_of::<T>() {
// niche optimization means the `T` is always stored at the same position as the Option.
0
} else {
assert!(mem::size_of::<Option<T>>() == mem::size_of::<Option<mem::MaybeUninit<T>>>());
let some_uninit = Some(mem::MaybeUninit::<T>::uninit());
// SAFETY: This gets the byte offset of the `Some(_)` value following the fact that
// niche optimization is not active, and thus Option<T> and Option<MaybeUninit<t>> share
// the same layout.
unsafe {
(some_uninit.as_ref().unwrap() as *const mem::MaybeUninit<T>)
.byte_offset_from(&some_uninit as *const Option<mem::MaybeUninit<T>>)
}
}
}
/// Returns a slice of the contained value, if any. If this is `None`, an
/// empty slice is returned. This can be useful to have a single type of
/// iterator over an `Option` or slice.
///
/// Note: Should you have an `Option<&T>` and wish to get a slice of `T`,
/// you can unpack it via `opt.map_or(&[], std::slice::from_ref)`.
///
/// # Examples
///
/// ```rust
/// #![feature(option_as_slice)]
///
/// assert_eq!(
/// [Some(1234).as_slice(), None.as_slice()],
/// [&[1234][..], &[][..]],
/// );
/// ```
///
/// The inverse of this function is (discounting
/// borrowing) [`[_]::first`](slice::first):
///
/// ```rust
/// #![feature(option_as_slice)]
///
/// for i in [Some(1234_u16), None] {
/// assert_eq!(i.as_ref(), i.as_slice().first());
/// }
/// ```
#[inline]
#[must_use]
#[unstable(feature = "option_as_slice", issue = "108545")]
pub fn as_slice(&self) -> &[T] {
// SAFETY: This is sound as long as `get_some_offset` returns the
// correct offset. Though in the `None` case, the slice may be located
// at a pointer pointing into padding, the fact that the slice is
// empty, and the padding is at a properly aligned position for a
// value of that type makes it sound.
unsafe {
slice::from_raw_parts(
(self as *const Option<T>).wrapping_byte_offset(Self::get_some_offset())
as *const T,
self.is_some() as usize,
)
}
}
/// Returns a mutable slice of the contained value, if any. If this is
/// `None`, an empty slice is returned. This can be useful to have a
/// single type of iterator over an `Option` or slice.
///
/// Note: Should you have an `Option<&mut T>` instead of a
/// `&mut Option<T>`, which this method takes, you can obtain a mutable
/// slice via `opt.map_or(&mut [], std::slice::from_mut)`.
///
/// # Examples
///
/// ```rust
/// #![feature(option_as_slice)]
///
/// assert_eq!(
/// [Some(1234).as_mut_slice(), None.as_mut_slice()],
/// [&mut [1234][..], &mut [][..]],
/// );
/// ```
///
/// The result is a mutable slice of zero or one items that points into
/// our original `Option`:
///
/// ```rust
/// #![feature(option_as_slice)]
///
/// let mut x = Some(1234);
/// x.as_mut_slice()[0] += 1;
/// assert_eq!(x, Some(1235));
/// ```
///
/// The inverse of this method (discounting borrowing)
/// is [`[_]::first_mut`](slice::first_mut):
///
/// ```rust
/// #![feature(option_as_slice)]
///
/// assert_eq!(Some(123).as_mut_slice().first_mut(), Some(&mut 123))
/// ```
#[inline]
#[must_use]
#[unstable(feature = "option_as_slice", issue = "108545")]
pub fn as_mut_slice(&mut self) -> &mut [T] {
// SAFETY: This is sound as long as `get_some_offset` returns the
// correct offset. Though in the `None` case, the slice may be located
// at a pointer pointing into padding, the fact that the slice is
// empty, and the padding is at a properly aligned position for a
// value of that type makes it sound.
unsafe {
slice::from_raw_parts_mut(
(self as *mut Option<T>).wrapping_byte_offset(Self::get_some_offset()) as *mut T,
self.is_some() as usize,
)
}
}
/////////////////////////////////////////////////////////////////////////
// Getting to contained values
/////////////////////////////////////////////////////////////////////////

View File

@ -0,0 +1,28 @@
// compile-flags: -O
// only-x86_64
#![crate_type = "lib"]
#![feature(option_as_slice)]
extern crate core;
use core::num::NonZeroU64;
use core::option::Option;
// CHECK-LABEL: @u64_opt_as_slice
#[no_mangle]
pub fn u64_opt_as_slice(o: &Option<u64>) -> &[u64] {
// CHECK: start:
// CHECK-NOT: select
// CHECK: ret
o.as_slice()
}
// CHECK-LABEL: @nonzero_u64_opt_as_slice
#[no_mangle]
pub fn nonzero_u64_opt_as_slice(o: &Option<NonZeroU64>) -> &[NonZeroU64] {
// CHECK: start:
// CHECK-NOT: select
// CHECK: ret
o.as_slice()
}