Auto merge of #83207 - oli-obk:valtree2, r=lcnr

normalize mir::Constant differently from ty::Const in preparation for valtrees

Valtrees are unable to represent many kind of constant values (this is on purpose). For constants that are used at runtime, we do not need a valtree representation and can thus use a different form of evaluation. In order to make this explicit and less fragile, I added a `fold_constant` method to `TypeFolder` and implemented it for normalization. Normalization can now, when it wants to eagerly evaluate a constant, normalize `mir::Constant` directly into a `mir::ConstantKind::Val` instead of relying on the `ty::Const` evaluation.

In the future we can get rid of the `ty::Const` in there entirely and add our own `Unevaluated` variant to `mir::ConstantKind`. This would allow us to remove the `promoted` field from `ty::ConstKind::Unevaluated`, as promoteds can never occur in the type system.

cc `@rust-lang/wg-const-eval`

r? `@lcnr`

compiler/rustc_infer/src/infer/resolve.rs

@@ -1,5 +1,6 @@
 use super::type_variable::{TypeVariableOrigin, TypeVariableOriginKind};
 use super::{FixupError, FixupResult, InferCtxt, Span};
+use rustc_middle::mir;
 use rustc_middle::ty::fold::{TypeFolder, TypeVisitor};
 use rustc_middle::ty::{self, Const, InferConst, Ty, TyCtxt, TypeFoldable};
 
@@ -46,6 +47,10 @@ impl<'a, 'tcx> TypeFolder<'tcx> for OpportunisticVarResolver<'a, 'tcx> {
             ct.super_fold_with(self)
         }
     }
+
+    fn fold_mir_const(&mut self, constant: mir::ConstantKind<'tcx>) -> mir::ConstantKind<'tcx> {
+        constant.super_fold_with(self)
+    }
 }
 
 /// The opportunistic region resolver opportunistically resolves regions

compiler/rustc_middle/src/mir/mod.rs

@@ -2392,7 +2392,8 @@ pub struct Constant<'tcx> {
     pub literal: ConstantKind<'tcx>,
 }
 
-#[derive(Clone, Copy, PartialEq, PartialOrd, TyEncodable, TyDecodable, Hash, HashStable, Debug)]
+#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, TyEncodable, TyDecodable, Hash, HashStable, Debug)]
+#[derive(Lift)]
 pub enum ConstantKind<'tcx> {
     /// This constant came from the type system
     Ty(&'tcx ty::Const<'tcx>),
@@ -2691,7 +2692,13 @@ impl<'tcx> Display for Constant<'tcx> {
             ty::FnDef(..) => {}
             _ => write!(fmt, "const ")?,
         }
-        match self.literal {
+        Display::fmt(&self.literal, fmt)
+    }
+}
+
+impl<'tcx> Display for ConstantKind<'tcx> {
+    fn fmt(&self, fmt: &mut Formatter<'_>) -> fmt::Result {
+        match *self {
             ConstantKind::Ty(c) => pretty_print_const(c, fmt, true),
             ConstantKind::Val(val, ty) => pretty_print_const_value(val, ty, fmt, true),
         }

compiler/rustc_middle/src/mir/type_foldable.rs

@@ -348,6 +348,11 @@ impl<'tcx> TypeFoldable<'tcx> for Constant<'tcx> {
 }
 
 impl<'tcx> TypeFoldable<'tcx> for ConstantKind<'tcx> {
+    #[inline(always)]
+    fn fold_with<F: TypeFolder<'tcx>>(self, folder: &mut F) -> Self {
+        folder.fold_mir_const(self)
+    }
+
     fn super_fold_with<F: TypeFolder<'tcx>>(self, folder: &mut F) -> Self {
         match self {
             ConstantKind::Ty(c) => ConstantKind::Ty(c.fold_with(folder)),

compiler/rustc_middle/src/query/mod.rs

@@ -1486,6 +1486,13 @@ rustc_queries! {
         desc { "normalizing `{}`", goal.value }
     }
 
+    /// Do not call this query directly: invoke `normalize_erasing_regions` instead.
+    query normalize_mir_const_after_erasing_regions(
+        goal: ParamEnvAnd<'tcx, mir::ConstantKind<'tcx>>
+    ) -> mir::ConstantKind<'tcx> {
+        desc { "normalizing `{}`", goal.value }
+    }
+
     query implied_outlives_bounds(
         goal: CanonicalTyGoal<'tcx>
     ) -> Result<

compiler/rustc_middle/src/ty/erase_regions.rs

@@ -1,3 +1,4 @@
+use crate::mir;
 use crate::ty::fold::{TypeFoldable, TypeFolder};
 use crate::ty::{self, Ty, TyCtxt, TypeFlags};
 
@@ -65,4 +66,8 @@ impl TypeFolder<'tcx> for RegionEraserVisitor<'tcx> {
             _ => self.tcx.lifetimes.re_erased,
         }
     }
+
+    fn fold_mir_const(&mut self, c: mir::ConstantKind<'tcx>) -> mir::ConstantKind<'tcx> {
+        c.super_fold_with(self)
+    }
 }

compiler/rustc_middle/src/ty/fold.rs

@@ -30,6 +30,7 @@
 //!
 //! These methods return true to indicate that the visitor has found what it is
 //! looking for, and does not need to visit anything else.
+use crate::mir;
 use crate::ty::{self, flags::FlagComputation, Binder, Ty, TyCtxt, TypeFlags};
 use rustc_hir as hir;
 use rustc_hir::def_id::DefId;
@@ -180,6 +181,10 @@ pub trait TypeFolder<'tcx>: Sized {
     fn fold_const(&mut self, c: &'tcx ty::Const<'tcx>) -> &'tcx ty::Const<'tcx> {
         c.super_fold_with(self)
     }
+
+    fn fold_mir_const(&mut self, c: mir::ConstantKind<'tcx>) -> mir::ConstantKind<'tcx> {
+        bug!("most type folders should not be folding MIR datastructures: {:?}", c)
+    }
 }
 
 pub trait TypeVisitor<'tcx>: Sized {

compiler/rustc_middle/src/ty/instance.rs

@@ -483,6 +483,7 @@ impl<'tcx> Instance<'tcx> {
         if let Some(substs) = self.substs_for_mir_body() { v.subst(tcx, substs) } else { *v }
     }
 
+    #[inline(always)]
     pub fn subst_mir_and_normalize_erasing_regions<T>(
         &self,
         tcx: TyCtxt<'tcx>,

compiler/rustc_middle/src/ty/normalize_erasing_regions.rs

@@ -7,6 +7,7 @@
 //! `normalize_generic_arg_after_erasing_regions` query for each type
 //! or constant found within. (This underlying query is what is cached.)
 
+use crate::mir;
 use crate::ty::fold::{TypeFoldable, TypeFolder};
 use crate::ty::subst::{Subst, SubstsRef};
 use crate::ty::{self, Ty, TyCtxt};
@@ -101,4 +102,10 @@ impl TypeFolder<'tcx> for NormalizeAfterErasingRegionsFolder<'tcx> {
         let arg = self.param_env.and(c.into());
         self.tcx.normalize_generic_arg_after_erasing_regions(arg).expect_const()
     }
+
+    #[inline]
+    fn fold_mir_const(&mut self, c: mir::ConstantKind<'tcx>) -> mir::ConstantKind<'tcx> {
+        let arg = self.param_env.and(c);
+        self.tcx.normalize_mir_const_after_erasing_regions(arg)
+    }
 }

compiler/rustc_middle/src/ty/subst.rs

@@ -1,5 +1,6 @@
 // Type substitutions.
 
+use crate::mir;
 use crate::ty::codec::{TyDecoder, TyEncoder};
 use crate::ty::fold::{TypeFoldable, TypeFolder, TypeVisitor};
 use crate::ty::sty::{ClosureSubsts, GeneratorSubsts};
@@ -506,6 +507,11 @@ impl<'a, 'tcx> TypeFolder<'tcx> for SubstFolder<'a, 'tcx> {
             c.super_fold_with(self)
         }
     }
+
+    #[inline]
+    fn fold_mir_const(&mut self, c: mir::ConstantKind<'tcx>) -> mir::ConstantKind<'tcx> {
+        c.super_fold_with(self)
+    }
 }
 
 impl<'a, 'tcx> SubstFolder<'a, 'tcx> {

compiler/rustc_mir/src/const_eval/mod.rs

@@ -110,7 +110,7 @@ fn const_to_valtree_inner<'tcx>(
 
             let variant = ecx.read_discriminant(&place.into()).unwrap().1;
 
-            branches(def.variants[variant].fields.len(), Some(variant))
+            branches(def.variants[variant].fields.len(), def.is_enum().then_some(variant))
         }
 
         ty::Never

compiler/rustc_mir/src/interpret/eval_context.rs

@@ -524,6 +524,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
         }
     }
 
+    #[inline(always)]
     pub fn layout_of_local(
         &self,
         frame: &Frame<'mir, 'tcx, M::PointerTag, M::FrameExtra>,

compiler/rustc_mir/src/interpret/intrinsics.rs

@@ -171,8 +171,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                 };
                 let val =
                     self.tcx.const_eval_global_id(self.param_env, gid, Some(self.tcx.span))?;
-                let const_ = ty::Const { val: ty::ConstKind::Value(val), ty };
-                let val = self.const_to_op(&const_, None)?;
+                let val = self.const_val_to_op(val, ty, Some(dest.layout))?;
                 self.copy_op(&val, dest)?;
             }
 

compiler/rustc_mir/src/interpret/operand.rs

@@ -570,7 +570,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
     ) -> InterpResult<'tcx, OpTy<'tcx, M::PointerTag>> {
         match val {
             mir::ConstantKind::Ty(ct) => self.const_to_op(ct, layout),
-            mir::ConstantKind::Val(val, ty) => self.const_val_to_op(*val, ty, None),
+            mir::ConstantKind::Val(val, ty) => self.const_val_to_op(*val, ty, layout),
         }
     }
 

compiler/rustc_mir/src/monomorphize/collector.rs

@@ -184,7 +184,6 @@ use rustc_hir::def_id::{DefId, DefIdMap, LocalDefId, LOCAL_CRATE};
 use rustc_hir::itemlikevisit::ItemLikeVisitor;
 use rustc_hir::lang_items::LangItem;
 use rustc_index::bit_set::GrowableBitSet;
-use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
 use rustc_middle::mir::interpret::{AllocId, ConstValue};
 use rustc_middle::mir::interpret::{ErrorHandled, GlobalAlloc, Scalar};
 use rustc_middle::mir::mono::{InstantiationMode, MonoItem};
@@ -193,6 +192,7 @@ use rustc_middle::mir::{self, Local, Location};
 use rustc_middle::ty::adjustment::{CustomCoerceUnsized, PointerCast};
 use rustc_middle::ty::subst::{GenericArgKind, InternalSubsts};
 use rustc_middle::ty::{self, GenericParamDefKind, Instance, Ty, TyCtxt, TypeFoldable};
+use rustc_middle::{middle::codegen_fn_attrs::CodegenFnAttrFlags, mir::visit::TyContext};
 use rustc_session::config::EntryFnType;
 use rustc_span::source_map::{dummy_spanned, respan, Span, Spanned, DUMMY_SP};
 use smallvec::SmallVec;
@@ -638,6 +638,35 @@ impl<'a, 'tcx> MirVisitor<'tcx> for MirNeighborCollector<'a, 'tcx> {
         self.super_rvalue(rvalue, location);
     }
 
+    /// This does not walk the constant, as it has been handled entirely here and trying
+    /// to walk it would attempt to evaluate the `ty::Const` inside, which doesn't necessarily
+    /// work, as some constants cannot be represented in the type system.
+    fn visit_constant(&mut self, constant: &mir::Constant<'tcx>, location: Location) {
+        let literal = self.monomorphize(constant.literal);
+        let val = match literal {
+            mir::ConstantKind::Val(val, _) => val,
+            mir::ConstantKind::Ty(ct) => match ct.val {
+                ty::ConstKind::Value(val) => val,
+                ty::ConstKind::Unevaluated(ct) => {
+                    let param_env = ty::ParamEnv::reveal_all();
+                    match self.tcx.const_eval_resolve(param_env, ct, None) {
+                        // The `monomorphize` call should have evaluated that constant already.
+                        Ok(val) => val,
+                        Err(ErrorHandled::Reported(ErrorReported) | ErrorHandled::Linted) => return,
+                        Err(ErrorHandled::TooGeneric) => span_bug!(
+                            self.body.source_info(location).span,
+                            "collection encountered polymorphic constant: {:?}",
+                            literal
+                        ),
+                    }
+                }
+                _ => return,
+            },
+        };
+        collect_const_value(self.tcx, val, self.output);
+        self.visit_ty(literal.ty(), TyContext::Location(location));
+    }
+
     fn visit_const(&mut self, constant: &&'tcx ty::Const<'tcx>, location: Location) {
         debug!("visiting const {:?} @ {:?}", *constant, location);
 
@@ -648,7 +677,13 @@ impl<'a, 'tcx> MirVisitor<'tcx> for MirNeighborCollector<'a, 'tcx> {
             ty::ConstKind::Value(val) => collect_const_value(self.tcx, val, self.output),
             ty::ConstKind::Unevaluated(unevaluated) => {
                 match self.tcx.const_eval_resolve(param_env, unevaluated, None) {
-                    Ok(val) => collect_const_value(self.tcx, val, self.output),
+                    // The `monomorphize` call should have evaluated that constant already.
+                    Ok(val) => span_bug!(
+                        self.body.source_info(location).span,
+                        "collection encountered the unevaluated constant {} which evaluated to {:?}",
+                        substituted_constant,
+                        val
+                    ),
                     Err(ErrorHandled::Reported(ErrorReported) | ErrorHandled::Linted) => {}
                     Err(ErrorHandled::TooGeneric) => span_bug!(
                         self.body.source_info(location).span,

compiler/rustc_mir/src/util/pretty.rs

@@ -452,7 +452,11 @@ impl Visitor<'tcx> for ExtraComments<'tcx> {
                 match literal {
                     ConstantKind::Ty(literal) => self.push(&format!("+ literal: {:?}", literal)),
                     ConstantKind::Val(val, ty) => {
-                        self.push(&format!("+ literal: {:?}, {}", val, ty))
+                        // To keep the diffs small, we render this almost like we render ty::Const
+                        self.push(&format!(
+                            "+ literal: Const {{ ty: {}, val: Value({:?}) }}",
+                            ty, val
+                        ))
                     }
                 }
             }
@@ -465,7 +469,21 @@ impl Visitor<'tcx> for ExtraComments<'tcx> {
         if use_verbose(ty) {
             self.push("ty::Const");
             self.push(&format!("+ ty: {:?}", ty));
-            self.push(&format!("+ val: {:?}", val));
+            let val = match val {
+                ty::ConstKind::Param(p) => format!("Param({})", p),
+                ty::ConstKind::Infer(infer) => format!("Infer({:?})", infer),
+                ty::ConstKind::Bound(idx, var) => format!("Bound({:?}, {:?})", idx, var),
+                ty::ConstKind::Placeholder(ph) => format!("PlaceHolder({:?})", ph),
+                ty::ConstKind::Unevaluated(uv) => format!(
+                    "Unevaluated({}, {:?}, {:?})",
+                    self.tcx.def_path_str(uv.def.did),
+                    uv.substs,
+                    uv.promoted
+                ),
+                ty::ConstKind::Value(val) => format!("Value({:?})", val),
+                ty::ConstKind::Error(_) => format!("Error"),
+            };
+            self.push(&format!("+ val: {}", val));
         }
     }
 

compiler/rustc_mir_build/src/thir/pattern/const_to_pat.rs

@@ -246,6 +246,18 @@ impl<'a, 'tcx> ConstToPat<'a, 'tcx> {
             })
     }
 
+    fn field_pats(
+        &self,
+        vals: impl Iterator<Item = &'tcx ty::Const<'tcx>>,
+    ) -> Result<Vec<FieldPat<'tcx>>, FallbackToConstRef> {
+        vals.enumerate()
+            .map(|(idx, val)| {
+                let field = Field::new(idx);
+                Ok(FieldPat { field, pattern: self.recur(val, false)? })
+            })
+            .collect()
+    }
+
     // Recursive helper for `to_pat`; invoke that (instead of calling this directly).
     fn recur(
         &self,
@@ -257,16 +269,6 @@ impl<'a, 'tcx> ConstToPat<'a, 'tcx> {
         let tcx = self.tcx();
         let param_env = self.param_env;
 
-        let field_pats = |vals: &[&'tcx ty::Const<'tcx>]| -> Result<_, _> {
-            vals.iter()
-                .enumerate()
-                .map(|(idx, val)| {
-                    let field = Field::new(idx);
-                    Ok(FieldPat { field, pattern: self.recur(val, false)? })
-                })
-                .collect()
-        };
-
         let kind = match cv.ty.kind() {
             ty::Float(_) => {
                 tcx.struct_span_lint_hir(
@@ -361,12 +363,12 @@ impl<'a, 'tcx> ConstToPat<'a, 'tcx> {
                     variant_index: destructured
                         .variant
                         .expect("destructed const of adt without variant id"),
-                    subpatterns: field_pats(destructured.fields)?,
+                    subpatterns: self.field_pats(destructured.fields.iter().copied())?,
                 }
             }
             ty::Tuple(_) | ty::Adt(_, _) => {
                 let destructured = tcx.destructure_const(param_env.and(cv));
-                PatKind::Leaf { subpatterns: field_pats(destructured.fields)? }
+                PatKind::Leaf { subpatterns: self.field_pats(destructured.fields.iter().copied())? }
             }
             ty::Array(..) => PatKind::Array {
                 prefix: tcx

compiler/rustc_query_impl/src/keys.rs

@@ -255,6 +255,15 @@ impl<'tcx> Key for GenericArg<'tcx> {
     }
 }
 
+impl<'tcx> Key for mir::ConstantKind<'tcx> {
+    fn query_crate(&self) -> CrateNum {
+        LOCAL_CRATE
+    }
+    fn default_span(&self, _: TyCtxt<'_>) -> Span {
+        DUMMY_SP
+    }
+}
+
 impl<'tcx> Key for &'tcx ty::Const<'tcx> {
     fn query_crate(&self) -> CrateNum {
         LOCAL_CRATE

compiler/rustc_trait_selection/src/traits/query/normalize.rs

@@ -10,6 +10,7 @@ use crate::traits::{Obligation, ObligationCause, PredicateObligation, Reveal};
 use rustc_data_structures::sso::SsoHashMap;
 use rustc_data_structures::stack::ensure_sufficient_stack;
 use rustc_infer::traits::Normalized;
+use rustc_middle::mir;
 use rustc_middle::ty::fold::{TypeFoldable, TypeFolder};
 use rustc_middle::ty::subst::Subst;
 use rustc_middle::ty::{self, Ty, TyCtxt};
@@ -214,4 +215,8 @@ impl<'cx, 'tcx> TypeFolder<'tcx> for QueryNormalizer<'cx, 'tcx> {
         let constant = constant.super_fold_with(self);
         constant.eval(self.infcx.tcx, self.param_env)
     }
+
+    fn fold_mir_const(&mut self, constant: mir::ConstantKind<'tcx>) -> mir::ConstantKind<'tcx> {
+        constant.super_fold_with(self)
+    }
 }

compiler/rustc_traits/src/normalize_erasing_regions.rs

@@ -1,24 +1,35 @@
 use rustc_infer::infer::TyCtxtInferExt;
 use rustc_middle::traits::query::NoSolution;
 use rustc_middle::ty::query::Providers;
-use rustc_middle::ty::subst::GenericArg;
 use rustc_middle::ty::{self, ParamEnvAnd, TyCtxt, TypeFoldable};
 use rustc_trait_selection::traits::query::normalize::AtExt;
 use rustc_trait_selection::traits::{Normalized, ObligationCause};
 use std::sync::atomic::Ordering;
 
 crate fn provide(p: &mut Providers) {
-    *p = Providers { normalize_generic_arg_after_erasing_regions, ..*p };
+    *p = Providers {
+        normalize_generic_arg_after_erasing_regions: |tcx, goal| {
+            debug!("normalize_generic_arg_after_erasing_regions(goal={:#?})", goal);
+
+            tcx.sess
+                .perf_stats
+                .normalize_generic_arg_after_erasing_regions
+                .fetch_add(1, Ordering::Relaxed);
+            normalize_after_erasing_regions(tcx, goal)
+        },
+        normalize_mir_const_after_erasing_regions: |tcx, goal| {
+            normalize_after_erasing_regions(tcx, goal)
+        },
+        ..*p
+    };
 }
 
-fn normalize_generic_arg_after_erasing_regions<'tcx>(
+#[instrument(level = "debug", skip(tcx))]
+fn normalize_after_erasing_regions<'tcx, T: TypeFoldable<'tcx> + PartialEq + Copy>(
     tcx: TyCtxt<'tcx>,
-    goal: ParamEnvAnd<'tcx, GenericArg<'tcx>>,
-) -> GenericArg<'tcx> {
-    debug!("normalize_generic_arg_after_erasing_regions(goal={:#?})", goal);
-
+    goal: ParamEnvAnd<'tcx, T>,
+) -> T {
     let ParamEnvAnd { param_env, value } = goal;
-    tcx.sess.perf_stats.normalize_generic_arg_after_erasing_regions.fetch_add(1, Ordering::Relaxed);
     tcx.infer_ctxt().enter(|infcx| {
         let cause = ObligationCause::dummy();
         match infcx.at(&cause, param_env).normalize(value) {