update const arg queries

Author: lcnr

src/librustc_middle/mir/query.rs

@@ -1,10 +1,10 @@
 //! Values computed by queries that use MIR.
 
-use crate::ty::{self, Ty};
+use crate::ty::{self, Ty, TyCtxt};
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::sync::Lrc;
 use rustc_hir as hir;
-use rustc_hir::def_id::DefId;
+use rustc_hir::def_id::{DefId, LocalDefId};
 use rustc_index::bit_set::BitMatrix;
 use rustc_index::vec::IndexVec;
 use rustc_span::{Span, Symbol};
@@ -323,3 +323,24 @@ pub struct CoverageInfo {
     /// The total number of coverage region counters added to the MIR `Body`.
     pub num_counters: u32,
 }
+
+impl<'tcx> TyCtxt<'tcx> {
+    pub fn mir_borrowck_opt_const_arg(
+        self,
+        def: ty::WithOptParam<LocalDefId>,
+    ) -> &'tcx BorrowCheckResult<'tcx> {
+        if let Some(param_did) = def.param_did {
+            self.mir_borrowck_const_arg((def.did, param_did))
+        } else {
+            self.mir_borrowck(def.did)
+        }
+    }
+
+    pub fn mir_const_qualif_opt_const_arg(self, def: ty::WithOptParam<LocalDefId>) -> ConstQualifs {
+        if let Some(param_did) = def.param_did {
+            self.mir_const_qualif_const_arg((def.did, param_did))
+        } else {
+            self.mir_const_qualif(def.did)
+        }
+    }
+}

src/librustc_middle/query/mod.rs

@@ -209,11 +209,11 @@ rustc_queries! {
             cache_on_disk_if { key.is_local() }
         }
         query mir_const_qualif_const_arg(
-            key: ty::WithOptParam<LocalDefId>
+            key: (LocalDefId, DefId)
         ) -> mir::ConstQualifs {
             desc {
-                |tcx| "const checking the potential const argument `{}`",
-                tcx.def_path_str(key.did.to_def_id())
+                |tcx| "const checking the const argument `{}`",
+                tcx.def_path_str(key.0.to_def_id())
             }
         }
 
@@ -257,10 +257,10 @@ rustc_queries! {
             desc { |tcx| "optimizing MIR for `{}`", tcx.def_path_str(key) }
             cache_on_disk_if { key.is_local() }
         }
-        query optimized_mir_of_const_arg(key: ty::WithOptParam<LocalDefId>) -> &'tcx mir::Body<'tcx> {
+        query optimized_mir_of_const_arg(key: (LocalDefId, DefId)) -> &'tcx mir::Body<'tcx> {
             desc {
-                |tcx| "optimizing MIR for the potential const argument `{}`",
-                tcx.def_path_str(key.did.to_def_id())
+                |tcx| "optimizing MIR for the const argument `{}`",
+                tcx.def_path_str(key.0.to_def_id())
             }
         }
 
@@ -280,7 +280,7 @@ rustc_queries! {
             key: ty::WithOptParam<LocalDefId>
         ) -> &'tcx IndexVec<mir::Promoted, mir::Body<'tcx>> {
             desc {
-                |tcx| "optimizing promoted MIR for the potential const argument `{}`",
+                |tcx| "optimizing promoted MIR for the const argument `{}`",
                 tcx.def_path_str(key.did.to_def_id()),
             }
         }
@@ -496,8 +496,8 @@ rustc_queries! {
             desc { |tcx| "unsafety-checking `{}`", tcx.def_path_str(key.to_def_id()) }
             cache_on_disk_if { true }
         }
-        query unsafety_check_result_const_arg(key: ty::WithOptParam<LocalDefId>) -> &'tcx mir::UnsafetyCheckResult {
-            desc { |tcx| "unsafety-checking the potential const arg `{}`", tcx.def_path_str(key.did.to_def_id()) }
+        query unsafety_check_result_const_arg(key: (LocalDefId, DefId)) -> &'tcx mir::UnsafetyCheckResult {
+            desc { |tcx| "unsafety-checking the const arg `{}`", tcx.def_path_str(key.0.to_def_id()) }
         }
 
         /// HACK: when evaluated, this reports a "unsafe derive on repr(packed)" error.
@@ -579,12 +579,12 @@ rustc_queries! {
             desc { |tcx| "type-checking `{}`", tcx.def_path_str(key.to_def_id()) }
             cache_on_disk_if { true }
         }
-        query _typeck_tables_of_const_arg(
-            key: ty::WithOptParam<LocalDefId>
+        query typeck_tables_of_const_arg(
+            key: (LocalDefId, DefId)
         ) -> &'tcx ty::TypeckTables<'tcx> {
             desc {
                 |tcx| "type-checking the const argument `{}`",
-                tcx.def_path_str(key.did.to_def_id()),
+                tcx.def_path_str(key.0.to_def_id()),
             }
         }
         query diagnostic_only_typeck_tables_of(key: LocalDefId) -> &'tcx ty::TypeckTables<'tcx> {
@@ -627,10 +627,10 @@ rustc_queries! {
                     || opt_result.map_or(false, |r| !r.concrete_opaque_types.is_empty())
             }
         }
-        query mir_borrowck_const_arg(key: ty::WithOptParam<LocalDefId>) -> &'tcx mir::BorrowCheckResult<'tcx> {
+        query mir_borrowck_const_arg(key: (LocalDefId, DefId)) -> &'tcx mir::BorrowCheckResult<'tcx> {
             desc {
-                |tcx| "borrow-checking the potential const argument`{}`",
-                tcx.def_path_str(key.did.to_def_id())
+                |tcx| "borrow-checking the const argument`{}`",
+                tcx.def_path_str(key.0.to_def_id())
             }
         }
     }
@@ -1501,11 +1501,11 @@ rustc_queries! {
         }
 
         query resolve_instance_of_const_arg(
-            key: ty::ParamEnvAnd<'tcx, (ty::WithOptParam<DefId>, SubstsRef<'tcx>)>
+            key: ty::ParamEnvAnd<'tcx, (LocalDefId, DefId, SubstsRef<'tcx>)>
         ) -> Result<Option<ty::Instance<'tcx>>, ErrorReported> {
             desc {
-                "resolving instance of the potential const argument `{}`",
-                ty::Instance::new(key.value.0.did, key.value.1),
+                "resolving instance of the const argument `{}`",
+                ty::Instance::new(key.value.0.to_def_id(), key.value.2),
             }
         }
     }

src/librustc_middle/ty/context.rs

@@ -980,12 +980,12 @@ pub struct GlobalCtxt<'tcx> {
 }
 
 impl<'tcx> TyCtxt<'tcx> {
-    pub fn typeck_tables_of_const_arg(
+    pub fn typeck_tables_of_opt_const_arg(
         self,
         def: ty::WithOptParam<LocalDefId>,
     ) -> &'tcx TypeckTables<'tcx> {
-        if def.param_did.is_some() {
-            self._typeck_tables_of_const_arg(def)
+        if let Some(param_did) = def.param_did {
+            self.typeck_tables_of_const_arg((def.did, param_did))
         } else {
             self.typeck_tables_of(def.did)
         }

src/librustc_middle/ty/instance.rs

@@ -359,7 +359,14 @@ impl<'tcx> Instance<'tcx> {
         substs: SubstsRef<'tcx>,
     ) -> Result<Option<Instance<'tcx>>, ErrorReported> {
         let substs = tcx.erase_regions(&substs);
-        tcx.resolve_instance_of_const_arg(tcx.erase_regions(&param_env.and((def, substs))))
+
+        if let Some((did, param_did)) = def.as_const_arg() {
+            tcx.resolve_instance_of_const_arg(
+                tcx.erase_regions(&param_env.and((did, param_did, substs))),
+            )
+        } else {
+            tcx.resolve_instance(tcx.erase_regions(&param_env.and((def.did, substs))))
+        }
     }
 
     pub fn resolve_for_fn_ptr(

src/librustc_middle/ty/mod.rs

@@ -1600,6 +1600,16 @@ impl WithOptParam<DefId> {
         self.did.as_local().map(|did| WithOptParam { did, param_did: self.param_did })
     }
 
+    pub fn as_const_arg(self) -> Option<(LocalDefId, DefId)> {
+        if let Some(param_did) = self.param_did {
+            if let Some(did) = self.did.as_local() {
+                return Some((did, param_did));
+            }
+        }
+
+        None
+    }
+
     pub fn expect_local(self) -> WithOptParam<LocalDefId> {
         self.as_local().unwrap()
     }
@@ -1611,10 +1621,6 @@ impl WithOptParam<DefId> {
     pub fn ty_def_id(self) -> DefId {
         self.param_did.unwrap_or(self.did)
     }
-
-    pub fn init_me_bby(tcx: TyCtxt<'_>, did: DefId) -> WithOptParam<DefId> {
-        WithOptParam { did, param_did: did.as_local().and_then(|did| tcx.opt_const_param_of(did)) }
-    }
 }
 
 /// When type checking, we use the `ParamEnv` to track
@@ -2889,8 +2895,9 @@ impl<'tcx> TyCtxt<'tcx> {
     pub fn instance_mir(self, instance: ty::InstanceDef<'tcx>) -> &'tcx Body<'tcx> {
         match instance {
             ty::InstanceDef::Item(def) => {
-                if let Some(def) = def.as_local() {
-                    self.optimized_mir_of_const_arg(def)
+                if let Some((did, param_did)) = def.as_const_arg() {
+                    // The `param_did` is only `Some` for local `DefId`s.
+                    self.optimized_mir_of_const_arg((did, param_did))
                 } else {
                     self.optimized_mir(def.did)
                 }

src/librustc_middle/ty/query/keys.rs

@@ -138,6 +138,17 @@ impl Key for (DefId, LocalDefId) {
     }
 }
 
+impl Key for (LocalDefId, DefId) {
+    type CacheSelector = DefaultCacheSelector;
+
+    fn query_crate(&self) -> CrateNum {
+        LOCAL_CRATE
+    }
+    fn default_span(&self, tcx: TyCtxt<'_>) -> Span {
+        self.0.default_span(tcx)
+    }
+}
+
 impl Key for (CrateNum, DefId) {
     type CacheSelector = DefaultCacheSelector;
 
@@ -182,14 +193,14 @@ impl<'tcx> Key for (DefId, SubstsRef<'tcx>) {
     }
 }
 
-impl<'tcx> Key for (ty::WithOptParam<DefId>, SubstsRef<'tcx>) {
+impl<'tcx> Key for (LocalDefId, DefId, SubstsRef<'tcx>) {
     type CacheSelector = DefaultCacheSelector;
 
     fn query_crate(&self) -> CrateNum {
         LOCAL_CRATE
     }
     fn default_span(&self, tcx: TyCtxt<'_>) -> Span {
-        self.0.did.default_span(tcx)
+        self.0.default_span(tcx)
     }
 }
 

src/librustc_middle/ty/structural_impls.rs

@@ -272,6 +272,7 @@ CloneTypeFoldableAndLiftImpls! {
     ::rustc_span::symbol::Symbol,
     ::rustc_hir::def::Res,
     ::rustc_hir::def_id::DefId,
+    ::rustc_hir::def_id::LocalDefId,
     ::rustc_hir::LlvmInlineAsmInner,
     ::rustc_hir::MatchSource,
     ::rustc_hir::Mutability,
@@ -719,6 +720,18 @@ impl<'tcx, T: TypeFoldable<'tcx>, U: TypeFoldable<'tcx>> TypeFoldable<'tcx> for
     }
 }
 
+impl<'tcx, A: TypeFoldable<'tcx>, B: TypeFoldable<'tcx>, C: TypeFoldable<'tcx>> TypeFoldable<'tcx>
+    for (A, B, C)
+{
+    fn super_fold_with<F: TypeFolder<'tcx>>(&self, folder: &mut F) -> (A, B, C) {
+        (self.0.fold_with(folder), self.1.fold_with(folder), self.2.fold_with(folder))
+    }
+
+    fn super_visit_with<V: TypeVisitor<'tcx>>(&self, visitor: &mut V) -> bool {
+        self.0.visit_with(visitor) || self.1.visit_with(visitor) || self.2.visit_with(visitor)
+    }
+}
+
 EnumTypeFoldableImpl! {
     impl<'tcx, T> TypeFoldable<'tcx> for Option<T> {
         (Some)(a),

src/librustc_mir/borrow_check/mod.rs

@@ -89,8 +89,8 @@ const DEREF_PROJECTION: &[PlaceElem<'_>; 1] = &[ProjectionElem::Deref];
 pub fn provide(providers: &mut Providers) {
     *providers = Providers {
         mir_borrowck: |tcx, did| mir_borrowck(tcx, ty::WithOptParam::dummy(did)),
-        mir_borrowck_const_arg: |tcx, def| {
-            if def.param_did.is_none() { tcx.mir_borrowck(def.did) } else { mir_borrowck(tcx, def) }
+        mir_borrowck_const_arg: |tcx, (did, param_did)| {
+            mir_borrowck(tcx, ty::WithOptParam { did, param_did: Some(param_did) })
         },
         ..*providers
     };
@@ -101,8 +101,8 @@ fn mir_borrowck<'tcx>(
     def: ty::WithOptParam<LocalDefId>,
 ) -> &'tcx BorrowCheckResult<'tcx> {
     if def.param_did.is_none() {
-        if let param_did @ Some(_) = tcx.opt_const_param_of(def.did) {
-            return tcx.mir_borrowck_const_arg(ty::WithOptParam { param_did, ..def });
+        if let Some(param_did) = tcx.opt_const_param_of(def.did) {
+            return tcx.mir_borrowck_const_arg((def.did, param_did));
         }
     }
 
@@ -150,7 +150,7 @@ fn do_mir_borrowck<'a, 'tcx>(
     }
 
     // Gather the upvars of a closure, if any.
-    let tables = tcx.typeck_tables_of_const_arg(def);
+    let tables = tcx.typeck_tables_of_opt_const_arg(def);
     if let Some(ErrorReported) = tables.tainted_by_errors {
         infcx.set_tainted_by_errors();
     }

src/librustc_mir/const_eval/eval_queries.rs

@@ -292,7 +292,8 @@ pub fn const_eval_raw_provider<'tcx>(
 
     if let Some(def) = def.as_local() {
         if tcx.has_typeck_tables(def.did) {
-            if let Some(error_reported) = tcx.typeck_tables_of_const_arg(def).tainted_by_errors {
+            if let Some(error_reported) = tcx.typeck_tables_of_opt_const_arg(def).tainted_by_errors
+            {
                 return Err(ErrorHandled::Reported(error_reported));
             }
         }

src/librustc_mir/interpret/eval_context.rs

@@ -398,7 +398,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
         if let Some(def) = def.as_local() {
             if self.tcx.has_typeck_tables(def.did) {
                 if let Some(error_reported) =
-                    self.tcx.typeck_tables_of_const_arg(def).tainted_by_errors
+                    self.tcx.typeck_tables_of_opt_const_arg(def).tainted_by_errors
                 {
                     throw_inval!(TypeckError(error_reported))
                 }
@@ -415,8 +415,8 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
         match instance {
             ty::InstanceDef::Item(def) => {
                 if self.tcx.is_mir_available(def.did) {
-                    if let Some(def) = def.as_local() {
-                        Ok(self.tcx.optimized_mir_of_const_arg(def))
+                    if let Some((did, param_did)) = def.as_const_arg() {
+                        Ok(self.tcx.optimized_mir_of_const_arg((did, param_did)))
                     } else {
                         Ok(self.tcx.optimized_mir(def.did))
                     }