interpret/write_discriminant: when encoding niched variant, ensure the stored value matches

compiler/rustc_const_eval/messages.ftl

@@ -170,6 +170,10 @@ const_eval_invalid_meta =
     invalid metadata in wide pointer: total size is bigger than largest supported object
 const_eval_invalid_meta_slice =
     invalid metadata in wide pointer: slice is bigger than largest supported object
+
+const_eval_invalid_niched_enum_variant_written =
+    trying to set discriminant of a {$ty} to the niched variant, but the value does not match
+
 const_eval_invalid_str =
     this string is not valid UTF-8: {$err}
 const_eval_invalid_tag =

compiler/rustc_const_eval/src/errors.rs

@@ -509,6 +509,9 @@ impl<'a> ReportErrorExt for UndefinedBehaviorInfo<'a> {
             ScalarSizeMismatch(_) => const_eval_scalar_size_mismatch,
             UninhabitedEnumVariantWritten(_) => const_eval_uninhabited_enum_variant_written,
             UninhabitedEnumVariantRead(_) => const_eval_uninhabited_enum_variant_read,
+            InvalidNichedEnumVariantWritten { .. } => {
+                const_eval_invalid_niched_enum_variant_written
+            }
             AbiMismatchArgument { .. } => const_eval_incompatible_types,
             AbiMismatchReturn { .. } => const_eval_incompatible_return_types,
         }
@@ -597,6 +600,9 @@ impl<'a> ReportErrorExt for UndefinedBehaviorInfo<'a> {
                 builder.arg("target_size", info.target_size);
                 builder.arg("data_size", info.data_size);
             }
+            InvalidNichedEnumVariantWritten { enum_ty } => {
+                builder.arg("ty", enum_ty.to_string());
+            }
             AbiMismatchArgument { caller_ty, callee_ty }
             | AbiMismatchReturn { caller_ty, callee_ty } => {
                 builder.arg("caller_ty", caller_ty.to_string());

compiler/rustc_const_eval/src/interpret/discriminant.rs

@@ -85,6 +85,14 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                     // Write result.
                     let niche_dest = self.project_field(dest, tag_field)?;
                     self.write_immediate(*tag_val, &niche_dest)?;
+                } else {
+                    // The untagged variant is implicitly encoded simply by having a value that is
+                    // outside the niche variants. But what if the data stored here does not
+                    // actually encode this variant? That would be bad! So let's double-check...
+                    let actual_variant = self.read_discriminant(&dest.to_op(self)?)?;
+                    if actual_variant != variant_index {
+                        throw_ub!(InvalidNichedEnumVariantWritten { enum_ty: dest.layout().ty });
+                    }
                 }
             }
         }

compiler/rustc_middle/src/mir/interpret/error.rs

@@ -356,6 +356,8 @@ pub enum UndefinedBehaviorInfo<'tcx> {
     UninhabitedEnumVariantWritten(VariantIdx),
     /// An uninhabited enum variant is projected.
     UninhabitedEnumVariantRead(VariantIdx),
+    /// Trying to set discriminant to the niched variant, but the value does not match.
+    InvalidNichedEnumVariantWritten { enum_ty: Ty<'tcx> },
     /// ABI-incompatible argument types.
     AbiMismatchArgument { caller_ty: Ty<'tcx>, callee_ty: Ty<'tcx> },
     /// ABI-incompatible return types.

src/tools/miri/tests/fail/enum-set-discriminant-niche-variant-wrong.rs

@@ -0,0 +1,31 @@
+#![feature(core_intrinsics)]
+#![feature(custom_mir)]
+
+use std::intrinsics::mir::*;
+use std::num::NonZeroI32;
+
+// We define our own option type so that we can control the varian indices.
+#[allow(unused)]
+enum Option<T> {
+    None,
+    Some(T),
+}
+use Option::*;
+
+#[custom_mir(dialect = "runtime", phase = "optimized")]
+fn set_discriminant(ptr: &mut Option<NonZeroI32>) {
+    mir! {
+        {
+            // We set the discriminant to `Some`, which is a NOP since this is the niched variant.
+            // However, the enum is actually encoding `None` currently! That's not good...
+            SetDiscriminant(*ptr, 1);
+            //~^ ERROR: trying to set discriminant of a Option<std::num::NonZero<i32>> to the niched variant, but the value does not match
+            Return()
+        }
+    }
+}
+
+pub fn main() {
+    let mut v = None;
+    set_discriminant(&mut v);
+}

src/tools/miri/tests/fail/enum-set-discriminant-niche-variant-wrong.stderr

@@ -0,0 +1,20 @@
+error: Undefined Behavior: trying to set discriminant of a Option<std::num::NonZero<i32>> to the niched variant, but the value does not match
+  --> $DIR/enum-set-discriminant-niche-variant-wrong.rs:LL:CC
+   |
+LL |             SetDiscriminant(*ptr, 1);
+   |             ^^^^^^^^^^^^^^^^^^^^^^^^ trying to set discriminant of a Option<std::num::NonZero<i32>> to the niched variant, but the value does not match
+   |
+   = help: this indicates a bug in the program: it performed an invalid operation, and caused Undefined Behavior
+   = help: see https://doc.rust-lang.org/nightly/reference/behavior-considered-undefined.html for further information
+   = note: BACKTRACE:
+   = note: inside `set_discriminant` at $DIR/enum-set-discriminant-niche-variant-wrong.rs:LL:CC
+note: inside `main`
+  --> $DIR/enum-set-discriminant-niche-variant-wrong.rs:LL:CC
+   |
+LL |     set_discriminant(&mut v);
+   |     ^^^^^^^^^^^^^^^^^^^^^^^^
+
+note: some details are omitted, run with `MIRIFLAGS=-Zmiri-backtrace=full` for a verbose backtrace
+
+error: aborting due to 1 previous error
+