Track (partial) niche information in `NaiveLayout`

Still more complexity, but this allows computing exact `NaiveLayout`s
for null-optimized enums, and thus allows calls like
`transmute::<Option<&T>, &U>()` to work in generic contexts.

5 files changed, 185 insertions, 47 deletions

diff --git a/compiler/rustc_abi/src/lib.rs b/compiler/rustc_abi/src/lib.rs
index be1db7311de..d396f18d59c 100644
--- a/compiler/rustc_abi/src/lib.rs
+++ b/compiler/rustc_abi/src/lib.rs
@@ -1062,9 +1062,15 @@ impl WrappingRange {
     /// Returns `true` if `size` completely fills the range.
     #[inline]
     pub fn is_full_for(&self, size: Size) -> bool {
+        debug_assert!(self.is_in_range_for(size));
+        self.start == (self.end.wrapping_add(1) & size.unsigned_int_max())
+    }
+
+    /// Returns `true` if the range is valid for `size`.
+    #[inline(always)]
+    pub fn is_in_range_for(&self, size: Size) -> bool {
         let max_value = size.unsigned_int_max();
-        debug_assert!(self.start <= max_value && self.end <= max_value);
-        self.start == (self.end.wrapping_add(1) & max_value)
+        self.start <= max_value && self.end <= max_value
     }
 }
 
diff --git a/compiler/rustc_middle/src/ty/layout.rs b/compiler/rustc_middle/src/ty/layout.rs
index 34a345aba04..26137e86fa0 100644
--- a/compiler/rustc_middle/src/ty/layout.rs
+++ b/compiler/rustc_middle/src/ty/layout.rs
@@ -655,6 +655,8 @@ impl std::ops::DerefMut for TyAndNaiveLayout<'_> {
 #[derive(Copy, Clone, Debug, HashStable)]
 pub struct NaiveLayout {
     pub abi: NaiveAbi,
+    /// Niche information, required for tracking non-null enum optimizations.
+    pub niches: NaiveNiches,
     /// An underestimate of the layout's size.
     pub size: Size,
     /// An underestimate of the layout's required alignment.
@@ -664,12 +666,19 @@ pub struct NaiveLayout {
 }
 
 #[derive(Copy, Clone, Debug, Eq, PartialEq, HashStable)]
+pub enum NaiveNiches {
+    None,
+    Some,
+    Maybe,
+}
+
+#[derive(Copy, Clone, Debug, Eq, PartialEq, HashStable)]
 pub enum NaiveAbi {
-    /// A scalar layout, always implies `exact`.
+    /// A scalar layout, always implies `exact` and a non-zero `size`.
     Scalar(Primitive),
-    /// An uninhabited layout. (needed to properly track `Scalar`)
+    /// An uninhabited layout. (needed to properly track `Scalar` and niches)
     Uninhabited,
-    /// An unsized aggregate. (needed to properly track `Scalar`)
+    /// An unsized aggregate. (needed to properly track `Scalar` and niches)
     Unsized,
     /// Any other sized layout.
     Sized,
@@ -687,8 +696,13 @@ impl NaiveAbi {
 
 impl NaiveLayout {
     /// The layout of an empty aggregate, e.g. `()`.
-    pub const EMPTY: Self =
-        Self { size: Size::ZERO, align: Align::ONE, exact: true, abi: NaiveAbi::Sized };
+    pub const EMPTY: Self = Self {
+        size: Size::ZERO,
+        align: Align::ONE,
+        exact: true,
+        abi: NaiveAbi::Sized,
+        niches: NaiveNiches::None,
+    };
 
     /// Returns whether `self` is a valid approximation of the given full `layout`.
     ///
@@ -699,12 +713,20 @@ impl NaiveLayout {
         }
 
         if let NaiveAbi::Scalar(prim) = self.abi {
-            assert!(self.exact);
-            if !matches!(layout.abi(), Abi::Scalar(s) if s.primitive() == prim) {
+            if !self.exact
+                || self.size == Size::ZERO
+                || !matches!(layout.abi(), Abi::Scalar(s) if s.primitive() == prim)
+            {
                 return false;
             }
         }
 
+        match (self.niches, layout.largest_niche()) {
+            (NaiveNiches::None, Some(_)) => return false,
+            (NaiveNiches::Some, None) => return false,
+            _ => (),
+        }
+
         !self.exact || (self.size, self.align) == (layout.size(), layout.align().abi)
     }
 
@@ -745,6 +767,15 @@ impl NaiveLayout {
         self
     }
 
+    /// Artificially makes this layout inexact.
+    #[must_use]
+    #[inline]
+    pub fn inexact(mut self) -> Self {
+        self.abi = self.abi.as_aggregate();
+        self.exact = false;
+        self
+    }
+
     /// Pads this layout so that its size is a multiple of `align`.
     #[must_use]
     #[inline]
@@ -777,11 +808,18 @@ impl NaiveLayout {
             // Default case.
             (_, _) => Sized,
         };
-        Some(Self { abi, size, align, exact })
+        let niches = match (self.niches, other.niches) {
+            (NaiveNiches::Some, _) | (_, NaiveNiches::Some) => NaiveNiches::Some,
+            (NaiveNiches::None, NaiveNiches::None) => NaiveNiches::None,
+            (_, _) => NaiveNiches::Maybe,
+        };
+        Some(Self { abi, size, align, exact, niches })
     }
 
     /// Returns the layout of `self` superposed with `other`, as in an `enum`
     /// or an `union`.
+    ///
+    /// Note: This always ignore niche information from `other`.
     #[must_use]
     #[inline]
     pub fn union(&self, other: &Self) -> Self {
@@ -793,7 +831,7 @@ impl NaiveLayout {
         let abi = match (self.abi, other.abi) {
             // The unsized ABI overrides everything.
             (Unsized, _) | (_, Unsized) => Unsized,
-            // A scalar union must have a single non ZST-field.
+            // A scalar union must have a single non ZST-field...
             (_, s @ Scalar(_)) if exact && self.size == Size::ZERO => s,
             (s @ Scalar(_), _) if exact && other.size == Size::ZERO => s,
             // ...or identical scalar fields.
@@ -802,7 +840,7 @@ impl NaiveLayout {
             (Uninhabited, Uninhabited) => Uninhabited,
             (_, _) => Sized,
         };
-        Self { abi, size, align, exact }
+        Self { abi, size, align, exact, niches: self.niches }
     }
 }
 
diff --git a/compiler/rustc_ty_utils/src/layout_naive.rs b/compiler/rustc_ty_utils/src/layout_naive.rs
index 501b27777fa..3070ab59d53 100644
--- a/compiler/rustc_ty_utils/src/layout_naive.rs
+++ b/compiler/rustc_ty_utils/src/layout_naive.rs
@@ -1,12 +1,14 @@
 use rustc_middle::query::Providers;
 use rustc_middle::ty::layout::{
-    IntegerExt, LayoutCx, LayoutError, LayoutOf, NaiveAbi, NaiveLayout, TyAndNaiveLayout,
+    IntegerExt, LayoutCx, LayoutError, LayoutOf, NaiveAbi, NaiveLayout, NaiveNiches,
+    TyAndNaiveLayout,
 };
 use rustc_middle::ty::{self, ReprOptions, Ty, TyCtxt, TypeVisitableExt};
-
 use rustc_span::DUMMY_SP;
 use rustc_target::abi::*;
 
+use std::ops::Bound;
+
 use crate::layout::{compute_array_count, ptr_metadata_scalar};
 
 pub fn provide(providers: &mut Providers) {
@@ -61,8 +63,9 @@ fn naive_layout_of_uncached<'tcx>(
     let tcx = cx.tcx;
     let dl = cx.data_layout();
 
-    let scalar = |value: Primitive| NaiveLayout {
+    let scalar = |niched: bool, value: Primitive| NaiveLayout {
         abi: NaiveAbi::Scalar(value),
+        niches: if niched { NaiveNiches::Some } else { NaiveNiches::None },
         size: value.size(dl),
         align: value.align(dl).abi,
         exact: true,
@@ -105,26 +108,30 @@ fn naive_layout_of_uncached<'tcx>(
 
     Ok(match *ty.kind() {
         // Basic scalars
-        ty::Bool => scalar(Int(I8, false)),
-        ty::Char => scalar(Int(I32, false)),
-        ty::Int(ity) => scalar(Int(Integer::from_int_ty(dl, ity), true)),
-        ty::Uint(ity) => scalar(Int(Integer::from_uint_ty(dl, ity), false)),
-        ty::Float(fty) => scalar(match fty {
-            ty::FloatTy::F32 => F32,
-            ty::FloatTy::F64 => F64,
-        }),
-        ty::FnPtr(_) => scalar(Pointer(dl.instruction_address_space)),
+        ty::Bool => scalar(true, Int(I8, false)),
+        ty::Char => scalar(true, Int(I32, false)),
+        ty::Int(ity) => scalar(false, Int(Integer::from_int_ty(dl, ity), true)),
+        ty::Uint(ity) => scalar(false, Int(Integer::from_uint_ty(dl, ity), false)),
+        ty::Float(fty) => scalar(
+            false,
+            match fty {
+                ty::FloatTy::F32 => F32,
+                ty::FloatTy::F64 => F64,
+            },
+        ),
+        ty::FnPtr(_) => scalar(true, Pointer(dl.instruction_address_space)),
 
         // The never type.
         ty::Never => NaiveLayout { abi: NaiveAbi::Uninhabited, ..NaiveLayout::EMPTY },
 
         // Potentially-wide pointers.
         ty::Ref(_, pointee, _) | ty::RawPtr(ty::TypeAndMut { ty: pointee, .. }) => {
-            let data_ptr = scalar(Pointer(AddressSpace::DATA));
+            let data_ptr = scalar(!ty.is_unsafe_ptr(), Pointer(AddressSpace::DATA));
             if let Some(metadata) = ptr_metadata_scalar(cx, pointee)? {
                 // Effectively a (ptr, meta) tuple.
+                let meta = scalar(!metadata.is_always_valid(dl), metadata.primitive());
                 let l = data_ptr
-                    .concat(&scalar(metadata.primitive()), dl)
+                    .concat(&meta, dl)
                     .ok_or_else(|| error(cx, LayoutError::SizeOverflow(ty)))?;
                 l.pad_to_align(l.align)
             } else {
@@ -134,8 +141,9 @@ fn naive_layout_of_uncached<'tcx>(
         }
 
         ty::Dynamic(_, _, ty::DynStar) => {
-            let ptr = scalar(Pointer(AddressSpace::DATA));
-            ptr.concat(&ptr, dl).ok_or_else(|| error(cx, LayoutError::SizeOverflow(ty)))?
+            let ptr = scalar(false, Pointer(AddressSpace::DATA));
+            let vtable = scalar(true, Pointer(AddressSpace::DATA));
+            ptr.concat(&vtable, dl).ok_or_else(|| error(cx, LayoutError::SizeOverflow(ty)))?
         }
 
         // Arrays and slices.
@@ -149,13 +157,16 @@ fn naive_layout_of_uncached<'tcx>(
                     .size
                     .checked_mul(count, cx)
                     .ok_or_else(|| error(cx, LayoutError::SizeOverflow(ty)))?,
+                niches: if count == 0 { NaiveNiches::None } else { element.niches },
                 ..*element
             }
         }
-        ty::Slice(element) => {
-            let element = cx.naive_layout_of(element)?;
-            NaiveLayout { abi: NaiveAbi::Unsized, size: Size::ZERO, ..*element }
-        }
+        ty::Slice(element) => NaiveLayout {
+            abi: NaiveAbi::Unsized,
+            size: Size::ZERO,
+            niches: NaiveNiches::None,
+            ..*cx.naive_layout_of(element)?
+        },
 
         ty::FnDef(..) => NaiveLayout::EMPTY,
 
@@ -166,7 +177,9 @@ fn naive_layout_of_uncached<'tcx>(
 
         // FIXME(reference_niches): try to actually compute a reasonable layout estimate,
         // without duplicating too much code from `generator_layout`.
-        ty::Generator(..) => NaiveLayout { exact: false, ..NaiveLayout::EMPTY },
+        ty::Generator(..) => {
+            NaiveLayout { exact: false, niches: NaiveNiches::Maybe, ..NaiveLayout::EMPTY }
+        }
 
         ty::Closure(_, ref substs) => {
             univariant(&mut substs.as_closure().upvar_tys(), &ReprOptions::default())?
@@ -175,6 +188,7 @@ fn naive_layout_of_uncached<'tcx>(
         ty::Tuple(tys) => univariant(&mut tys.iter(), &ReprOptions::default())?,
 
         ty::Adt(def, substs) if def.is_union() => {
+            assert_eq!(def.variants().len(), 1, "union should have a single variant");
             let repr = def.repr();
             let pack = repr.pack.unwrap_or(Align::MAX);
             if repr.pack.is_some() && repr.align.is_some() {
@@ -182,7 +196,12 @@ fn naive_layout_of_uncached<'tcx>(
                 return Err(error(cx, LayoutError::Unknown(ty)));
             }
 
-            let mut layout = NaiveLayout::EMPTY;
+            let mut layout = NaiveLayout {
+                // Unions never have niches.
+                niches: NaiveNiches::None,
+                ..NaiveLayout::EMPTY
+            };
+
             for f in &def.variants()[FIRST_VARIANT].fields {
                 let field = cx.naive_layout_of(f.ty(tcx, substs))?;
                 layout = layout.union(&field.packed(pack));
@@ -201,24 +220,87 @@ fn naive_layout_of_uncached<'tcx>(
 
         ty::Adt(def, substs) => {
             let repr = def.repr();
-            let base = NaiveLayout {
-                // For simplicity, assume that any enum has its discriminant field (if it exists)
-                // niched inside one of the variants; this will underestimate the size (and sometimes
-                // alignment) of enums. We also doesn't compute exact alignment for SIMD structs.
-                // FIXME(reference_niches): Be smarter here.
-                // Also consider adding a special case for null-optimized enums, so that we can have
-                // `Option<&T>: PointerLike` in generic contexts.
-                exact: !def.is_enum() && !repr.simd(),
+            let mut layout = NaiveLayout {
                 // An ADT with no inhabited variants should have an uninhabited ABI.
                 abi: NaiveAbi::Uninhabited,
                 ..NaiveLayout::EMPTY
             };
 
-            let layout = def.variants().iter().try_fold(base, |layout, v| {
+            let mut empty_variants = 0;
+            for v in def.variants() {
                 let mut fields = v.fields.iter().map(|f| f.ty(tcx, substs));
                 let vlayout = univariant(&mut fields, &repr)?;
-                Ok(layout.union(&vlayout))
-            })?;
+
+                if vlayout.size == Size::ZERO && vlayout.exact {
+                    empty_variants += 1;
+                } else {
+                    // Remember the niches of the last seen variant.
+                    layout.niches = vlayout.niches;
+                }
+
+                layout = layout.union(&vlayout);
+            }
+
+            if def.is_enum() {
+                let may_need_discr = match def.variants().len() {
+                    0 | 1 => false,
+                    // Simple Option-like niche optimization.
+                    // Handling this special case allows enums like `Option<&T>`
+                    // to be recognized as `PointerLike` and to be transmutable
+                    // in generic contexts.
+                    2 if empty_variants == 1 && layout.niches == NaiveNiches::Some => {
+                        layout.niches = NaiveNiches::Maybe; // fill up the niche.
+                        false
+                    }
+                    _ => true,
+                };
+
+                if may_need_discr || repr.inhibit_enum_layout_opt() {
+                    // For simplicity, assume that the discriminant always get niched.
+                    // This will be wrong in many cases, which will cause the size (and
+                    // sometimes the alignment) to be underestimated.
+                    // FIXME(reference_niches): Be smarter here.
+                    layout.niches = NaiveNiches::Maybe;
+                    layout = layout.inexact();
+                }
+            } else {
+                assert_eq!(def.variants().len(), 1, "struct should have a single variant");
+
+                // We don't compute exact alignment for SIMD structs.
+                if repr.simd() {
+                    layout = layout.inexact();
+                }
+
+                // `UnsafeCell` hides all niches.
+                if def.is_unsafe_cell() {
+                    layout.niches = NaiveNiches::None;
+                }
+            }
+
+            let valid_range = tcx.layout_scalar_valid_range(def.did());
+            if valid_range != (Bound::Unbounded, Bound::Unbounded) {
+                let get = |bound, default| match bound {
+                    Bound::Unbounded => default,
+                    Bound::Included(v) => v,
+                    Bound::Excluded(_) => bug!("exclusive `layout_scalar_valid_range` bound"),
+                };
+
+                let valid_range = WrappingRange {
+                    start: get(valid_range.0, 0),
+                    // FIXME: this is wrong for scalar-pair ABIs. Fortunately, the
+                    // only type this could currently affect is`NonNull<T: !Sized>`,
+                    // and the `NaiveNiches` result still ends up correct.
+                    end: get(valid_range.1, layout.size.unsigned_int_max()),
+                };
+                assert!(
+                    valid_range.is_in_range_for(layout.size),
+                    "`layout_scalar_valid_range` values are out of bounds",
+                );
+                if !valid_range.is_full_for(layout.size) {
+                    layout.niches = NaiveNiches::Some;
+                }
+            }
+
             layout.pad_to_align(layout.align)
         }
 
diff --git a/tests/ui/layout/valid_range_oob.stderr b/tests/ui/layout/valid_range_oob.stderr
index a3a514fb830..772113fa5fb 100644
--- a/tests/ui/layout/valid_range_oob.stderr
+++ b/tests/ui/layout/valid_range_oob.stderr
@@ -1,6 +1,6 @@
 error: the compiler unexpectedly panicked. this is a bug.
 
 query stack during panic:
-#0 [layout_of] computing layout of `Foo`
-#1 [eval_to_allocation_raw] const-evaluating + checking `FOO`
+#0 [naive_layout_of] computing layout (naive) of `Foo`
+#1 [layout_of] computing layout of `Foo`
 end of query stack
diff --git a/tests/ui/transmute/transmute-fat-pointers.rs b/tests/ui/transmute/transmute-fat-pointers.rs
index 7c1beffd14e..d373ff5f24a 100644
--- a/tests/ui/transmute/transmute-fat-pointers.rs
+++ b/tests/ui/transmute/transmute-fat-pointers.rs
@@ -30,4 +30,16 @@ fn f<T, U: ?Sized>(x: &T) -> &U {
     unsafe { transmute(x) } //~ ERROR cannot transmute between types of different sizes
 }
 
+fn g<T, U>(x: &T) -> Option<&U> {
+    unsafe { transmute(x) }
+}
+
+fn h<T>(x: &[T]) -> Option<&dyn Send> {
+    unsafe { transmute(x) }
+}
+
+fn i<T>(x: [usize; 1]) -> Option<&'static T> {
+    unsafe { transmute(x) }
+}
+
 fn main() { }