Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Move Option::as_slice to an always-sound implementation #108623

Merged
merged 2 commits into from
Mar 13, 2023
Merged

Move Option::as_slice to an always-sound implementation #108623

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f6a57c1
Move `Option::as_slice` to an always-sound implementation
scottmcm Mar 1, 2023
e975057
Clarify the text of some comments
scottmcm Mar 12, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
120 changes: 83 additions & 37 deletions library/core/src/option.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
the8472 marked this conversation as resolved.
Show resolved Hide resolved
///
/// 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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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) }
}

/////////////////////////////////////////////////////////////////////////
Expand Down
18 changes: 13 additions & 5 deletions tests/codegen/option-as-slice.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
// compile-flags: -O
// compile-flags: -O -Z randomize-layout=no
// only-x86_64

#![crate_type = "lib"]
Expand All @@ -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()
Copy link
Member Author

@scottmcm scottmcm Mar 1, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Output of the codegen test on my machine, in case it helps a reviewer:

; Function Attrs: mustprogress nofree nosync nounwind willreturn uwtable
define { ptr, i64 } @u64_opt_as_slice(ptr noalias noundef readonly align 8 dereferenceable(16) %o) unnamed_addr #0 {
start:
  %_3.i = load i64, ptr %o, align 8, !range !1, !alias.scope !2, !noundef !5
  %payload.i = getelementptr inbounds { i64, i64 }, ptr %o, i64 0, i32 1
  %0 = insertvalue { ptr, i64 } undef, ptr %payload.i, 0
  %1 = insertvalue { ptr, i64 } %0, i64 %_3.i, 1
  ret { ptr, i64 } %1
}

; Function Attrs: mustprogress nofree nosync nounwind willreturn uwtable
define { ptr, i64 } @nonzero_u64_opt_as_slice(ptr noalias noundef readonly align 8 dereferenceable(8) %o) unnamed_addr #0 {
start:
  %0 = load i64, ptr %o, align 8, !alias.scope !6, !noundef !5
  %.not3.i = icmp ne i64 %0, 0
  %len.i = zext i1 %.not3.i to i64
  %1 = insertvalue { ptr, i64 } undef, ptr %o, 0
  %2 = insertvalue { ptr, i64 } %1, i64 %len.i, 1
  ret { ptr, i64 } %2
}

!1 = !{i64 0, i64 2}

}