From 928bd3b4e031d827ec880c7fc3ca06ce5b8243b6 Mon Sep 17 00:00:00 2001
From: Ralf Jung <post@ralfj.de>
Date: Sat, 23 Mar 2024 10:44:28 +0100
Subject: [PATCH] tag_for_variant follow-ups

---
 .../src/const_eval/eval_queries.rs            | 20 +--------------
 .../rustc_const_eval/src/const_eval/mod.rs    | 25 ++++++++++++++++---
 .../src/interpret/discriminant.rs             |  8 +++---
 3 files changed, 28 insertions(+), 25 deletions(-)

diff --git a/compiler/rustc_const_eval/src/const_eval/eval_queries.rs b/compiler/rustc_const_eval/src/const_eval/eval_queries.rs
index 16bd0296247..5a1c7cc4209 100644
--- a/compiler/rustc_const_eval/src/const_eval/eval_queries.rs
+++ b/compiler/rustc_const_eval/src/const_eval/eval_queries.rs
@@ -3,7 +3,7 @@
 use rustc_hir::def::DefKind;
 use rustc_middle::mir::interpret::{AllocId, ErrorHandled, InterpErrorInfo};
 use rustc_middle::mir::{self, ConstAlloc, ConstValue};
-use rustc_middle::query::{Key, TyCtxtAt};
+use rustc_middle::query::TyCtxtAt;
 use rustc_middle::traits::Reveal;
 use rustc_middle::ty::layout::LayoutOf;
 use rustc_middle::ty::print::with_no_trimmed_paths;
@@ -243,24 +243,6 @@ pub(crate) fn turn_into_const_value<'tcx>(
     op_to_const(&ecx, &mplace.into(), /* for diagnostics */ false)
 }
 
-/// Computes the tag (if any) for a given type and variant.
-#[instrument(skip(tcx), level = "debug")]
-pub fn tag_for_variant_provider<'tcx>(
-    tcx: TyCtxt<'tcx>,
-    (ty, variant_index): (Ty<'tcx>, abi::VariantIdx),
-) -> Option<ty::ScalarInt> {
-    assert!(ty.is_enum());
-
-    let ecx = InterpCx::new(
-        tcx,
-        ty.default_span(tcx),
-        ty::ParamEnv::reveal_all(),
-        crate::const_eval::DummyMachine,
-    );
-
-    ecx.tag_for_variant(ty, variant_index).unwrap().map(|(tag, _tag_field)| tag)
-}
-
 #[instrument(skip(tcx), level = "debug")]
 pub fn eval_to_const_value_raw_provider<'tcx>(
     tcx: TyCtxt<'tcx>,
diff --git a/compiler/rustc_const_eval/src/const_eval/mod.rs b/compiler/rustc_const_eval/src/const_eval/mod.rs
index b768c429070..8efc67bcb0c 100644
--- a/compiler/rustc_const_eval/src/const_eval/mod.rs
+++ b/compiler/rustc_const_eval/src/const_eval/mod.rs
@@ -2,10 +2,11 @@
 
 use rustc_middle::mir;
 use rustc_middle::mir::interpret::InterpErrorInfo;
-use rustc_middle::query::TyCtxtAt;
-use rustc_middle::ty::{self, Ty};
+use rustc_middle::query::{Key, TyCtxtAt};
+use rustc_middle::ty::{self, Ty, TyCtxt};
+use rustc_target::abi::VariantIdx;
 
-use crate::interpret::format_interp_error;
+use crate::interpret::{format_interp_error, InterpCx};
 
 mod dummy_machine;
 mod error;
@@ -77,3 +78,21 @@ pub(crate) fn try_destructure_mir_constant_for_user_output<'tcx>(
 
     Some(mir::DestructuredConstant { variant, fields })
 }
+
+/// Computes the tag (if any) for a given type and variant.
+#[instrument(skip(tcx), level = "debug")]
+pub fn tag_for_variant_provider<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    (ty, variant_index): (Ty<'tcx>, VariantIdx),
+) -> Option<ty::ScalarInt> {
+    assert!(ty.is_enum());
+
+    let ecx = InterpCx::new(
+        tcx,
+        ty.default_span(tcx),
+        ty::ParamEnv::reveal_all(),
+        crate::const_eval::DummyMachine,
+    );
+
+    ecx.tag_for_variant(ty, 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 40469c6632c..704f597cfdb 100644
--- a/compiler/rustc_const_eval/src/interpret/discriminant.rs
+++ b/compiler/rustc_const_eval/src/interpret/discriminant.rs
@@ -38,7 +38,7 @@ pub fn write_discriminant(
             }
             None => {
                 // No need to write the tag here, because an untagged variant is
-                // implicitly encoded. For `Niche`-optimized enums, it's by
+                // implicitly encoded. For `Niche`-optimized enums, this works by
                 // 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...
@@ -227,8 +227,10 @@ pub fn discriminant_for_variant(
         Ok(ImmTy::from_scalar(discr_value, discr_layout))
     }
 
-    /// Computes the tag value and its field number (if any) of a given variant
-    /// of type `ty`.
+    /// Computes how to write the tag of a given variant of enum `ty`:
+    /// - `None` means that nothing needs to be done as the variant is encoded implicitly
+    /// - `Some((val, field_idx))` means that the given integer value needs to be stored at the
+    ///   given field index.
     pub(crate) fn tag_for_variant(
         &self,
         ty: Ty<'tcx>,