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.
This commit is contained in:
Scott McMurray 2023-03-01 10:57:46 -08:00
parent 5423745db8
commit f6a57c1955
2 changed files with 92 additions and 42 deletions

View File

@ -735,22 +735,43 @@ impl<T> Option<T> {
} }
} }
const fn get_some_offset() -> isize { /// This is a guess at how many bytes into the option the payload can be found.
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. /// For niche-optimized types it's correct because it's pigeon-holed to only
0 /// 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.
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 {
// The offset is at least inside the object, so let's try it.
offset
} else { } else {
assert!(mem::size_of::<Option<T>>() == mem::size_of::<Option<mem::MaybeUninit<T>>>()); // The offset guess is definitely wrong, so use the address
let some_uninit = Some(mem::MaybeUninit::<T>::uninit()); // of the original option since we have it already.
// SAFETY: This gets the byte offset of the `Some(_)` value following the fact that // This also correctly handles the case of layout-optimized enums
// niche optimization is not active, and thus Option<T> and Option<MaybeUninit<t>> share // where `max_offset == 0` and thus this is the only possibility.
// the same layout. 0
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 /// 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 /// empty slice is returned. This can be useful to have a single type of
@ -784,18 +805,28 @@ impl<T> Option<T> {
#[must_use] #[must_use]
#[unstable(feature = "option_as_slice", issue = "108545")] #[unstable(feature = "option_as_slice", issue = "108545")]
pub fn as_slice(&self) -> &[T] { pub fn as_slice(&self) -> &[T] {
// SAFETY: This is sound as long as `get_some_offset` returns the let payload_ptr: *const T =
// correct offset. Though in the `None` case, the slice may be located // The goal here is that both arms here are calculating exactly
// at a pointer pointing into padding, the fact that the slice is // the same pointer, and thus it'll be folded away when the guessed
// empty, and the padding is at a properly aligned position for a // offset is correct, but if the guess is wrong for some reason
// value of that type makes it sound. // it'll at least still be sound, just no longer optimal.
unsafe { if let Some(payload) = self {
slice::from_raw_parts( payload
(self as *const Option<T>).wrapping_byte_offset(Self::get_some_offset()) } else {
as *const T, let self_ptr: *const Self = self;
self.is_some() as usize, // 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 /// Returns a mutable slice of the contained value, if any. If this is
@ -840,17 +871,28 @@ impl<T> Option<T> {
#[must_use] #[must_use]
#[unstable(feature = "option_as_slice", issue = "108545")] #[unstable(feature = "option_as_slice", issue = "108545")]
pub fn as_mut_slice(&mut self) -> &mut [T] { pub fn as_mut_slice(&mut self) -> &mut [T] {
// SAFETY: This is sound as long as `get_some_offset` returns the let payload_ptr: *mut T =
// correct offset. Though in the `None` case, the slice may be located // The goal here is that both arms here are calculating exactly
// at a pointer pointing into padding, the fact that the slice is // the same pointer, and thus it'll be folded away when the guessed
// empty, and the padding is at a properly aligned position for a // offset is correct, but if the guess is wrong for some reason
// value of that type makes it sound. // it'll at least still be sound, just no longer optimal.
unsafe { if let Some(payload) = self {
slice::from_raw_parts_mut( payload
(self as *mut Option<T>).wrapping_byte_offset(Self::get_some_offset()) as *mut T, } else {
self.is_some() as usize, 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) }
} }
///////////////////////////////////////////////////////////////////////// /////////////////////////////////////////////////////////////////////////

View File

@ -1,4 +1,4 @@
// compile-flags: -O // compile-flags: -O -Z randomize-layout=no
// only-x86_64 // only-x86_64
#![crate_type = "lib"] #![crate_type = "lib"]
@ -12,17 +12,25 @@ use core::option::Option;
// CHECK-LABEL: @u64_opt_as_slice // CHECK-LABEL: @u64_opt_as_slice
#[no_mangle] #[no_mangle]
pub fn u64_opt_as_slice(o: &Option<u64>) -> &[u64] { pub fn u64_opt_as_slice(o: &Option<u64>) -> &[u64] {
// CHECK: start:
// CHECK-NOT: select // CHECK-NOT: select
// CHECK: ret // CHECK-NOT: br
// CHECK-NOT: switch
// CHECK-NOT: icmp
o.as_slice() o.as_slice()
} }
// CHECK-LABEL: @nonzero_u64_opt_as_slice // CHECK-LABEL: @nonzero_u64_opt_as_slice
#[no_mangle] #[no_mangle]
pub fn nonzero_u64_opt_as_slice(o: &Option<NonZeroU64>) -> &[NonZeroU64] { pub fn nonzero_u64_opt_as_slice(o: &Option<NonZeroU64>) -> &[NonZeroU64] {
// CHECK: start:
// CHECK-NOT: select // 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() o.as_slice()
} }