report UB when the niche value refers to the untagged variant

4 files changed, 29 insertions, 21 deletions

diff --git a/compiler/rustc_const_eval/src/const_eval/mod.rs b/compiler/rustc_const_eval/src/const_eval/mod.rs
index 8cbdcd68e13..34f795bda75 100644
--- a/compiler/rustc_const_eval/src/const_eval/mod.rs
+++ b/compiler/rustc_const_eval/src/const_eval/mod.rs
@@ -2,6 +2,7 @@
 
 use rustc_abi::VariantIdx;
 use rustc_middle::query::{Key, TyCtxtAt};
+use rustc_middle::ty::layout::LayoutOf;
 use rustc_middle::ty::{self, Ty, TyCtxt};
 use rustc_middle::{bug, mir};
 use tracing::instrument;
@@ -85,5 +86,6 @@ pub fn tag_for_variant_provider<'tcx>(
         crate::const_eval::DummyMachine,
     );
 
-    ecx.tag_for_variant(ty, variant_index).unwrap().map(|(tag, _tag_field)| tag)
+    let layout = ecx.layout_of(ty).unwrap();
+    ecx.tag_for_variant(layout, variant_index).unwrap().map(|(tag, _tag_field)| tag)
 }
diff --git a/compiler/rustc_const_eval/src/interpret/discriminant.rs b/compiler/rustc_const_eval/src/interpret/discriminant.rs
index f94d0cbb42b..c7c8a2902e2 100644
--- a/compiler/rustc_const_eval/src/interpret/discriminant.rs
+++ b/compiler/rustc_const_eval/src/interpret/discriminant.rs
@@ -1,7 +1,7 @@
 //! Functions for reading and writing discriminants of multi-variant layouts (enums and coroutines).
 
 use rustc_abi::{self as abi, TagEncoding, VariantIdx, Variants};
-use rustc_middle::ty::layout::{LayoutOf, PrimitiveExt};
+use rustc_middle::ty::layout::{LayoutOf, PrimitiveExt, TyAndLayout};
 use rustc_middle::ty::{self, CoroutineArgsExt, ScalarInt, Ty};
 use rustc_middle::{mir, span_bug};
 use tracing::{instrument, trace};
@@ -21,17 +21,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
         variant_index: VariantIdx,
         dest: &impl Writeable<'tcx, M::Provenance>,
     ) -> InterpResult<'tcx> {
-        // Layout computation excludes uninhabited variants from consideration
-        // therefore there's no way to represent those variants in the given layout.
-        // Essentially, uninhabited variants do not have a tag that corresponds to their
-        // discriminant, so we cannot do anything here.
-        // When evaluating we will always error before even getting here, but ConstProp 'executes'
-        // dead code, so we cannot ICE here.
-        if dest.layout().for_variant(self, variant_index).is_uninhabited() {
-            throw_ub!(UninhabitedEnumVariantWritten(variant_index))
-        }
-
-        match self.tag_for_variant(dest.layout().ty, variant_index)? {
+        match self.tag_for_variant(dest.layout(), variant_index)? {
             Some((tag, tag_field)) => {
                 // No need to validate that the discriminant here because the
                 // `TyAndLayout::for_variant()` call earlier already checks the
@@ -188,6 +178,11 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
                             let variants =
                                 ty.ty_adt_def().expect("tagged layout for non adt").variants();
                             assert!(variant_index < variants.next_index());
+                            if variant_index == untagged_variant {
+                                // The untagged variant can be in the niche range, but even then it
+                                // is not a valid encoding.
+                                throw_ub!(InvalidTag(Scalar::from_uint(tag_bits, tag_layout.size)))
+                            }
                             variant_index
                         } else {
                             untagged_variant
@@ -236,10 +231,18 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
     ///   given field index.
     pub(crate) fn tag_for_variant(
         &self,
-        ty: Ty<'tcx>,
+        layout: TyAndLayout<'tcx>,
         variant_index: VariantIdx,
     ) -> InterpResult<'tcx, Option<(ScalarInt, usize)>> {
-        match self.layout_of(ty)?.variants {
+        // Layout computation excludes uninhabited variants from consideration.
+        // Therefore, there's no way to represent those variants in the given layout.
+        // Essentially, uninhabited variants do not have a tag that corresponds to their
+        // discriminant, so we have to bail out here.
+        if layout.for_variant(self, variant_index).is_uninhabited() {
+            throw_ub!(UninhabitedEnumVariantWritten(variant_index))
+        }
+
+        match layout.variants {
             abi::Variants::Single { .. } => {
                 // The tag of a `Single` enum is like the tag of the niched
                 // variant: there's no tag as the discriminant is encoded
@@ -260,7 +263,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
                 // raw discriminants for enums are isize or bigger during
                 // their computation, but the in-memory tag is the smallest possible
                 // representation
-                let discr = self.discriminant_for_variant(ty, variant_index)?;
+                let discr = self.discriminant_for_variant(layout.ty, variant_index)?;
                 let discr_size = discr.layout.size;
                 let discr_val = discr.to_scalar().to_bits(discr_size)?;
                 let tag_size = tag_layout.size(self);
@@ -286,11 +289,13 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
                 ..
             } => {
                 assert!(variant_index != untagged_variant);
+                // We checked that this variant is inhabited, so it must be in the niche range.
+                assert!(
+                    niche_variants.contains(&variant_index),
+                    "invalid variant index for this enum"
+                );
                 let variants_start = niche_variants.start().as_u32();
-                let variant_index_relative = variant_index
-                    .as_u32()
-                    .checked_sub(variants_start)
-                    .expect("overflow computing relative variant idx");
+                let variant_index_relative = variant_index.as_u32().strict_sub(variants_start);
                 // We need to use machine arithmetic when taking into account `niche_start`:
                 // tag_val = variant_index_relative + niche_start_val
                 let tag_layout = self.layout_of(tag_layout.primitive().to_int_ty(*self.tcx))?;
diff --git a/compiler/rustc_const_eval/src/lib.rs b/compiler/rustc_const_eval/src/lib.rs
index 2a7408f1c70..b5adf06b300 100644
--- a/compiler/rustc_const_eval/src/lib.rs
+++ b/compiler/rustc_const_eval/src/lib.rs
@@ -10,6 +10,7 @@
 #![feature(never_type)]
 #![feature(rustdoc_internals)]
 #![feature(slice_ptr_get)]
+#![feature(strict_overflow_ops)]
 #![feature(trait_alias)]
 #![feature(try_blocks)]
 #![feature(unqualified_local_imports)]
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)
 }