Rollup merge of #133681 - RalfJung:niches, r=wesleywiser

improve TagEncoding::Niche docs, sanity check, and UB checks

Turns out the `niche_variants` range can actually contain the `untagged_variant`. We should report this as UB in Miri, so this PR implements that.

Also rename `partially_check_layout` to `layout_sanity_check` for better consistency with how similar functions are called in other parts of the compiler.

Turns out my adjustments to the transmutation logic also fix https://github.com/rust-lang/rust/issues/126267.

2 files changed, 14 insertions, 4 deletions

diff --git a/compiler/rustc_ty_utils/src/layout.rs b/compiler/rustc_ty_utils/src/layout.rs
index 66134b81b2a..0d656f1b63b 100644
--- a/compiler/rustc_ty_utils/src/layout.rs
+++ b/compiler/rustc_ty_utils/src/layout.rs
@@ -81,7 +81,7 @@ fn layout_of<'tcx>(
         record_layout_for_printing(&cx, layout);
     }
 
-    invariant::partially_check_layout(&cx, &layout);
+    invariant::layout_sanity_check(&cx, &layout);
 
     Ok(layout)
 }
diff --git a/compiler/rustc_ty_utils/src/layout/invariant.rs b/compiler/rustc_ty_utils/src/layout/invariant.rs
index 26ea81daf78..f39b87622f4 100644
--- a/compiler/rustc_ty_utils/src/layout/invariant.rs
+++ b/compiler/rustc_ty_utils/src/layout/invariant.rs
@@ -1,11 +1,11 @@
 use std::assert_matches::assert_matches;
 
-use rustc_abi::{BackendRepr, FieldsShape, Scalar, Size, Variants};
+use rustc_abi::{BackendRepr, FieldsShape, Scalar, Size, TagEncoding, Variants};
 use rustc_middle::bug;
 use rustc_middle::ty::layout::{HasTyCtxt, LayoutCx, TyAndLayout};
 
 /// Enforce some basic invariants on layouts.
-pub(super) fn partially_check_layout<'tcx>(cx: &LayoutCx<'tcx>, layout: &TyAndLayout<'tcx>) {
+pub(super) fn layout_sanity_check<'tcx>(cx: &LayoutCx<'tcx>, layout: &TyAndLayout<'tcx>) {
     let tcx = cx.tcx();
 
     // Type-level uninhabitedness should always imply ABI uninhabitedness.
@@ -241,7 +241,17 @@ pub(super) fn partially_check_layout<'tcx>(cx: &LayoutCx<'tcx>, layout: &TyAndLa
 
     check_layout_abi(cx, layout);
 
-    if let Variants::Multiple { variants, .. } = &layout.variants {
+    if let Variants::Multiple { variants, tag, tag_encoding, .. } = &layout.variants {
+        if let TagEncoding::Niche { niche_start, untagged_variant, niche_variants } = tag_encoding {
+            let niche_size = tag.size(cx);
+            assert!(*niche_start <= niche_size.unsigned_int_max());
+            for (idx, variant) in variants.iter_enumerated() {
+                // Ensure all inhabited variants are accounted for.
+                if !variant.is_uninhabited() {
+                    assert!(idx == *untagged_variant || niche_variants.contains(&idx));
+                }
+            }
+        }
         for variant in variants.iter() {
             // No nested "multiple".
             assert_matches!(variant.variants, Variants::Single { .. });