Return a Result for query cache.

Author: cjgillot

compiler/rustc_query_system/src/query/caches.rs

@@ -31,17 +31,15 @@ pub trait QueryCache: QueryStorage {
     /// 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<D, Q, R, OnHit, OnMiss>(
+    fn lookup<'s, D, Q, R, OnHit>(
         &self,
-        state: &QueryState<D, Q, Self>,
-        key: Self::Key,
+        state: &'s QueryState<D, Q, Self>,
+        key: &Self::Key,
         // `on_hit` can be called while holding a lock to the query state shard.
         on_hit: OnHit,
-        on_miss: OnMiss,
-    ) -> R
+    ) -> Result<R, QueryLookup<'s, D, Q, Self::Key, Self::Sharded>>
     where
-        OnHit: FnOnce(&Self::Stored, DepNodeIndex) -> R,
-        OnMiss: FnOnce(Self::Key, QueryLookup<'_, D, Q, Self::Key, Self::Sharded>) -> R;
+        OnHit: FnOnce(&Self::Stored, DepNodeIndex) -> R;
 
     fn complete(
         &self,
@@ -95,23 +93,24 @@ where
     type Sharded = FxHashMap<K, (V, DepNodeIndex)>;
 
     #[inline(always)]
-    fn lookup<D, Q, R, OnHit, OnMiss>(
+    fn lookup<'s, D, Q, R, OnHit>(
         &self,
-        state: &QueryState<D, Q, Self>,
-        key: K,
+        state: &'s QueryState<D, Q, Self>,
+        key: &K,
         on_hit: OnHit,
-        on_miss: OnMiss,
-    ) -> R
+    ) -> Result<R, QueryLookup<'s, D, Q, K, Self::Sharded>>
     where
         OnHit: FnOnce(&V, DepNodeIndex) -> R,
-        OnMiss: FnOnce(K, QueryLookup<'_, D, Q, K, Self::Sharded>) -> R,
     {
-        let mut lookup = state.get_lookup(&key);
-        let lock = &mut *lookup.lock;
+        let lookup = state.get_lookup(key);
+        let result = lookup.lock.cache.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) }
+        if let Some((_, value)) = result {
+            let hit_result = on_hit(&value.0, value.1);
+            Ok(hit_result)
+        } else {
+            Err(lookup)
+        }
     }
 
     #[inline]
@@ -177,26 +176,23 @@ where
     type Sharded = FxHashMap<K, &'tcx (V, DepNodeIndex)>;
 
     #[inline(always)]
-    fn lookup<D, Q, R, OnHit, OnMiss>(
+    fn lookup<'s, D, Q, R, OnHit>(
         &self,
-        state: &QueryState<D, Q, Self>,
-        key: K,
+        state: &'s QueryState<D, Q, Self>,
+        key: &K,
         on_hit: OnHit,
-        on_miss: OnMiss,
-    ) -> R
+    ) -> Result<R, QueryLookup<'s, D, Q, K, Self::Sharded>>
     where
         OnHit: FnOnce(&&'tcx V, DepNodeIndex) -> R,
-        OnMiss: FnOnce(K, QueryLookup<'_, D, Q, K, Self::Sharded>) -> R,
     {
-        let mut lookup = state.get_lookup(&key);
-        let lock = &mut *lookup.lock;
-
-        let result = lock.cache.raw_entry().from_key_hashed_nocheck(lookup.key_hash, &key);
+        let lookup = state.get_lookup(key);
+        let result = lookup.lock.cache.raw_entry().from_key_hashed_nocheck(lookup.key_hash, key);
 
         if let Some((_, value)) = result {
-            on_hit(&&value.0, value.1)
+            let hit_result = on_hit(&&value.0, value.1);
+            Ok(hit_result)
         } else {
-            on_miss(key, lookup)
+            Err(lookup)
         }
     }
 

compiler/rustc_query_system/src/query/plumbing.rs

@@ -248,13 +248,8 @@ where
                 return TryGetJob::Cycle(value);
             }
 
-            let cached = try_get_cached(
-                tcx,
-                state,
-                (*key).clone(),
-                |value, index| (value.clone(), index),
-                |_, _| panic!("value must be in cache after waiting"),
-            );
+            let cached = try_get_cached(tcx, state, key, |value, index| (value.clone(), index))
+                .unwrap_or_else(|_| panic!("value must be in cache after waiting"));
 
             if let Some(prof_timer) = _query_blocked_prof_timer.take() {
                 prof_timer.finish_with_query_invocation_id(cached.1.into());
@@ -356,35 +351,28 @@ where
 /// 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 try_get_cached<CTX, C, R, OnHit, OnMiss>(
+fn try_get_cached<'a, CTX, C, R, OnHit>(
     tcx: CTX,
-    state: &QueryState<CTX::DepKind, CTX::Query, C>,
-    key: C::Key,
+    state: &'a QueryState<CTX::DepKind, CTX::Query, C>,
+    key: &C::Key,
     // `on_hit` can be called while holding a lock to the query cache
     on_hit: OnHit,
-    on_miss: OnMiss,
-) -> R
+) -> Result<R, QueryLookup<'a, CTX::DepKind, CTX::Query, C::Key, C::Sharded>>
 where
     C: QueryCache,
     CTX: QueryContext,
     OnHit: FnOnce(&C::Stored, DepNodeIndex) -> R,
-    OnMiss: FnOnce(C::Key, QueryLookup<'_, CTX::DepKind, CTX::Query, C::Key, C::Sharded>) -> R,
 {
-    state.cache.lookup(
-        state,
-        key,
-        |value, index| {
-            if unlikely!(tcx.profiler().enabled()) {
-                tcx.profiler().query_cache_hit(index.into());
-            }
-            #[cfg(debug_assertions)]
-            {
-                state.cache_hits.fetch_add(1, Ordering::Relaxed);
-            }
-            on_hit(value, index)
-        },
-        on_miss,
-    )
+    state.cache.lookup(state, &key, |value, index| {
+        if unlikely!(tcx.profiler().enabled()) {
+            tcx.profiler().query_cache_hit(index.into());
+        }
+        #[cfg(debug_assertions)]
+        {
+            state.cache_hits.fetch_add(1, Ordering::Relaxed);
+        }
+        on_hit(value, index)
+    })
 }
 
 fn try_execute_query<CTX, C>(
@@ -626,16 +614,14 @@ where
     C: QueryCache,
     C::Key: crate::dep_graph::DepNodeParams<CTX>,
 {
-    try_get_cached(
-        tcx,
-        state,
-        key,
-        |value, index| {
-            tcx.dep_graph().read_index(index);
-            value.clone()
-        },
-        |key, lookup| try_execute_query(tcx, state, span, key, lookup, query),
-    )
+    let cached = try_get_cached(tcx, state, &key, |value, index| {
+        tcx.dep_graph().read_index(index);
+        value.clone()
+    });
+    match cached {
+        Ok(value) => value,
+        Err(lookup) => try_execute_query(tcx, state, span, key, lookup, query),
+    }
 }
 
 /// Ensure that either this query has all green inputs or been executed.
@@ -694,25 +680,24 @@ fn force_query_impl<CTX, C>(
     // We may be concurrently trying both execute and force a query.
     // Ensure that only one of them runs the query.
 
-    try_get_cached(
-        tcx,
-        state,
-        key,
-        |_, _| {
-            // Cache hit, do nothing
-        },
-        |key, lookup| {
-            let job = match JobOwner::<'_, CTX::DepKind, CTX::Query, C>::try_start(
-                tcx, state, span, &key, lookup, query,
-            ) {
-                TryGetJob::NotYetStarted(job) => job,
-                TryGetJob::Cycle(_) => return,
-                #[cfg(parallel_compiler)]
-                TryGetJob::JobCompleted(_) => return,
-            };
-            force_query_with_job(tcx, key, job, dep_node, query);
-        },
-    );
+    let cached = try_get_cached(tcx, state, &key, |_, _| {
+        // Cache hit, do nothing
+    });
+
+    let lookup = match cached {
+        Ok(()) => return,
+        Err(lookup) => lookup,
+    };
+
+    let job = match JobOwner::<'_, CTX::DepKind, CTX::Query, C>::try_start(
+        tcx, state, span, &key, lookup, query,
+    ) {
+        TryGetJob::NotYetStarted(job) => job,
+        TryGetJob::Cycle(_) => return,
+        #[cfg(parallel_compiler)]
+        TryGetJob::JobCompleted(_) => return,
+    };
+    force_query_with_job(tcx, key, job, dep_node, query);
 }
 
 pub enum QueryMode {