Auto merge of #105871 - llogiq:option-as-slice, r=scottmcm

Add `Option::as_`(`mut_`)`slice`

This adds the following functions:

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

The `as_slice` and `as_mut_slice_mut` functions benefit from an optimization that makes them completely branch-free. ~~Unfortunately, this optimization is not available on by-value Options, therefore the `into_slice` implementations use the plain `match` + `slice::from_ref` approach.~~

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

The idea has been discussed on [Zulip](https://rust-lang.zulipchat.com/#narrow/stream/219381-t-libs/topic/Option.3A.3Aas_slice). Notably the idea for the `as_slice_mut` and `into_slice´ methods came from `@cuviper` and `@Sp00ph` hardened the optimization against niche-optimized Options.

The [rust playground](https://play.rust-lang.org/?version=nightly&mode=release&edition=2021&gist=74f8e4239a19f454c183aaf7b4a969e0) shows that the generated assembly of the optimized method is basically only a copy while the naive method generates code containing a `test dx, dx` on x86_64.

---

EDIT from reviewer: ACP is rust-lang/libs-team#150

Author: bors

library/core/src/lib.rs

@@ -136,6 +136,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)]

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<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
     /////////////////////////////////////////////////////////////////////////

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>) -> &[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()
+}