Add intrinsic for Option::Some(_) offset

Author: llogiq

compiler/rustc_codegen_ssa/src/mir/intrinsic.rs

@@ -13,7 +13,7 @@ use rustc_middle::ty::{self, Ty, TyCtxt};
 use rustc_span::{sym, Span};
 use rustc_target::abi::{
     call::{FnAbi, PassMode},
-    WrappingRange,
+    FieldsShape, Variants, WrappingRange,
 };
 
 fn copy_intrinsic<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
@@ -104,6 +104,14 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     bx.const_usize(bx.layout_of(tp_ty).align.abi.bytes())
                 }
             }
+            sym::option_some_offset => {
+                let ty = substs.type_at(0);
+                let layout = bx.layout_of(ty);
+                let Variants::Multiple { variants, .. } = layout.layout.variants() else { bug!() };
+                let Some(variant) = variants.iter().last() else { bug!() };
+                let FieldsShape::Arbitrary { offsets, .. } = &variant.fields else { bug!() };
+                bx.const_usize(offsets[0].bytes())
+            }
             sym::vtable_size | sym::vtable_align => {
                 let vtable = args[0].immediate();
                 let idx = match name {

compiler/rustc_const_eval/src/interpret/intrinsics.rs

@@ -15,7 +15,7 @@ use rustc_middle::ty::layout::{LayoutOf as _, ValidityRequirement};
 use rustc_middle::ty::subst::SubstsRef;
 use rustc_middle::ty::{Ty, TyCtxt};
 use rustc_span::symbol::{sym, Symbol};
-use rustc_target::abi::{Abi, Align, Primitive, Size};
+use rustc_target::abi::{Abi, Align, FieldsShape, Primitive, Size, Variants};
 
 use super::{
     util::ensure_monomorphic_enough, CheckInAllocMsg, ImmTy, InterpCx, Machine, OpTy, PlaceTy,
@@ -106,6 +106,13 @@ pub(crate) fn eval_nullary_intrinsic<'tcx>(
             | ty::Tuple(_)
             | ty::Error(_) => ConstValue::from_target_usize(0u64, &tcx),
         },
+        sym::option_some_offset => {
+            let layout = tcx.layout_of(param_env.and(tp_ty)).map_err(|e| err_inval!(Layout(e)))?;
+            let Variants::Multiple { variants, .. } = layout.layout.variants() else { bug!() };
+            let Some(variant) = variants.iter().last() else { bug!() };
+            let FieldsShape::Arbitrary { offsets, .. } = &variant.fields else { bug!() };
+            ConstValue::from_target_usize(offsets[0].bytes(), &tcx)
+        }
         other => bug!("`{}` is not a zero arg intrinsic", other),
     })
 }

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -110,6 +110,7 @@ pub fn intrinsic_operation_unsafety(tcx: TyCtxt<'_>, intrinsic_id: DefId) -> hir
         | sym::rustc_peek
         | sym::maxnumf64
         | sym::type_name
+        | sym::option_some_offset
         | sym::forget
         | sym::black_box
         | sym::variant_count
@@ -214,6 +215,7 @@ pub fn check_intrinsic_type(tcx: TyCtxt<'_>, it: &hir::ForeignItem<'_>) {
 
             sym::type_name => (1, Vec::new(), tcx.mk_static_str()),
             sym::type_id => (1, Vec::new(), tcx.types.u64),
+            sym::option_some_offset => (1, Vec::new(), tcx.types.usize),
             sym::offset | sym::arith_offset => (
                 1,
                 vec![

compiler/rustc_span/src/symbol.rs

@@ -1044,6 +1044,7 @@ symbols! {
         optin_builtin_traits,
         option,
         option_env,
+        option_some_offset,
         options,
         or,
         or_patterns,

library/core/src/intrinsics.rs

@@ -1303,6 +1303,16 @@ extern "rust-intrinsic" {
     #[rustc_const_stable(feature = "const_ptr_offset", since = "1.61.0")]
     pub fn arith_offset<T>(dst: *const T, offset: isize) -> *const T;
 
+    #[cfg(not(bootstrap))]
+    /// The offset of the `Some(_)` value of an `Option<T>`. Used internally for
+    /// `Option::as_slice` and `Option::as_mut_slice`. This needs to be called with an
+    /// `Option<_>` type.
+    // FIXME: This should be replaced once we get a stable public method for getting
+    // the offset of enum variant's fields (e.g. an extension of `offset_of!` to enums)
+    #[rustc_const_stable(feature = "const_option_some_offset", since = "1.68.0")]
+    #[rustc_safe_intrinsic]
+    pub fn option_some_offset<T>() -> usize;
+
     /// Masks out bits of the pointer according to a mask.
     ///
     /// Note that, unlike most intrinsics, this is safe to call;

library/core/src/option.rs

@@ -736,46 +736,15 @@ impl<T> Option<T> {
     }
 
     /// 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 {
-            // 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
-        }
-    };
+    /// As this version will only be ever used to compile rustc and the performance
+    /// penalty is negligible, use a minimal implementation here.
+    #[cfg(bootstrap)]
+    const SOME_BYTE_OFFSET_GUESS: usize = 0;
+
+    // FIXME: replace this with whatever stable method to get the offset of an enum
+    // field may appear first.
+    #[cfg(not(bootstrap))]
+    const SOME_BYTE_OFFSET_GUESS: usize = crate::intrinsics::option_some_offset::<Option<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
@@ -820,7 +789,7 @@ impl<T> Option<T> {
                 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() }
+                unsafe { self_ptr.byte_add(Self::SOME_BYTE_OFFSET_GUESS).cast() }
             };
         let len = usize::from(self.is_some());
 
@@ -886,7 +855,7 @@ impl<T> Option<T> {
                 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() }
+                unsafe { self_ptr.byte_add(Self::SOME_BYTE_OFFSET_GUESS).cast() }
             };
         let len = usize::from(self.is_some());