safe transmute: support `Variants::Single` enums

Previously, the implementation of `Tree::from_enum` incorrectly
treated enums with `Variants::Single` and `Variants::Multiple`
identically. This is incorrect for `Variants::Single` enums,
which delegate their layout to that of a variant with a particular
index (or no variant at all if the enum is empty).

This flaw manifested first as an ICE. `Tree::from_enum` attempted
to compute the tag of variants other than the one at
`Variants::Single`'s `index`, and fell afoul of a sanity-checking
assertion in `compiler/rustc_const_eval/src/interpret/discriminant.rs`.
This assertion is non-load-bearing, and can be removed; the routine
its in is well-behaved even without it.

With the assertion removed, the proximate issue becomes apparent:
calling `Tree::from_variant` on a variant that does not exist is
ill-defined. A sanity check the given variant has
`FieldShapes::Arbitrary` fails, and the analysis is (correctly)
aborted with `Err::NotYetSupported`.

This commit corrects this chain of failures by ensuring that
`Tree::from_variant` is not called on variants that are, as far as
layout is concerned, nonexistent. Specifically, the implementation
of `Tree::from_enum` is now partitioned into three cases:

  1. enums that are uninhabited
  2. enums for which all but one variant is uninhabited
  3. enums with multiple inhabited variants

`Tree::from_variant` is now only invoked in the third case. In the
first case, `Tree::uninhabited()` is produced. In the second case,
the layout is delegated to `Variants::Single`'s index.

Fixes #125811

4 files changed, 81 insertions, 59 deletions

diff --git a/compiler/rustc_const_eval/src/interpret/discriminant.rs b/compiler/rustc_const_eval/src/interpret/discriminant.rs
index 0dbee8c1d94..a50b50d231d 100644
--- a/compiler/rustc_const_eval/src/interpret/discriminant.rs
+++ b/compiler/rustc_const_eval/src/interpret/discriminant.rs
@@ -241,10 +241,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
         variant_index: VariantIdx,
     ) -> InterpResult<'tcx, Option<(ScalarInt, usize)>> {
         match self.layout_of(ty)?.variants {
-            abi::Variants::Single { index } => {
-                assert_eq!(index, variant_index);
-                Ok(None)
-            }
+            abi::Variants::Single { .. } => Ok(None),
 
             abi::Variants::Multiple {
                 tag_encoding: TagEncoding::Direct,
diff --git a/compiler/rustc_transmute/src/layout/tree.rs b/compiler/rustc_transmute/src/layout/tree.rs
index eae1a9dfaa2..241381f5875 100644
--- a/compiler/rustc_transmute/src/layout/tree.rs
+++ b/compiler/rustc_transmute/src/layout/tree.rs
@@ -331,30 +331,63 @@ pub(crate) mod rustc {
             assert!(def.is_enum());
             let layout = ty_and_layout.layout;
 
-            if let Variants::Multiple { tag_field, .. } = layout.variants() {
-                // For enums (but not coroutines), the tag field is
-                // currently always the first field of the layout.
-                assert_eq!(*tag_field, 0);
-            }
+            // Computes the variant of a given index.
+            let layout_of_variant = |index| {
+                let tag = cx.tcx.tag_for_variant((ty_and_layout.ty, index));
+                let variant_def = Def::Variant(def.variant(index));
+                let variant_ty_and_layout = ty_and_layout.for_variant(&cx, index);
+                Self::from_variant(variant_def, tag, variant_ty_and_layout, layout.size, cx)
+            };
 
-            let variants = def.discriminants(cx.tcx()).try_fold(
-                Self::uninhabited(),
-                |variants, (idx, ref discriminant)| {
-                    let tag = cx.tcx.tag_for_variant((ty_and_layout.ty, idx));
-                    let variant_def = Def::Variant(def.variant(idx));
-                    let variant_ty_and_layout = ty_and_layout.for_variant(&cx, idx);
-                    let variant = Self::from_variant(
-                        variant_def,
-                        tag,
-                        variant_ty_and_layout,
-                        layout.size,
-                        cx,
+            // We consider three kinds of enums, each demanding a different
+            // treatment of their layout computation:
+            // 1. enums that are uninhabited
+            // 2. enums for which all but one variant is uninhabited
+            // 3. enums with multiple inhabited variants
+            match layout.variants() {
+                _ if layout.abi.is_uninhabited() => {
+                    // Uninhabited enums are usually (always?) zero-sized. In
+                    // the (unlikely?) event that an uninhabited enum is
+                    // non-zero-sized, this assert will trigger an ICE, and this
+                    // code should be modified such that a `layout.size` amount
+                    // of uninhabited bytes is returned instead.
+                    //
+                    // Uninhabited enums are currently implemented such that
+                    // their layout is described with `Variants::Single`, even
+                    // though they don't necessarily have a 'single' variant to
+                    // defer to. That said, we don't bother specifically
+                    // matching on `Variants::Single` in this arm because the
+                    // behavioral principles here remain true even if, for
+                    // whatever reason, the compiler describes an uninhabited
+                    // enum with `Variants::Multiple`.
+                    assert_eq!(layout.size, Size::ZERO);
+                    Ok(Self::uninhabited())
+                }
+                Variants::Single { index } => {
+                    // `Variants::Single` on non-uninhabited enums denotes that
+                    // the enum delegates its layout to the variant at `index`.
+                    layout_of_variant(*index)
+                }
+                Variants::Multiple { tag_field, .. } => {
+                    // `Variants::Multiple` denotes an enum with multiple
+                    // inhabited variants. The layout of such an enum is the
+                    // disjunction of the layouts of its tagged variants.
+
+                    // For enums (but not coroutines), the tag field is
+                    // currently always the first field of the layout.
+                    assert_eq!(*tag_field, 0);
+
+                    let variants = def.discriminants(cx.tcx()).try_fold(
+                        Self::uninhabited(),
+                        |variants, (idx, ref discriminant)| {
+                            let variant = layout_of_variant(idx)?;
+                            Result::<Self, Err>::Ok(variants.or(variant))
+                        },
                     )?;
-                    Result::<Self, Err>::Ok(variants.or(variant))
-                },
-            )?;
 
-            return Ok(Self::def(Def::Adt(def)).then(variants));
+                    return Ok(Self::def(Def::Adt(def)).then(variants));
+                }
+            }
         }
 
         /// Constructs a `Tree` from a 'variant-like' layout.
diff --git a/tests/crashes/125811.rs b/tests/crashes/125811.rs
deleted file mode 100644
index eb764e8d152..00000000000
--- a/tests/crashes/125811.rs
+++ /dev/null
@@ -1,34 +0,0 @@
-//@ known-bug: rust-lang/rust#125811
-mod assert {
-    use std::mem::{Assume, BikeshedIntrinsicFrom};
-
-    pub fn is_transmutable<Src, Dst>()
-    where
-        Dst: BikeshedIntrinsicFrom<Src>,
-    {
-    }
-}
-
-#[repr(C)]
-struct Zst;
-
-enum V0 {
-    B(!),
-}
-
-enum V2 {
-    V = 2,
-}
-
-enum Lopsided {
-    Smol(Zst),
-    Lorg(V0),
-}
-
-#[repr(C)]
-#[repr(C)]
-struct Dst(Lopsided, V2);
-
-fn should_pad_variants() {
-    assert::is_transmutable::<Src, Dst>();
-}
diff --git a/tests/ui/transmutability/enums/uninhabited_optimization.rs b/tests/ui/transmutability/enums/uninhabited_optimization.rs
new file mode 100644
index 00000000000..04a8eb40c8b
--- /dev/null
+++ b/tests/ui/transmutability/enums/uninhabited_optimization.rs
@@ -0,0 +1,26 @@
+//@ check-pass
+//! Tests that we do not regress rust-lang/rust#125811
+#![feature(transmutability)]
+
+fn assert_transmutable<T>()
+where
+    (): std::mem::BikeshedIntrinsicFrom<T>
+{}
+
+enum Uninhabited {}
+
+enum SingleInhabited {
+    X,
+    Y(Uninhabited)
+}
+
+enum SingleUninhabited {
+    X(Uninhabited),
+    Y(Uninhabited),
+}
+
+fn main() {
+    assert_transmutable::<Uninhabited>();
+    assert_transmutable::<SingleInhabited>();
+    assert_transmutable::<SingleUninhabited>();
+}