From 41da875faef58e618cafc7dfdc5f3985a58f1e98 Mon Sep 17 00:00:00 2001 From: Andre Bogus Date: Sun, 18 Dec 2022 15:40:46 +0100 Subject: [PATCH] Add `Option::as_slice`(`_mut`) This adds the following functions: * `Option::as_slice(&self) -> &[T]` * `Option::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` is equal to that of `Option>`. --- library/core/src/lib.rs | 1 + library/core/src/option.rs | 119 +++++++++++++++++++++++++++++++ tests/codegen/option-as-slice.rs | 28 ++++++++ 3 files changed, 148 insertions(+) create mode 100644 tests/codegen/option-as-slice.rs diff --git a/library/core/src/lib.rs b/library/core/src/lib.rs index dc0702c467a4e..7b8f3880d3e50 100644 --- a/library/core/src/lib.rs +++ b/library/core/src/lib.rs @@ -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)] diff --git a/library/core/src/option.rs b/library/core/src/option.rs index 5d5e95590344a..994c08d1fb50d 100644 --- a/library/core/src/option.rs +++ b/library/core/src/option.rs @@ -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 Option { } } + const fn get_some_offset() -> isize { + if mem::size_of::>() == mem::size_of::() { + // niche optimization means the `T` is always stored at the same position as the Option. + 0 + } else { + assert!(mem::size_of::>() == mem::size_of::>>()); + let some_uninit = Some(mem::MaybeUninit::::uninit()); + // SAFETY: This gets the byte offset of the `Some(_)` value following the fact that + // niche optimization is not active, and thus Option and Option> share + // the same layout. + unsafe { + (some_uninit.as_ref().unwrap() as *const mem::MaybeUninit) + .byte_offset_from(&some_uninit as *const Option>) + } + } + } + + /// 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).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`, 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).wrapping_byte_offset(Self::get_some_offset()) as *mut T, + self.is_some() as usize, + ) + } + } + ///////////////////////////////////////////////////////////////////////// // Getting to contained values ///////////////////////////////////////////////////////////////////////// diff --git a/tests/codegen/option-as-slice.rs b/tests/codegen/option-as-slice.rs new file mode 100644 index 0000000000000..d5077dbf6ccd9 --- /dev/null +++ b/tests/codegen/option-as-slice.rs @@ -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] { + // 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] { + // CHECK: start: + // CHECK-NOT: select + // CHECK: ret + o.as_slice() +}