compiler/rustc_query_system/src/query/caches.rs

   1 use crate::dep_graph::DepNodeIndex;
   2
   3 use rustc_arena::TypedArena;
   4 use rustc_data_structures::fx::FxHashMap;
   5 use rustc_data_structures::sharded;
   6 #[cfg(parallel_compiler)]
   7 use rustc_data_structures::sharded::Sharded;
   8 #[cfg(not(parallel_compiler))]
   9 use rustc_data_structures::sync::Lock;
  10 use rustc_data_structures::sync::WorkerLocal;
  11 use rustc_index::vec::{Idx, IndexVec};
  12 use std::default::Default;
  13 use std::fmt::Debug;
  14 use std::hash::Hash;
  15 use std::marker::PhantomData;
  16
  17 pub trait CacheSelector<'tcx, V> {
  18     type Cache
  19     where
  20         V: Clone;
  21     type ArenaCache;
  22 }
  23
  24 pub trait QueryStorage {
  25     type Value: Debug;
  26     type Stored: Clone;
  27
  28     /// Store a value without putting it in the cache.
  29     /// This is meant to be used with cycle errors.
  30     fn store_nocache(&self, value: Self::Value) -> Self::Stored;
  31 }
  32
  33 pub trait QueryCache: QueryStorage + Sized {
  34     type Key: Hash + Eq + Clone + Debug;
  35
  36     /// Checks if the query is already computed and in the cache.
  37     /// It returns the shard index and a lock guard to the shard,
  38     /// which will be used if the query is not in the cache and we need
  39     /// to compute it.
  40     fn lookup<R, OnHit>(
  41         &self,
  42         key: &Self::Key,
  43         // `on_hit` can be called while holding a lock to the query state shard.
  44         on_hit: OnHit,
  45     ) -> Result<R, ()>
  46     where
  47         OnHit: FnOnce(&Self::Stored, DepNodeIndex) -> R;
  48
  49     fn complete(&self, key: Self::Key, value: Self::Value, index: DepNodeIndex) -> Self::Stored;
  50
  51     fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex));
  52 }
  53
  54 pub struct DefaultCacheSelector<K>(PhantomData<K>);
  55
  56 impl<'tcx, K: Eq + Hash, V: 'tcx> CacheSelector<'tcx, V> for DefaultCacheSelector<K> {
  57     type Cache = DefaultCache<K, V>
  58     where
  59         V: Clone;
  60     type ArenaCache = ArenaCache<'tcx, K, V>;
  61 }
  62
  63 pub struct DefaultCache<K, V> {
  64     #[cfg(parallel_compiler)]
  65     cache: Sharded<FxHashMap<K, (V, DepNodeIndex)>>,
  66     #[cfg(not(parallel_compiler))]
  67     cache: Lock<FxHashMap<K, (V, DepNodeIndex)>>,
  68 }
  69
  70 impl<K, V> Default for DefaultCache<K, V> {
  71     fn default() -> Self {
  72         DefaultCache { cache: Default::default() }
  73     }
  74 }
  75
  76 impl<K: Eq + Hash, V: Clone + Debug> QueryStorage for DefaultCache<K, V> {
  77     type Value = V;
  78     type Stored = V;
  79
  80     #[inline]
  81     fn store_nocache(&self, value: Self::Value) -> Self::Stored {
  82         // We have no dedicated storage
  83         value
  84     }
  85 }
  86
  87 impl<K, V> QueryCache for DefaultCache<K, V>
  88 where
  89     K: Eq + Hash + Clone + Debug,
  90     V: Clone + Debug,
  91 {
  92     type Key = K;
  93
  94     #[inline(always)]
  95     fn lookup<R, OnHit>(&self, key: &K, on_hit: OnHit) -> Result<R, ()>
  96     where
  97         OnHit: FnOnce(&V, DepNodeIndex) -> R,
  98     {
  99         let key_hash = sharded::make_hash(key);
 100         #[cfg(parallel_compiler)]
 101         let lock = self.cache.get_shard_by_hash(key_hash).lock();
 102         #[cfg(not(parallel_compiler))]
 103         let lock = self.cache.lock();
 104         let result = lock.raw_entry().from_key_hashed_nocheck(key_hash, key);
 105
 106         if let Some((_, value)) = result {
 107             let hit_result = on_hit(&value.0, value.1);
 108             Ok(hit_result)
 109         } else {
 110             Err(())
 111         }
 112     }
 113
 114     #[inline]
 115     fn complete(&self, key: K, value: V, index: DepNodeIndex) -> Self::Stored {
 116         #[cfg(parallel_compiler)]
 117         let mut lock = self.cache.get_shard_by_value(&key).lock();
 118         #[cfg(not(parallel_compiler))]
 119         let mut lock = self.cache.lock();
 120         lock.insert(key, (value.clone(), index));
 121         value
 122     }
 123
 124     fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
 125         #[cfg(parallel_compiler)]
 126         {
 127             let shards = self.cache.lock_shards();
 128             for shard in shards.iter() {
 129                 for (k, v) in shard.iter() {
 130                     f(k, &v.0, v.1);
 131                 }
 132             }
 133         }
 134         #[cfg(not(parallel_compiler))]
 135         {
 136             let map = self.cache.lock();
 137             for (k, v) in map.iter() {
 138                 f(k, &v.0, v.1);
 139             }
 140         }
 141     }
 142 }
 143
 144 pub struct ArenaCache<'tcx, K, V> {
 145     arena: WorkerLocal<TypedArena<(V, DepNodeIndex)>>,
 146     #[cfg(parallel_compiler)]
 147     cache: Sharded<FxHashMap<K, &'tcx (V, DepNodeIndex)>>,
 148     #[cfg(not(parallel_compiler))]
 149     cache: Lock<FxHashMap<K, &'tcx (V, DepNodeIndex)>>,
 150 }
 151
 152 impl<'tcx, K, V> Default for ArenaCache<'tcx, K, V> {
 153     fn default() -> Self {
 154         ArenaCache { arena: WorkerLocal::new(|_| TypedArena::default()), cache: Default::default() }
 155     }
 156 }
 157
 158 impl<'tcx, K: Eq + Hash, V: Debug + 'tcx> QueryStorage for ArenaCache<'tcx, K, V> {
 159     type Value = V;
 160     type Stored = &'tcx V;
 161
 162     #[inline]
 163     fn store_nocache(&self, value: Self::Value) -> Self::Stored {
 164         let value = self.arena.alloc((value, DepNodeIndex::INVALID));
 165         let value = unsafe { &*(&value.0 as *const _) };
 166         &value
 167     }
 168 }
 169
 170 impl<'tcx, K, V: 'tcx> QueryCache for ArenaCache<'tcx, K, V>
 171 where
 172     K: Eq + Hash + Clone + Debug,
 173     V: Debug,
 174 {
 175     type Key = K;
 176
 177     #[inline(always)]
 178     fn lookup<R, OnHit>(&self, key: &K, on_hit: OnHit) -> Result<R, ()>
 179     where
 180         OnHit: FnOnce(&&'tcx V, DepNodeIndex) -> R,
 181     {
 182         let key_hash = sharded::make_hash(key);
 183         #[cfg(parallel_compiler)]
 184         let lock = self.cache.get_shard_by_hash(key_hash).lock();
 185         #[cfg(not(parallel_compiler))]
 186         let lock = self.cache.lock();
 187         let result = lock.raw_entry().from_key_hashed_nocheck(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(())
 194         }
 195     }
 196
 197     #[inline]
 198     fn complete(&self, key: K, value: V, index: DepNodeIndex) -> Self::Stored {
 199         let value = self.arena.alloc((value, index));
 200         let value = unsafe { &*(value as *const _) };
 201         #[cfg(parallel_compiler)]
 202         let mut lock = self.cache.get_shard_by_value(&key).lock();
 203         #[cfg(not(parallel_compiler))]
 204         let mut lock = self.cache.lock();
 205         lock.insert(key, value);
 206         &value.0
 207     }
 208
 209     fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
 210         #[cfg(parallel_compiler)]
 211         {
 212             let shards = self.cache.lock_shards();
 213             for shard in shards.iter() {
 214                 for (k, v) in shard.iter() {
 215                     f(k, &v.0, v.1);
 216                 }
 217             }
 218         }
 219         #[cfg(not(parallel_compiler))]
 220         {
 221             let map = self.cache.lock();
 222             for (k, v) in map.iter() {
 223                 f(k, &v.0, v.1);
 224             }
 225         }
 226     }
 227 }
 228
 229 pub struct VecCacheSelector<K>(PhantomData<K>);
 230
 231 impl<'tcx, K: Idx, V: 'tcx> CacheSelector<'tcx, V> for VecCacheSelector<K> {
 232     type Cache = VecCache<K, V>
 233     where
 234         V: Clone;
 235     type ArenaCache = VecArenaCache<'tcx, K, V>;
 236 }
 237
 238 pub struct VecCache<K: Idx, V> {
 239     #[cfg(parallel_compiler)]
 240     cache: Sharded<IndexVec<K, Option<(V, DepNodeIndex)>>>,
 241     #[cfg(not(parallel_compiler))]
 242     cache: Lock<IndexVec<K, Option<(V, DepNodeIndex)>>>,
 243 }
 244
 245 impl<K: Idx, V> Default for VecCache<K, V> {
 246     fn default() -> Self {
 247         VecCache { cache: Default::default() }
 248     }
 249 }
 250
 251 impl<K: Eq + Idx, V: Clone + Debug> QueryStorage for VecCache<K, V> {
 252     type Value = V;
 253     type Stored = V;
 254
 255     #[inline]
 256     fn store_nocache(&self, value: Self::Value) -> Self::Stored {
 257         // We have no dedicated storage
 258         value
 259     }
 260 }
 261
 262 impl<K, V> QueryCache for VecCache<K, V>
 263 where
 264     K: Eq + Idx + Clone + Debug,
 265     V: Clone + Debug,
 266 {
 267     type Key = K;
 268
 269     #[inline(always)]
 270     fn lookup<R, OnHit>(&self, key: &K, on_hit: OnHit) -> Result<R, ()>
 271     where
 272         OnHit: FnOnce(&V, DepNodeIndex) -> R,
 273     {
 274         #[cfg(parallel_compiler)]
 275         let lock = self.cache.get_shard_by_hash(key.index() as u64).lock();
 276         #[cfg(not(parallel_compiler))]
 277         let lock = self.cache.lock();
 278         if let Some(Some(value)) = lock.get(*key) {
 279             let hit_result = on_hit(&value.0, value.1);
 280             Ok(hit_result)
 281         } else {
 282             Err(())
 283         }
 284     }
 285
 286     #[inline]
 287     fn complete(&self, key: K, value: V, index: DepNodeIndex) -> Self::Stored {
 288         #[cfg(parallel_compiler)]
 289         let mut lock = self.cache.get_shard_by_hash(key.index() as u64).lock();
 290         #[cfg(not(parallel_compiler))]
 291         let mut lock = self.cache.lock();
 292         lock.insert(key, (value.clone(), index));
 293         value
 294     }
 295
 296     fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
 297         #[cfg(parallel_compiler)]
 298         {
 299             let shards = self.cache.lock_shards();
 300             for shard in shards.iter() {
 301                 for (k, v) in shard.iter_enumerated() {
 302                     if let Some(v) = v {
 303                         f(&k, &v.0, v.1);
 304                     }
 305                 }
 306             }
 307         }
 308         #[cfg(not(parallel_compiler))]
 309         {
 310             let map = self.cache.lock();
 311             for (k, v) in map.iter_enumerated() {
 312                 if let Some(v) = v {
 313                     f(&k, &v.0, v.1);
 314                 }
 315             }
 316         }
 317     }
 318 }
 319
 320 pub struct VecArenaCache<'tcx, K: Idx, V> {
 321     arena: WorkerLocal<TypedArena<(V, DepNodeIndex)>>,
 322     #[cfg(parallel_compiler)]
 323     cache: Sharded<IndexVec<K, Option<&'tcx (V, DepNodeIndex)>>>,
 324     #[cfg(not(parallel_compiler))]
 325     cache: Lock<IndexVec<K, Option<&'tcx (V, DepNodeIndex)>>>,
 326 }
 327
 328 impl<'tcx, K: Idx, V> Default for VecArenaCache<'tcx, K, V> {
 329     fn default() -> Self {
 330         VecArenaCache {
 331             arena: WorkerLocal::new(|_| TypedArena::default()),
 332             cache: Default::default(),
 333         }
 334     }
 335 }
 336
 337 impl<'tcx, K: Eq + Idx, V: Debug + 'tcx> QueryStorage for VecArenaCache<'tcx, K, V> {
 338     type Value = V;
 339     type Stored = &'tcx V;
 340
 341     #[inline]
 342     fn store_nocache(&self, value: Self::Value) -> Self::Stored {
 343         let value = self.arena.alloc((value, DepNodeIndex::INVALID));
 344         let value = unsafe { &*(&value.0 as *const _) };
 345         &value
 346     }
 347 }
 348
 349 impl<'tcx, K, V: 'tcx> QueryCache for VecArenaCache<'tcx, K, V>
 350 where
 351     K: Eq + Idx + Clone + Debug,
 352     V: Debug,
 353 {
 354     type Key = K;
 355
 356     #[inline(always)]
 357     fn lookup<R, OnHit>(&self, key: &K, on_hit: OnHit) -> Result<R, ()>
 358     where
 359         OnHit: FnOnce(&&'tcx V, DepNodeIndex) -> R,
 360     {
 361         #[cfg(parallel_compiler)]
 362         let lock = self.cache.get_shard_by_hash(key.index() as u64).lock();
 363         #[cfg(not(parallel_compiler))]
 364         let lock = self.cache.lock();
 365         if let Some(Some(value)) = lock.get(*key) {
 366             let hit_result = on_hit(&&value.0, value.1);
 367             Ok(hit_result)
 368         } else {
 369             Err(())
 370         }
 371     }
 372
 373     #[inline]
 374     fn complete(&self, key: K, value: V, index: DepNodeIndex) -> Self::Stored {
 375         let value = self.arena.alloc((value, index));
 376         let value = unsafe { &*(value as *const _) };
 377         #[cfg(parallel_compiler)]
 378         let mut lock = self.cache.get_shard_by_hash(key.index() as u64).lock();
 379         #[cfg(not(parallel_compiler))]
 380         let mut lock = self.cache.lock();
 381         lock.insert(key, value);
 382         &value.0
 383     }
 384
 385     fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
 386         #[cfg(parallel_compiler)]
 387         {
 388             let shards = self.cache.lock_shards();
 389             for shard in shards.iter() {
 390                 for (k, v) in shard.iter_enumerated() {
 391                     if let Some(v) = v {
 392                         f(&k, &v.0, v.1);
 393                     }
 394                 }
 395             }
 396         }
 397         #[cfg(not(parallel_compiler))]
 398         {
 399             let map = self.cache.lock();
 400             for (k, v) in map.iter_enumerated() {
 401                 if let Some(v) = v {
 402                     f(&k, &v.0, v.1);
 403                 }
 404             }
 405         }
 406     }
 407 }