src/librustc/ty/query/caches.rs

   1 use crate::dep_graph::DepNodeIndex;
   2 use crate::ty::query::config::QueryContext;
   3 use crate::ty::query::plumbing::{QueryLookup, QueryState, QueryStateShard};
   4
   5 use rustc_data_structures::fx::FxHashMap;
   6 use rustc_data_structures::sharded::Sharded;
   7 use std::default::Default;
   8 use std::hash::Hash;
   9 use std::marker::PhantomData;
  10
  11 pub(crate) trait CacheSelector<CTX: QueryContext, K, V> {
  12     type Cache: QueryCache<CTX, Key = K, Value = V>;
  13 }
  14
  15 pub(crate) trait QueryCache<CTX: QueryContext>: Default {
  16     type Key;
  17     type Value;
  18     type Sharded: Default;
  19
  20     /// Checks if the query is already computed and in the cache.
  21     /// It returns the shard index and a lock guard to the shard,
  22     /// which will be used if the query is not in the cache and we need
  23     /// to compute it.
  24     fn lookup<R, GetCache, OnHit, OnMiss>(
  25         &self,
  26         state: &QueryState<CTX, Self>,
  27         get_cache: GetCache,
  28         key: Self::Key,
  29         // `on_hit` can be called while holding a lock to the query state shard.
  30         on_hit: OnHit,
  31         on_miss: OnMiss,
  32     ) -> R
  33     where
  34         GetCache: for<'a> Fn(
  35             &'a mut QueryStateShard<CTX, Self::Key, Self::Sharded>,
  36         ) -> &'a mut Self::Sharded,
  37         OnHit: FnOnce(&Self::Value, DepNodeIndex) -> R,
  38         OnMiss: FnOnce(Self::Key, QueryLookup<'_, CTX, Self::Key, Self::Sharded>) -> R;
  39
  40     fn complete(
  41         &self,
  42         tcx: CTX,
  43         lock_sharded_storage: &mut Self::Sharded,
  44         key: Self::Key,
  45         value: Self::Value,
  46         index: DepNodeIndex,
  47     );
  48
  49     fn iter<R, L>(
  50         &self,
  51         shards: &Sharded<L>,
  52         get_shard: impl Fn(&mut L) -> &mut Self::Sharded,
  53         f: impl for<'a> FnOnce(
  54             Box<dyn Iterator<Item = (&'a Self::Key, &'a Self::Value, DepNodeIndex)> + 'a>,
  55         ) -> R,
  56     ) -> R;
  57 }
  58
  59 pub struct DefaultCacheSelector;
  60
  61 impl<CTX: QueryContext, K: Eq + Hash, V: Clone> CacheSelector<CTX, K, V> for DefaultCacheSelector {
  62     type Cache = DefaultCache<K, V>;
  63 }
  64
  65 pub struct DefaultCache<K, V>(PhantomData<(K, V)>);
  66
  67 impl<K, V> Default for DefaultCache<K, V> {
  68     fn default() -> Self {
  69         DefaultCache(PhantomData)
  70     }
  71 }
  72
  73 impl<CTX: QueryContext, K: Eq + Hash, V: Clone> QueryCache<CTX> for DefaultCache<K, V> {
  74     type Key = K;
  75     type Value = V;
  76     type Sharded = FxHashMap<K, (V, DepNodeIndex)>;
  77
  78     #[inline(always)]
  79     fn lookup<R, GetCache, OnHit, OnMiss>(
  80         &self,
  81         state: &QueryState<CTX, Self>,
  82         get_cache: GetCache,
  83         key: K,
  84         on_hit: OnHit,
  85         on_miss: OnMiss,
  86     ) -> R
  87     where
  88         GetCache:
  89             for<'a> Fn(&'a mut QueryStateShard<CTX, K, Self::Sharded>) -> &'a mut Self::Sharded,
  90         OnHit: FnOnce(&V, DepNodeIndex) -> R,
  91         OnMiss: FnOnce(K, QueryLookup<'_, CTX, K, Self::Sharded>) -> R,
  92     {
  93         let mut lookup = state.get_lookup(&key);
  94         let lock = &mut *lookup.lock;
  95
  96         let result = get_cache(lock).raw_entry().from_key_hashed_nocheck(lookup.key_hash, &key);
  97
  98         if let Some((_, value)) = result { on_hit(&value.0, value.1) } else { on_miss(key, lookup) }
  99     }
 100
 101     #[inline]
 102     fn complete(
 103         &self,
 104         _: CTX,
 105         lock_sharded_storage: &mut Self::Sharded,
 106         key: K,
 107         value: V,
 108         index: DepNodeIndex,
 109     ) {
 110         lock_sharded_storage.insert(key, (value, index));
 111     }
 112
 113     fn iter<R, L>(
 114         &self,
 115         shards: &Sharded<L>,
 116         get_shard: impl Fn(&mut L) -> &mut Self::Sharded,
 117         f: impl for<'a> FnOnce(Box<dyn Iterator<Item = (&'a K, &'a V, DepNodeIndex)> + 'a>) -> R,
 118     ) -> R {
 119         let mut shards = shards.lock_shards();
 120         let mut shards: Vec<_> = shards.iter_mut().map(|shard| get_shard(shard)).collect();
 121         let results = shards.iter_mut().flat_map(|shard| shard.iter()).map(|(k, v)| (k, &v.0, v.1));
 122         f(Box::new(results))
 123     }
 124 }