Use Key impl to select cache.

4 files changed, 47 insertions, 13 deletions

diff --git a/compiler/rustc_middle/src/query/keys.rs b/compiler/rustc_middle/src/query/keys.rs
index e4a6877e01b..880632561b9 100644
--- a/compiler/rustc_middle/src/query/keys.rs
+++ b/compiler/rustc_middle/src/query/keys.rs
@@ -8,12 +8,15 @@ use crate::ty::subst::{GenericArg, SubstsRef};
 use crate::ty::{self, layout::TyAndLayout, Ty, TyCtxt};
 use rustc_hir::def_id::{CrateNum, DefId, LocalDefId, LOCAL_CRATE};
 use rustc_hir::hir_id::{HirId, OwnerId};
+use rustc_query_system::query::{DefaultCacheSelector, VecCacheSelector};
 use rustc_span::symbol::{Ident, Symbol};
 use rustc_span::{Span, DUMMY_SP};
 
 /// The `Key` trait controls what types can legally be used as the key
 /// for a query.
-pub trait Key {
+pub trait Key: Sized {
+    type CacheSelector = DefaultCacheSelector<Self>;
+
     /// Given an instance of this key, what crate is it referring to?
     /// This is used to find the provider.
     fn query_crate_is_local(&self) -> bool;
@@ -100,6 +103,8 @@ impl<'tcx> Key for mir::interpret::LitToConstInput<'tcx> {
 }
 
 impl Key for CrateNum {
+    type CacheSelector = VecCacheSelector<Self>;
+
     #[inline(always)]
     fn query_crate_is_local(&self) -> bool {
         *self == LOCAL_CRATE
@@ -110,6 +115,8 @@ impl Key for CrateNum {
 }
 
 impl Key for OwnerId {
+    type CacheSelector = VecCacheSelector<Self>;
+
     #[inline(always)]
     fn query_crate_is_local(&self) -> bool {
         true
@@ -123,6 +130,8 @@ impl Key for OwnerId {
 }
 
 impl Key for LocalDefId {
+    type CacheSelector = VecCacheSelector<Self>;
+
     #[inline(always)]
     fn query_crate_is_local(&self) -> bool {
         true
diff --git a/compiler/rustc_middle/src/ty/query.rs b/compiler/rustc_middle/src/ty/query.rs
index 985d396a2d0..0f62da9992f 100644
--- a/compiler/rustc_middle/src/ty/query.rs
+++ b/compiler/rustc_middle/src/ty/query.rs
@@ -15,6 +15,7 @@ use crate::mir::interpret::{
 };
 use crate::mir::interpret::{LitToConstError, LitToConstInput};
 use crate::mir::mono::CodegenUnit;
+use crate::query::Key;
 use crate::thir;
 use crate::traits::query::{
     CanonicalPredicateGoal, CanonicalProjectionGoal, CanonicalTyGoal,
@@ -120,17 +121,15 @@ macro_rules! query_helper_param_ty {
 }
 
 macro_rules! query_storage {
-    // FIXME(cjgillot) this macro-based way to perform type-based dispatch is clearly brittle.
-    // It should probably be replaced by an associated type on the `Key` trait.
-    ([][CrateNum, $V:ty]) => { VecCache<CrateNum, $V> };
-    ([(arena_cache) $($rest:tt)*][CrateNum, $V:ty]) => { VecArenaCache<'tcx, CrateNum, $V> };
-    ([][LocalDefId, $V:ty]) => { VecCache<LocalDefId, $V> };
-    ([(arena_cache) $($rest:tt)*][LocalDefId, $V:ty]) => { VecArenaCache<'tcx, LocalDefId, $V> };
-    ([][hir::OwnerId, $V:ty]) => { VecCache<hir::OwnerId, $V> };
-    ([(arena_cache) $($rest:tt)*][hir::OwnerId, $V:ty]) => { VecArenaCache<'tcx, hir::OwnerId, $V> };
-    ([][$K:ty, $V:ty]) => { DefaultCache<$K, $V> };
-    ([(arena_cache) $($rest:tt)*][$K:ty, $V:ty]) => { ArenaCache<'tcx, $K, $V> };
-    ([$other:tt $($modifiers:tt)*][$($args:tt)*]) => { query_storage!([$($modifiers)*][$($args)*]) };
+    ([][$K:ty, $V:ty]) => {
+        <<$K as Key>::CacheSelector as CacheSelector<'tcx, $V>>::Cache
+    };
+    ([(arena_cache) $($rest:tt)*][$K:ty, $V:ty]) => {
+        <<$K as Key>::CacheSelector as CacheSelector<'tcx, $V>>::ArenaCache
+    };
+    ([$other:tt $($modifiers:tt)*][$($args:tt)*]) => {
+        query_storage!([$($modifiers)*][$($args)*])
+    };
 }
 
 macro_rules! separate_provide_extern_decl {
diff --git a/compiler/rustc_query_system/src/query/caches.rs b/compiler/rustc_query_system/src/query/caches.rs
index 0a473f91267..cdd43572422 100644
--- a/compiler/rustc_query_system/src/query/caches.rs
+++ b/compiler/rustc_query_system/src/query/caches.rs
@@ -12,6 +12,14 @@ use rustc_index::vec::{Idx, IndexVec};
 use std::default::Default;
 use std::fmt::Debug;
 use std::hash::Hash;
+use std::marker::PhantomData;
+
+pub trait CacheSelector<'tcx, V> {
+    type Cache
+    where
+        V: Clone;
+    type ArenaCache;
+}
 
 pub trait QueryStorage {
     type Value: Debug;
@@ -43,6 +51,15 @@ pub trait QueryCache: QueryStorage + Sized {
     fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex));
 }
 
+pub struct DefaultCacheSelector<K>(PhantomData<K>);
+
+impl<'tcx, K: Eq + Hash, V: 'tcx> CacheSelector<'tcx, V> for DefaultCacheSelector<K> {
+    type Cache = DefaultCache<K, V>
+    where
+        V: Clone;
+    type ArenaCache = ArenaCache<'tcx, K, V>;
+}
+
 pub struct DefaultCache<K, V> {
     #[cfg(parallel_compiler)]
     cache: Sharded<FxHashMap<K, (V, DepNodeIndex)>>,
@@ -209,6 +226,15 @@ where
     }
 }
 
+pub struct VecCacheSelector<K>(PhantomData<K>);
+
+impl<'tcx, K: Idx, V: 'tcx> CacheSelector<'tcx, V> for VecCacheSelector<K> {
+    type Cache = VecCache<K, V>
+    where
+        V: Clone;
+    type ArenaCache = VecArenaCache<'tcx, K, V>;
+}
+
 pub struct VecCache<K: Idx, V> {
     #[cfg(parallel_compiler)]
     cache: Sharded<IndexVec<K, Option<(V, DepNodeIndex)>>>,
diff --git a/compiler/rustc_query_system/src/query/mod.rs b/compiler/rustc_query_system/src/query/mod.rs
index bf14cd8de37..1d3b46aad7b 100644
--- a/compiler/rustc_query_system/src/query/mod.rs
+++ b/compiler/rustc_query_system/src/query/mod.rs
@@ -8,7 +8,7 @@ pub use self::job::{print_query_stack, QueryInfo, QueryJob, QueryJobId, QueryJob
 
 mod caches;
 pub use self::caches::{
-    ArenaCache, DefaultCache, QueryCache, QueryStorage, VecArenaCache, VecCache,
+    CacheSelector, DefaultCacheSelector, QueryCache, QueryStorage, VecCacheSelector,
 };
 
 mod config;