compiler/rustc_query_system/src/query/caches.rs

   1 use crate::dep_graph::DepNodeIndex;
   2 use crate::query::plumbing::{QueryCacheStore, QueryLookup};
   3
   4 use rustc_arena::TypedArena;
   5 use rustc_data_structures::fx::FxHashMap;
   6 use rustc_data_structures::sharded::Sharded;
   7 use rustc_data_structures::sync::WorkerLocal;
   8 use std::default::Default;
   9 use std::fmt::Debug;
  10 use std::hash::Hash;
  11 use std::marker::PhantomData;
  12
  13 pub trait CacheSelector<K, V> {
  14     type Cache;
  15 }
  16
  17 pub trait QueryStorage {
  18     type Value: Debug;
  19     type Stored: Clone;
  20
  21     /// Store a value without putting it in the cache.
  22     /// This is meant to be used with cycle errors.
  23     fn store_nocache(&self, value: Self::Value) -> Self::Stored;
  24 }
  25
  26 pub trait QueryCache: QueryStorage + Sized {
  27     type Key: Hash + Eq + Clone + Debug;
  28     type Sharded: Default;
  29
  30     /// Checks if the query is already computed and in the cache.
  31     /// It returns the shard index and a lock guard to the shard,
  32     /// which will be used if the query is not in the cache and we need
  33     /// to compute it.
  34     fn lookup<'s, R, OnHit>(
  35         &self,
  36         state: &'s QueryCacheStore<Self>,
  37         key: &Self::Key,
  38         // `on_hit` can be called while holding a lock to the query state shard.
  39         on_hit: OnHit,
  40     ) -> Result<R, QueryLookup>
  41     where
  42         OnHit: FnOnce(&Self::Stored, DepNodeIndex) -> R;
  43
  44     fn complete(
  45         &self,
  46         lock_sharded_storage: &mut Self::Sharded,
  47         key: Self::Key,
  48         value: Self::Value,
  49         index: DepNodeIndex,
  50     ) -> Self::Stored;
  51
  52     fn iter(
  53         &self,
  54         shards: &Sharded<Self::Sharded>,
  55         f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex),
  56     );
  57 }
  58
  59 pub struct DefaultCacheSelector;
  60
  61 impl<K: Eq + Hash, V: Clone> CacheSelector<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<K: Eq + Hash, V: Clone + Debug> QueryStorage for DefaultCache<K, V> {
  74     type Value = V;
  75     type Stored = V;
  76
  77     #[inline]
  78     fn store_nocache(&self, value: Self::Value) -> Self::Stored {
  79         // We have no dedicated storage
  80         value
  81     }
  82 }
  83
  84 impl<K, V> QueryCache for DefaultCache<K, V>
  85 where
  86     K: Eq + Hash + Clone + Debug,
  87     V: Clone + Debug,
  88 {
  89     type Key = K;
  90     type Sharded = FxHashMap<K, (V, DepNodeIndex)>;
  91
  92     #[inline(always)]
  93     fn lookup<'s, R, OnHit>(
  94         &self,
  95         state: &'s QueryCacheStore<Self>,
  96         key: &K,
  97         on_hit: OnHit,
  98     ) -> Result<R, QueryLookup>
  99     where
 100         OnHit: FnOnce(&V, DepNodeIndex) -> R,
 101     {
 102         let (lookup, lock) = state.get_lookup(key);
 103         let result = lock.raw_entry().from_key_hashed_nocheck(lookup.key_hash, key);
 104
 105         if let Some((_, value)) = result {
 106             let hit_result = on_hit(&value.0, value.1);
 107             Ok(hit_result)
 108         } else {
 109             Err(lookup)
 110         }
 111     }
 112
 113     #[inline]
 114     fn complete(
 115         &self,
 116         lock_sharded_storage: &mut Self::Sharded,
 117         key: K,
 118         value: V,
 119         index: DepNodeIndex,
 120     ) -> Self::Stored {
 121         lock_sharded_storage.insert(key, (value.clone(), index));
 122         value
 123     }
 124
 125     fn iter(
 126         &self,
 127         shards: &Sharded<Self::Sharded>,
 128         f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex),
 129     ) {
 130         let shards = shards.lock_shards();
 131         for shard in shards.iter() {
 132             for (k, v) in shard.iter() {
 133                 f(k, &v.0, v.1);
 134             }
 135         }
 136     }
 137 }
 138
 139 pub struct ArenaCacheSelector<'tcx>(PhantomData<&'tcx ()>);
 140
 141 impl<'tcx, K: Eq + Hash, V: 'tcx> CacheSelector<K, V> for ArenaCacheSelector<'tcx> {
 142     type Cache = ArenaCache<'tcx, K, V>;
 143 }
 144
 145 pub struct ArenaCache<'tcx, K, V> {
 146     arena: WorkerLocal<TypedArena<(V, DepNodeIndex)>>,
 147     phantom: PhantomData<(K, &'tcx V)>,
 148 }
 149
 150 impl<'tcx, K, V> Default for ArenaCache<'tcx, K, V> {
 151     fn default() -> Self {
 152         ArenaCache { arena: WorkerLocal::new(|_| TypedArena::default()), phantom: PhantomData }
 153     }
 154 }
 155
 156 impl<'tcx, K: Eq + Hash, V: Debug + 'tcx> QueryStorage for ArenaCache<'tcx, K, V> {
 157     type Value = V;
 158     type Stored = &'tcx V;
 159
 160     #[inline]
 161     fn store_nocache(&self, value: Self::Value) -> Self::Stored {
 162         let value = self.arena.alloc((value, DepNodeIndex::INVALID));
 163         let value = unsafe { &*(&value.0 as *const _) };
 164         &value
 165     }
 166 }
 167
 168 impl<'tcx, K, V: 'tcx> QueryCache for ArenaCache<'tcx, K, V>
 169 where
 170     K: Eq + Hash + Clone + Debug,
 171     V: Debug,
 172 {
 173     type Key = K;
 174     type Sharded = FxHashMap<K, &'tcx (V, DepNodeIndex)>;
 175
 176     #[inline(always)]
 177     fn lookup<'s, R, OnHit>(
 178         &self,
 179         state: &'s QueryCacheStore<Self>,
 180         key: &K,
 181         on_hit: OnHit,
 182     ) -> Result<R, QueryLookup>
 183     where
 184         OnHit: FnOnce(&&'tcx V, DepNodeIndex) -> R,
 185     {
 186         let (lookup, lock) = state.get_lookup(key);
 187         let result = lock.raw_entry().from_key_hashed_nocheck(lookup.key_hash, key);
 188
 189         if let Some((_, value)) = result {
 190             let hit_result = on_hit(&&value.0, value.1);
 191             Ok(hit_result)
 192         } else {
 193             Err(lookup)
 194         }
 195     }
 196
 197     #[inline]
 198     fn complete(
 199         &self,
 200         lock_sharded_storage: &mut Self::Sharded,
 201         key: K,
 202         value: V,
 203         index: DepNodeIndex,
 204     ) -> Self::Stored {
 205         let value = self.arena.alloc((value, index));
 206         let value = unsafe { &*(value as *const _) };
 207         lock_sharded_storage.insert(key, value);
 208         &value.0
 209     }
 210
 211     fn iter(
 212         &self,
 213         shards: &Sharded<Self::Sharded>,
 214         f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex),
 215     ) {
 216         let shards = shards.lock_shards();
 217         for shard in shards.iter() {
 218             for (k, v) in shard.iter() {
 219                 f(k, &v.0, v.1);
 220             }
 221         }
 222     }
 223 }