1 files changed, 35 insertions, 26 deletions
diff --git a/compiler/rustc_query_system/src/query/plumbing.rs b/compiler/rustc_query_system/src/query/plumbing.rs
index 3a9d80280c2..6ea8e3b9200 100644
--- a/compiler/rustc_query_system/src/query/plumbing.rs
+++ b/compiler/rustc_query_system/src/query/plumbing.rs
@@ -16,7 +16,7 @@ use rustc_errors::{Diag, FatalError, StashKey};
 use rustc_span::{DUMMY_SP, Span};
 use tracing::instrument;
 
-use super::QueryConfig;
+use super::{QueryConfig, QueryStackFrameExtra};
 use crate::HandleCycleError;
 use crate::dep_graph::{DepContext, DepGraphData, DepNode, DepNodeIndex, DepNodeParams};
 use crate::ich::StableHashingContext;
@@ -29,23 +29,23 @@ fn equivalent_key<K: Eq, V>(k: &K) -> impl Fn(&(K, V)) -> bool + '_ {
     move |x| x.0 == *k
 }
 
-pub struct QueryState<K> {
-    active: Sharded<hashbrown::HashTable<(K, QueryResult)>>,
+pub struct QueryState<K, I> {
+    active: Sharded<hashbrown::HashTable<(K, QueryResult<I>)>>,
 }
 
 /// Indicates the state of a query for a given key in a query map.
-enum QueryResult {
+enum QueryResult<I> {
     /// An already executing query. The query job can be used to await for its completion.
-    Started(QueryJob),
+    Started(QueryJob<I>),
 
     /// The query panicked. Queries trying to wait on this will raise a fatal error which will
     /// silently panic.
     Poisoned,
 }
 
-impl QueryResult {
+impl<I> QueryResult<I> {
     /// Unwraps the query job expecting that it has started.
-    fn expect_job(self) -> QueryJob {
+    fn expect_job(self) -> QueryJob<I> {
         match self {
             Self::Started(job) => job,
             Self::Poisoned => {
@@ -55,7 +55,7 @@ impl QueryResult {
     }
 }
 
-impl<K> QueryState<K>
+impl<K, I> QueryState<K, I>
 where
     K: Eq + Hash + Copy + Debug,
 {
@@ -66,8 +66,8 @@ where
     pub fn try_collect_active_jobs<Qcx: Copy>(
         &self,
         qcx: Qcx,
-        make_query: fn(Qcx, K) -> QueryStackFrame,
-        jobs: &mut QueryMap,
+        make_query: fn(Qcx, K) -> QueryStackFrame<I>,
+        jobs: &mut QueryMap<I>,
     ) -> Option<()> {
         let mut active = Vec::new();
 
@@ -76,7 +76,7 @@ where
         for shard in self.active.try_lock_shards() {
             for (k, v) in shard?.iter() {
                 if let QueryResult::Started(ref job) = *v {
-                    active.push((*k, job.clone()));
+                    active.push((*k, (*job).clone()));
                 }
             }
         }
@@ -92,19 +92,19 @@ where
     }
 }
 
-impl<K> Default for QueryState<K> {
-    fn default() -> QueryState<K> {
+impl<K, I> Default for QueryState<K, I> {
+    fn default() -> QueryState<K, I> {
         QueryState { active: Default::default() }
     }
 }
 
 /// A type representing the responsibility to execute the job in the `job` field.
 /// This will poison the relevant query if dropped.
-struct JobOwner<'tcx, K>
+struct JobOwner<'tcx, K, I>
 where
     K: Eq + Hash + Copy,
 {
-    state: &'tcx QueryState<K>,
+    state: &'tcx QueryState<K, I>,
     key: K,
 }
 
@@ -146,7 +146,7 @@ where
         }
         Stash => {
             let guar = if let Some(root) = cycle_error.cycle.first()
-                && let Some(span) = root.query.span
+                && let Some(span) = root.query.info.span
             {
                 error.stash(span, StashKey::Cycle).unwrap()
             } else {
@@ -157,7 +157,7 @@ where
     }
 }
 
-impl<'tcx, K> JobOwner<'tcx, K>
+impl<'tcx, K, I> JobOwner<'tcx, K, I>
 where
     K: Eq + Hash + Copy,
 {
@@ -194,7 +194,7 @@ where
     }
 }
 
-impl<'tcx, K> Drop for JobOwner<'tcx, K>
+impl<'tcx, K, I> Drop for JobOwner<'tcx, K, I>
 where
     K: Eq + Hash + Copy,
 {
@@ -222,10 +222,19 @@ where
 }
 
 #[derive(Clone, Debug)]
-pub struct CycleError {
+pub struct CycleError<I = QueryStackFrameExtra> {
     /// The query and related span that uses the cycle.
-    pub usage: Option<(Span, QueryStackFrame)>,
-    pub cycle: Vec<QueryInfo>,
+    pub usage: Option<(Span, QueryStackFrame<I>)>,
+    pub cycle: Vec<QueryInfo<I>>,
+}
+
+impl<I> CycleError<I> {
+    fn lift<Qcx: QueryContext<QueryInfo = I>>(&self, qcx: Qcx) -> CycleError<QueryStackFrameExtra> {
+        CycleError {
+            usage: self.usage.as_ref().map(|(span, frame)| (*span, frame.lift(qcx))),
+            cycle: self.cycle.iter().map(|info| info.lift(qcx)).collect(),
+        }
+    }
 }
 
 /// Checks whether there is already a value for this key in the in-memory
@@ -262,10 +271,10 @@ where
 {
     // Ensure there was no errors collecting all active jobs.
     // We need the complete map to ensure we find a cycle to break.
-    let query_map = qcx.collect_active_jobs().expect("failed to collect active queries");
+    let query_map = qcx.collect_active_jobs().ok().expect("failed to collect active queries");
 
     let error = try_execute.find_cycle_in_stack(query_map, &qcx.current_query_job(), span);
-    (mk_cycle(query, qcx, error), None)
+    (mk_cycle(query, qcx, error.lift(qcx)), None)
 }
 
 #[inline(always)]
@@ -274,7 +283,7 @@ fn wait_for_query<Q, Qcx>(
     qcx: Qcx,
     span: Span,
     key: Q::Key,
-    latch: QueryLatch,
+    latch: QueryLatch<Qcx::QueryInfo>,
     current: Option<QueryJobId>,
 ) -> (Q::Value, Option<DepNodeIndex>)
 where
@@ -314,7 +323,7 @@ where
 
             (v, Some(index))
         }
-        Err(cycle) => (mk_cycle(query, qcx, cycle), None),
+        Err(cycle) => (mk_cycle(query, qcx, cycle.lift(qcx)), None),
     }
 }
 
@@ -392,7 +401,7 @@ where
 fn execute_job<Q, Qcx, const INCR: bool>(
     query: Q,
     qcx: Qcx,
-    state: &QueryState<Q::Key>,
+    state: &QueryState<Q::Key, Qcx::QueryInfo>,
     key: Q::Key,
     key_hash: u64,
     id: QueryJobId,