Generalise QueryLatch.

diff --git a/src/librustc/ty/query/config.rs b/src/librustc/ty/query/config.rs
index c3b0103cc0c21788772b5a9f54a3d4d598171f08..a4fb70f281c801c2b489ea75669c70e253be6070 100644 (file)
--- a/src/librustc/ty/query/config.rs
+++ b/src/librustc/ty/query/config.rs
@@ -2,6 +2,7 @@
 
 use crate::dep_graph::SerializedDepNodeIndex;
 use crate::ty::query::caches::QueryCache;
+use crate::ty::query::job::QueryJobId;
 use crate::ty::query::plumbing::CycleError;
 use crate::ty::query::QueryState;
 use rustc_data_structures::profiling::ProfileCategory;
@@ -30,6 +31,9 @@ pub trait QueryContext: DepContext {
 
     /// Get string representation from DefPath.
     fn def_path_str(&self, def_id: DefId) -> String;
+
+    /// Get the query information from the TLS context.
+    fn read_query_job<R>(&self, op: impl FnOnce(Option<QueryJobId<Self::DepKind>>) -> R) -> R;
 }
 
 pub(crate) trait QueryAccessors<CTX: QueryContext>: QueryConfig<CTX> {

diff --git a/src/librustc/ty/query/job.rs b/src/librustc/ty/query/job.rs
index 488615c7443f29f448dfd29e5062d38c4a085b3c..b4d986074a487f194f98dd024389e57b0461f1ea 100644 (file)
--- a/src/librustc/ty/query/job.rs
+++ b/src/librustc/ty/query/job.rs
@@ -53,11 +53,13 @@ pub struct QueryJobId<K> {
     pub kind: K,
 }
 
-impl QueryJobId<DepKind> {
-    pub fn new(job: QueryShardJobId, shard: usize, kind: DepKind) -> Self {
+impl<K> QueryJobId<K> {
+    pub fn new(job: QueryShardJobId, shard: usize, kind: K) -> Self {
         QueryJobId { job, shard: u16::try_from(shard).unwrap(), kind }
     }
+}
 
+impl QueryJobId<DepKind> {
     fn query<'tcx>(self, map: &QueryMap<'tcx>) -> Query<'tcx> {
         map.get(&self).unwrap().info.query.clone()
     }
@@ -223,16 +225,16 @@ fn new() -> Self {
 }
 
 #[cfg(parallel_compiler)]
-impl<'tcx> QueryLatch<TyCtxt<'tcx>> {
+impl<K, CTX> QueryLatch<CTX>
+where
+    K: rustc_query_system::dep_graph::DepKind,
+    CTX: QueryContext<DepKind = K>,
+{
     /// Awaits for the query job to complete.
-    pub(super) fn wait_on(
-        &self,
-        tcx: TyCtxt<'tcx>,
-        span: Span,
-    ) -> Result<(), CycleError<TyCtxt<'tcx>>> {
-        tls::with_related_context(tcx, move |icx| {
+    pub(super) fn wait_on(&self, tcx: CTX, span: Span) -> Result<(), CycleError<CTX>> {
+        tcx.read_query_job(move |query| {
             let waiter = Lrc::new(QueryWaiter {
-                query: icx.query,
+                query,
                 span,
                 cycle: Lock::new(None),
                 condvar: Condvar::new(),

diff --git a/src/librustc/ty/query/plumbing.rs b/src/librustc/ty/query/plumbing.rs
index ede1ccdbc04b9bdde985ea7c440d7b0db7a984fc..6d01f34c630897e565daf53413a7dee71a66235a 100644 (file)
--- a/src/librustc/ty/query/plumbing.rs
+++ b/src/librustc/ty/query/plumbing.rs
@@ -367,6 +367,10 @@ fn session(&self) -> &Session {
     fn def_path_str(&self, def_id: DefId) -> String {
         TyCtxt::def_path_str(*self, def_id)
     }
+
+    fn read_query_job<R>(&self, op: impl FnOnce(Option<QueryJobId<Self::DepKind>>) -> R) -> R {
+        tls::with_related_context(*self, move |icx| op(icx.query))
+    }
 }
 
 impl<'tcx> TyCtxt<'tcx> {