Auto merge of #108623 - scottmcm:try-different-as-slice-impl, r=the8472

Move `Option::as_slice` to an always-sound implementation

This approach depends on CSE to not have any branches or selects when the guessed offset is correct -- which it always will be right now -- but to also be *sound* (just less efficient) if the layout algorithms change such that the guess is incorrect.

The codegen test confirms that CSE handles this as expected, leaving the optimal codegen.

cc JakobDegen #108545

Author: bors

library/core/src/option.rs

@@ -735,22 +735,47 @@ 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
+    /// This is a guess at how many bytes into the option the payload can be found.
+    ///
+    /// For niche-optimized types it's correct because it's pigeon-holed to only
+    /// one possible place.  For other types, it's usually correct today, but
+    /// tweaks to the layout algorithm (particularly expansions of
+    /// `-Z randomize-layout`) might make it incorrect at any point.
+    ///
+    /// It's guaranteed to be a multiple of alignment (so will always give a
+    /// correctly-aligned location) and to be within the allocated object, so
+    /// is valid to use with `offset` and to use for a zero-sized read.
+    ///
+    /// FIXME: This is a horrible hack, but allows a nice optimization.  It should
+    /// be replaced with `offset_of!` once that works on enum variants.
+    const SOME_BYTE_OFFSET_GUESS: isize = {
+        let some_uninit = Some(mem::MaybeUninit::<T>::uninit());
+        let payload_ref = some_uninit.as_ref().unwrap();
+        // SAFETY: `as_ref` gives an address inside the existing `Option`,
+        // so both pointers are derived from the same thing and the result
+        // cannot overflow an `isize`.
+        let offset = unsafe { <*const _>::byte_offset_from(payload_ref, &some_uninit) };
+
+        // The offset is into the object, so it's guaranteed to be non-negative.
+        assert!(offset >= 0);
+
+        // The payload and the overall option are aligned,
+        // so the offset will be a multiple of the alignment too.
+        assert!((offset as usize) % mem::align_of::<T>() == 0);
+
+        let max_offset = mem::size_of::<Self>() - mem::size_of::<T>();
+        if offset as usize <= max_offset {
+            // There's enough space after this offset for a `T` to exist without
+            // overflowing the bounds of the object, so let's try it.
+            offset
         } 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>>)
-            }
+            // The offset guess is definitely wrong, so use the address
+            // of the original option since we have it already.
+            // This also correctly handles the case of layout-optimized enums
+            // where `max_offset == 0` and thus this is the only possibility.
+            0
         }
-    }
+    };
 
     /// 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
@@ -784,18 +809,28 @@ impl<T> Option<T> {
     #[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,
-            )
-        }
+        let payload_ptr: *const T =
+            // The goal here is that both arms here are calculating exactly
+            // the same pointer, and thus it'll be folded away when the guessed
+            // offset is correct, but if the guess is wrong for some reason
+            // it'll at least still be sound, just no longer optimal.
+            if let Some(payload) = self {
+                payload
+            } else {
+                let self_ptr: *const Self = self;
+                // SAFETY: `SOME_BYTE_OFFSET_GUESS` guarantees that its value is
+                // such that this will be in-bounds of the object.
+                unsafe { self_ptr.byte_offset(Self::SOME_BYTE_OFFSET_GUESS).cast() }
+            };
+        let len = usize::from(self.is_some());
+
+        // SAFETY: When the `Option` is `Some`, we're using the actual pointer
+        // to the payload, with a length of 1, so this is equivalent to
+        // `slice::from_ref`, and thus is safe.
+        // When the `Option` is `None`, the length used is 0, so to be safe it
+        // just needs to be aligned, which it is because `&self` is aligned and
+        // the offset used is a multiple of alignment.
+        unsafe { slice::from_raw_parts(payload_ptr, len) }
     }
 
     /// Returns a mutable slice of the contained value, if any. If this is
@@ -840,17 +875,28 @@ impl<T> Option<T> {
     #[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,
-            )
-        }
+        let payload_ptr: *mut T =
+            // The goal here is that both arms here are calculating exactly
+            // the same pointer, and thus it'll be folded away when the guessed
+            // offset is correct, but if the guess is wrong for some reason
+            // it'll at least still be sound, just no longer optimal.
+            if let Some(payload) = self {
+                payload
+            } else {
+                let self_ptr: *mut Self = self;
+                // SAFETY: `SOME_BYTE_OFFSET_GUESS` guarantees that its value is
+                // such that this will be in-bounds of the object.
+                unsafe { self_ptr.byte_offset(Self::SOME_BYTE_OFFSET_GUESS).cast() }
+            };
+        let len = usize::from(self.is_some());
+
+        // SAFETY: When the `Option` is `Some`, we're using the actual pointer
+        // to the payload, with a length of 1, so this is equivalent to
+        // `slice::from_mut`, and thus is safe.
+        // When the `Option` is `None`, the length used is 0, so to be safe it
+        // just needs to be aligned, which it is because `&self` is aligned and
+        // the offset used is a multiple of alignment.
+        unsafe { slice::from_raw_parts_mut(payload_ptr, len) }
     }
 
     /////////////////////////////////////////////////////////////////////////

tests/codegen/option-as-slice.rs

@@ -1,4 +1,4 @@
-// compile-flags: -O
+// compile-flags: -O -Z randomize-layout=no
 // only-x86_64
 
 #![crate_type = "lib"]
@@ -12,17 +12,25 @@ 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
+    // CHECK-NOT: br
+    // CHECK-NOT: switch
+    // CHECK-NOT: icmp
     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
+    // CHECK-NOT: br
+    // CHECK-NOT: switch
+    // CHECK-NOT: icmp
+    // CHECK: %[[NZ:.+]] = icmp ne i64 %{{.+}}, 0
+    // CHECK-NEXT: zext i1 %[[NZ]] to i64
+    // CHECK-NOT: select
+    // CHECK-NOT: br
+    // CHECK-NOT: switch
+    // CHECK-NOT: icmp
     o.as_slice()
 }