Access QueryStateShard directly.

Author: cjgillot

src/librustc_query_system/query/caches.rs

@@ -1,6 +1,6 @@
 use crate::dep_graph::DepNodeIndex;
 use crate::query::config::QueryContext;
-use crate::query::plumbing::{QueryLookup, QueryState, QueryStateShard};
+use crate::query::plumbing::{QueryLookup, QueryState};
 
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::sharded::Sharded;
@@ -21,19 +21,15 @@ pub trait QueryCache<CTX: QueryContext>: Default {
     /// It returns the shard index and a lock guard to the shard,
     /// which will be used if the query is not in the cache and we need
     /// to compute it.
-    fn lookup<R, GetCache, OnHit, OnMiss>(
+    fn lookup<R, OnHit, OnMiss>(
         &self,
         state: &QueryState<CTX, Self>,
-        get_cache: GetCache,
         key: Self::Key,
         // `on_hit` can be called while holding a lock to the query state shard.
         on_hit: OnHit,
         on_miss: OnMiss,
     ) -> R
     where
-        GetCache: for<'a> Fn(
-            &'a mut QueryStateShard<CTX, Self::Key, Self::Sharded>,
-        ) -> &'a mut Self::Sharded,
         OnHit: FnOnce(&Self::Value, DepNodeIndex) -> R,
         OnMiss: FnOnce(Self::Key, QueryLookup<'_, CTX, Self::Key, Self::Sharded>) -> R;
 
@@ -76,24 +72,21 @@ impl<CTX: QueryContext, K: Eq + Hash, V: Clone> QueryCache<CTX> for DefaultCache
     type Sharded = FxHashMap<K, (V, DepNodeIndex)>;
 
     #[inline(always)]
-    fn lookup<R, GetCache, OnHit, OnMiss>(
+    fn lookup<R, OnHit, OnMiss>(
         &self,
         state: &QueryState<CTX, Self>,
-        get_cache: GetCache,
         key: K,
         on_hit: OnHit,
         on_miss: OnMiss,
     ) -> R
     where
-        GetCache:
-            for<'a> Fn(&'a mut QueryStateShard<CTX, K, Self::Sharded>) -> &'a mut Self::Sharded,
         OnHit: FnOnce(&V, DepNodeIndex) -> R,
         OnMiss: FnOnce(K, QueryLookup<'_, CTX, K, Self::Sharded>) -> R,
     {
         let mut lookup = state.get_lookup(&key);
         let lock = &mut *lookup.lock;
 
-        let result = get_cache(lock).raw_entry().from_key_hashed_nocheck(lookup.key_hash, &key);
+        let result = lock.cache.raw_entry().from_key_hashed_nocheck(lookup.key_hash, &key);
 
         if let Some((_, value)) = result { on_hit(&value.0, value.1) } else { on_miss(key, lookup) }
     }

src/librustc_query_system/query/plumbing.rs

@@ -30,19 +30,13 @@ use std::ptr;
 use std::sync::atomic::{AtomicUsize, Ordering};
 
 pub struct QueryStateShard<CTX: QueryContext, K, C> {
-    cache: C,
+    pub(super) cache: C,
     active: FxHashMap<K, QueryResult<CTX>>,
 
     /// Used to generate unique ids for active jobs.
     jobs: u32,
 }
 
-impl<CTX: QueryContext, K, C> QueryStateShard<CTX, K, C> {
-    fn get_cache(&mut self) -> &mut C {
-        &mut self.cache
-    }
-}
-
 impl<CTX: QueryContext, K, C: Default> Default for QueryStateShard<CTX, K, C> {
     fn default() -> QueryStateShard<CTX, K, C> {
         QueryStateShard { cache: Default::default(), active: Default::default(), jobs: 0 }
@@ -372,7 +366,6 @@ where
 {
     state.cache.lookup(
         state,
-        QueryStateShard::<CTX, C::Key, C::Sharded>::get_cache,
         key,
         |value, index| {
             if unlikely!(tcx.profiler().enabled()) {