Introduce an abstraction for EvaluationCache and SelectionCache.

Author: cjgillot

src/librustc_middle/traits/select.rs

@@ -6,29 +6,18 @@ use self::EvaluationResult::*;
 
 use super::{SelectionError, SelectionResult};
 
-use crate::dep_graph::DepNodeIndex;
-use crate::ty::{self, TyCtxt};
+use crate::ty;
 
-use rustc_data_structures::fx::FxHashMap;
-use rustc_data_structures::sync::Lock;
 use rustc_hir::def_id::DefId;
+use rustc_query_system::cache::Cache;
 
-#[derive(Clone, Default)]
-pub struct SelectionCache<'tcx> {
-    pub hashmap: Lock<
-        FxHashMap<
-            ty::ParamEnvAnd<'tcx, ty::TraitRef<'tcx>>,
-            WithDepNode<SelectionResult<'tcx, SelectionCandidate<'tcx>>>,
-        >,
-    >,
-}
+pub type SelectionCache<'tcx> = Cache<
+    ty::ParamEnvAnd<'tcx, ty::TraitRef<'tcx>>,
+    SelectionResult<'tcx, SelectionCandidate<'tcx>>,
+>;
 
-impl<'tcx> SelectionCache<'tcx> {
-    /// Actually frees the underlying memory in contrast to what stdlib containers do on `clear`
-    pub fn clear(&self) {
-        *self.hashmap.borrow_mut() = Default::default();
-    }
-}
+pub type EvaluationCache<'tcx> =
+    Cache<ty::ParamEnvAnd<'tcx, ty::PolyTraitRef<'tcx>>, EvaluationResult>;
 
 /// The selection process begins by considering all impls, where
 /// clauses, and so forth that might resolve an obligation. Sometimes
@@ -265,37 +254,6 @@ impl<'tcx> From<OverflowError> for SelectionError<'tcx> {
     }
 }
 
-#[derive(Clone, Default)]
-pub struct EvaluationCache<'tcx> {
-    pub hashmap: Lock<
-        FxHashMap<ty::ParamEnvAnd<'tcx, ty::PolyTraitRef<'tcx>>, WithDepNode<EvaluationResult>>,
-    >,
-}
-
-impl<'tcx> EvaluationCache<'tcx> {
-    /// Actually frees the underlying memory in contrast to what stdlib containers do on `clear`
-    pub fn clear(&self) {
-        *self.hashmap.borrow_mut() = Default::default();
-    }
-}
-
-#[derive(Clone, Eq, PartialEq)]
-pub struct WithDepNode<T> {
-    dep_node: DepNodeIndex,
-    cached_value: T,
-}
-
-impl<T: Clone> WithDepNode<T> {
-    pub fn new(dep_node: DepNodeIndex, cached_value: T) -> Self {
-        WithDepNode { dep_node, cached_value }
-    }
-
-    pub fn get(&self, tcx: TyCtxt<'_>) -> T {
-        tcx.dep_graph.read_index(self.dep_node);
-        self.cached_value.clone()
-    }
-}
-
 #[derive(Clone, Debug)]
 pub enum IntercrateAmbiguityCause {
     DownstreamCrate { trait_desc: String, self_desc: Option<String> },

src/librustc_query_system/cache.rs

@@ -0,0 +1,62 @@
+//! Cache for candidate selection.
+
+use crate::dep_graph::DepNodeIndex;
+use crate::query::QueryContext;
+
+use rustc_data_structures::fx::FxHashMap;
+use rustc_data_structures::sync::HashMapExt;
+use rustc_data_structures::sync::Lock;
+
+use std::hash::Hash;
+
+#[derive(Clone)]
+pub struct Cache<Key, Value> {
+    hashmap: Lock<FxHashMap<Key, WithDepNode<Value>>>,
+}
+
+impl<Key, Value> Default for Cache<Key, Value> {
+    fn default() -> Self {
+        Self { hashmap: Default::default() }
+    }
+}
+
+impl<Key, Value> Cache<Key, Value> {
+    /// Actually frees the underlying memory in contrast to what stdlib containers do on `clear`
+    pub fn clear(&self) {
+        *self.hashmap.borrow_mut() = Default::default();
+    }
+}
+
+impl<Key: Eq + Hash, Value: Clone> Cache<Key, Value> {
+    pub fn get<CTX: QueryContext>(&self, key: &Key, tcx: CTX) -> Option<Value> {
+        Some(self.hashmap.borrow().get(key)?.get(tcx))
+    }
+
+    pub fn insert(&self, key: Key, dep_node: DepNodeIndex, value: Value) {
+        self.hashmap.borrow_mut().insert(key, WithDepNode::new(dep_node, value));
+    }
+
+    pub fn insert_same(&self, key: Key, dep_node: DepNodeIndex, value: Value)
+    where
+        Value: Eq,
+    {
+        self.hashmap.borrow_mut().insert_same(key, WithDepNode::new(dep_node, value));
+    }
+}
+
+#[derive(Clone, Eq, PartialEq)]
+pub struct WithDepNode<T> {
+    dep_node: DepNodeIndex,
+    cached_value: T,
+}
+
+impl<T: Clone> WithDepNode<T> {
+    pub fn new(dep_node: DepNodeIndex, cached_value: T) -> Self {
+        WithDepNode { dep_node, cached_value }
+    }
+
+    pub fn get<CTX: QueryContext>(&self, tcx: CTX) -> T {
+        tcx.dep_graph().read_index(self.dep_node);
+        self.cached_value.clone()
+    }
+}

src/librustc_query_system/lib.rs

@@ -11,5 +11,6 @@ extern crate log;
 #[macro_use]
 extern crate rustc_data_structures;
 
+pub mod cache;
 pub mod dep_graph;
 pub mod query;

src/librustc_trait_selection/traits/select/mod.rs

@@ -833,17 +833,11 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
     ) -> Option<EvaluationResult> {
         let tcx = self.tcx();
         if self.can_use_global_caches(param_env) {
-            let cache = tcx.evaluation_cache.hashmap.borrow();
-            if let Some(cached) = cache.get(&param_env.and(trait_ref)) {
-                return Some(cached.get(tcx));
+            if let Some(res) = tcx.evaluation_cache.get(&param_env.and(trait_ref), tcx) {
+                return Some(res);
             }
         }
-        self.infcx
-            .evaluation_cache
-            .hashmap
-            .borrow()
-            .get(&param_env.and(trait_ref))
-            .map(|v| v.get(tcx))
+        self.infcx.evaluation_cache.get(&param_env.and(trait_ref), tcx)
     }
 
     fn insert_evaluation_cache(
@@ -869,21 +863,13 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
                 // FIXME: Due to #50507 this overwrites the different values
                 // This should be changed to use HashMapExt::insert_same
                 // when that is fixed
-                self.tcx()
-                    .evaluation_cache
-                    .hashmap
-                    .borrow_mut()
-                    .insert(param_env.and(trait_ref), WithDepNode::new(dep_node, result));
+                self.tcx().evaluation_cache.insert(param_env.and(trait_ref), dep_node, result);
                 return;
             }
         }
 
         debug!("insert_evaluation_cache(trait_ref={:?}, candidate={:?})", trait_ref, result,);
-        self.infcx
-            .evaluation_cache
-            .hashmap
-            .borrow_mut()
-            .insert(param_env.and(trait_ref), WithDepNode::new(dep_node, result));
+        self.infcx.evaluation_cache.insert(param_env.and(trait_ref), dep_node, result);
     }
 
     /// For various reasons, it's possible for a subobligation
@@ -1180,17 +1166,11 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
         let tcx = self.tcx();
         let trait_ref = &cache_fresh_trait_pred.skip_binder().trait_ref;
         if self.can_use_global_caches(param_env) {
-            let cache = tcx.selection_cache.hashmap.borrow();
-            if let Some(cached) = cache.get(&param_env.and(*trait_ref)) {
-                return Some(cached.get(tcx));
+            if let Some(res) = tcx.selection_cache.get(&param_env.and(*trait_ref), tcx) {
+                return Some(res);
             }
         }
-        self.infcx
-            .selection_cache
-            .hashmap
-            .borrow()
-            .get(&param_env.and(*trait_ref))
-            .map(|v| v.get(tcx))
+        self.infcx.selection_cache.get(&param_env.and(*trait_ref), tcx)
     }
 
     /// Determines whether can we safely cache the result
@@ -1248,10 +1228,7 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
                         trait_ref, candidate,
                     );
                     // This may overwrite the cache with the same value.
-                    tcx.selection_cache
-                        .hashmap
-                        .borrow_mut()
-                        .insert(param_env.and(trait_ref), WithDepNode::new(dep_node, candidate));
+                    tcx.selection_cache.insert(param_env.and(trait_ref), dep_node, candidate);
                     return;
                 }
             }
@@ -1261,11 +1238,7 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
             "insert_candidate_cache(trait_ref={:?}, candidate={:?}) local",
             trait_ref, candidate,
         );
-        self.infcx
-            .selection_cache
-            .hashmap
-            .borrow_mut()
-            .insert(param_env.and(trait_ref), WithDepNode::new(dep_node, candidate));
+        self.infcx.selection_cache.insert(param_env.and(trait_ref), dep_node, candidate);
     }
 
     fn match_projection_obligation_against_definition_bounds(