Rollup merge of #153492 - Zoxc:querykeyid, r=nnethercote

Add a `TaggedQueryKey` to identify a query instance

This adds back a `TaggedQueryKey` enum which represents a query kind and the associated key. This is used to replace `QueryStackDeferred` and `QueryStackFrameExtra` and the associated lifting operation for cycle errors

This approach has a couple of benefits:
- We only run description queries when printing the query stack trace in the panic handler
- The unsafe code for `QueryStackDeferred` is removed
- Cycle handles have access to query keys, which may be handy

Some further work could be replacing `QueryStackFrame` with `TaggedQueryKey` as the extra information can be derived from it.

Author: Zalathar

Cargo.lock

@@ -4575,7 +4575,6 @@ dependencies = [
  "rustc_macros",
  "rustc_middle",
  "rustc_serialize",
- "rustc_session",
  "rustc_span",
  "rustc_thread_pool",
  "tracing",

compiler/rustc_middle/src/query/job.rs

@@ -7,21 +7,15 @@ use parking_lot::{Condvar, Mutex};
 use rustc_span::Span;
 
 use crate::query::plumbing::CycleError;
-use crate::query::stack::{QueryStackDeferred, QueryStackFrame, QueryStackFrameExtra};
+use crate::query::stack::QueryStackFrame;
 use crate::ty::TyCtxt;
 
 /// Represents a span and a query key.
 #[derive(Clone, Debug)]
-pub struct QueryInfo<I> {
+pub struct QueryInfo<'tcx> {
     /// The span corresponding to the reason for which this query was required.
     pub span: Span,
-    pub frame: QueryStackFrame<I>,
-}
-
-impl<'tcx> QueryInfo<QueryStackDeferred<'tcx>> {
-    pub(crate) fn lift(&self) -> QueryInfo<QueryStackFrameExtra> {
-        QueryInfo { span: self.span, frame: self.frame.lift() }
-    }
+    pub frame: QueryStackFrame<'tcx>,
 }
 
 /// A value uniquely identifying an active query job.
@@ -74,7 +68,7 @@ pub struct QueryWaiter<'tcx> {
     pub query: Option<QueryJobId>,
     pub condvar: Condvar,
     pub span: Span,
-    pub cycle: Mutex<Option<CycleError<QueryStackDeferred<'tcx>>>>,
+    pub cycle: Mutex<Option<CycleError<'tcx>>>,
 }
 
 #[derive(Clone, Debug)]
@@ -94,7 +88,7 @@ impl<'tcx> QueryLatch<'tcx> {
         tcx: TyCtxt<'tcx>,
         query: Option<QueryJobId>,
         span: Span,
-    ) -> Result<(), CycleError<QueryStackDeferred<'tcx>>> {
+    ) -> Result<(), CycleError<'tcx>> {
         let mut waiters_guard = self.waiters.lock();
         let Some(waiters) = &mut *waiters_guard else {
             return Ok(()); // already complete

compiler/rustc_middle/src/query/mod.rs

@@ -7,7 +7,7 @@ pub use self::plumbing::{
     ActiveKeyStatus, CycleError, CycleErrorHandling, EnsureMode, IntoQueryParam, QueryMode,
     QueryState, TyCtxtAt, TyCtxtEnsureDone, TyCtxtEnsureOk, TyCtxtEnsureResult,
 };
-pub use self::stack::{QueryStackDeferred, QueryStackFrame, QueryStackFrameExtra};
+pub use self::stack::QueryStackFrame;
 pub use crate::queries::Providers;
 use crate::ty::TyCtxt;
 

compiler/rustc_middle/src/query/plumbing.rs

@@ -12,9 +12,9 @@ pub use sealed::IntoQueryParam;
 
 use crate::dep_graph::{DepKind, DepNodeIndex, SerializedDepNodeIndex};
 use crate::ich::StableHashingContext;
-use crate::queries::{ExternProviders, Providers, QueryArenas, QueryVTables};
+use crate::queries::{ExternProviders, Providers, QueryArenas, QueryVTables, TaggedQueryKey};
 use crate::query::on_disk_cache::OnDiskCache;
-use crate::query::stack::{QueryStackDeferred, QueryStackFrame, QueryStackFrameExtra};
+use crate::query::stack::QueryStackFrame;
 use crate::query::{QueryCache, QueryInfo, QueryJob};
 use crate::ty::TyCtxt;
 
@@ -60,19 +60,10 @@ pub enum CycleErrorHandling {
 }
 
 #[derive(Clone, Debug)]
-pub struct CycleError<I = QueryStackFrameExtra> {
+pub struct CycleError<'tcx> {
     /// The query and related span that uses the cycle.
-    pub usage: Option<(Span, QueryStackFrame<I>)>,
-    pub cycle: Vec<QueryInfo<I>>,
-}
-
-impl<'tcx> CycleError<QueryStackDeferred<'tcx>> {
-    pub fn lift(&self) -> CycleError<QueryStackFrameExtra> {
-        CycleError {
-            usage: self.usage.as_ref().map(|(span, frame)| (*span, frame.lift())),
-            cycle: self.cycle.iter().map(|info| info.lift()).collect(),
-        }
-    }
+    pub usage: Option<(Span, QueryStackFrame<'tcx>)>,
+    pub cycle: Vec<QueryInfo<'tcx>>,
 }
 
 #[derive(Debug)]
@@ -139,16 +130,12 @@ pub struct QueryVTable<'tcx, C: QueryCache> {
     pub value_from_cycle_error: fn(
         tcx: TyCtxt<'tcx>,
         key: C::Key,
-        cycle_error: CycleError,
+        cycle_error: CycleError<'tcx>,
         guar: ErrorGuaranteed,
     ) -> C::Value,
     pub format_value: fn(&C::Value) -> String,
 
-    /// Formats a human-readable description of this query and its key, as
-    /// specified by the `desc` query modifier.
-    ///
-    /// Used when reporting query cycle errors and similar problems.
-    pub description_fn: fn(TyCtxt<'tcx>, C::Key) -> String,
+    pub create_tagged_key: fn(C::Key) -> TaggedQueryKey<'tcx>,
 
     /// Function pointer that is called by the query methods on [`TyCtxt`] and
     /// friends[^1], after they have checked the in-memory cache and found no
@@ -523,6 +510,69 @@ macro_rules! define_callbacks {
             }
         )*
 
+        /// Identifies a query by kind and key. This is in contrast to `QueryJobId` which is just a number.
+        #[allow(non_camel_case_types)]
+        #[derive(Clone, Debug)]
+        pub enum TaggedQueryKey<'tcx> {
+            $(
+                $name($name::Key<'tcx>),
+            )*
+        }
+
+        impl<'tcx> TaggedQueryKey<'tcx> {
+            /// Formats a human-readable description of this query and its key, as
+            /// specified by the `desc` query modifier.
+            ///
+            /// Used when reporting query cycle errors and similar problems.
+            pub fn description(&self, tcx: TyCtxt<'tcx>) -> String {
+                let (name, description) = ty::print::with_no_queries!(match self {
+                    $(
+                        TaggedQueryKey::$name(key) => (stringify!($name), _description_fns::$name(tcx, *key)),
+                    )*
+                });
+                if tcx.sess.verbose_internals() {
+                    format!("{description} [{name:?}]")
+                } else {
+                    description
+                }
+            }
+
+            /// Returns the default span for this query if `span` is a dummy span.
+            pub fn default_span(&self, tcx: TyCtxt<'tcx>, span: Span) -> Span {
+                if !span.is_dummy() {
+                    return span
+                }
+                if let TaggedQueryKey::def_span(..) = self {
+                    // The `def_span` query is used to calculate `default_span`,
+                    // so exit to avoid infinite recursion.
+                    return DUMMY_SP
+                }
+                match self {
+                    $(
+                        TaggedQueryKey::$name(key) => crate::query::QueryKey::default_span(key, tcx),
+                    )*
+                }
+            }
+
+            pub fn def_kind(&self, tcx: TyCtxt<'tcx>) -> Option<DefKind> {
+                // This is used to reduce code generation as it
+                // can be reused for queries with the same key type.
+                fn inner<'tcx>(key: &impl crate::query::QueryKey, tcx: TyCtxt<'tcx>) -> Option<DefKind> {
+                    key.key_as_def_id().and_then(|def_id| def_id.as_local()).map(|def_id| tcx.def_kind(def_id))
+                }
+
+                if let TaggedQueryKey::def_kind(..) = self {
+                    // Try to avoid infinite recursion.
+                    return None
+                }
+                match self {
+                    $(
+                        TaggedQueryKey::$name(key) => inner(key, tcx),
+                    )*
+                }
+            }
+        }
+
         /// Holds a `QueryVTable` for each query.
         pub struct QueryVTables<'tcx> {
             $(

compiler/rustc_middle/src/query/stack.rs

@@ -1,106 +1,18 @@
 use std::fmt::Debug;
-use std::marker::PhantomData;
-use std::mem::transmute;
-use std::sync::Arc;
 
-use rustc_data_structures::sync::{DynSend, DynSync};
-use rustc_hir::def::DefKind;
-use rustc_span::Span;
 use rustc_span::def_id::DefId;
 
 use crate::dep_graph::DepKind;
+use crate::queries::TaggedQueryKey;
 
 /// Description of a frame in the query stack.
 ///
 /// This is mostly used in case of cycles for error reporting.
 #[derive(Clone, Debug)]
-pub struct QueryStackFrame<I> {
-    /// This field initially stores a `QueryStackDeferred` during collection,
-    /// but can later be changed to `QueryStackFrameExtra` containing concrete information
-    /// by calling `lift`. This is done so that collecting query does not need to invoke
-    /// queries, instead `lift` will call queries in a more appropriate location.
-    pub info: I,
-
+pub struct QueryStackFrame<'tcx> {
+    pub tagged_key: TaggedQueryKey<'tcx>,
     pub dep_kind: DepKind,
     pub def_id: Option<DefId>,
     /// A def-id that is extracted from a `Ty` in a query key
     pub def_id_for_ty_in_cycle: Option<DefId>,
 }
-
-impl<'tcx> QueryStackFrame<QueryStackDeferred<'tcx>> {
-    #[inline]
-    pub fn new(
-        info: QueryStackDeferred<'tcx>,
-        dep_kind: DepKind,
-        def_id: Option<DefId>,
-        def_id_for_ty_in_cycle: Option<DefId>,
-    ) -> Self {
-        Self { info, def_id, dep_kind, def_id_for_ty_in_cycle }
-    }
-
-    pub fn lift(&self) -> QueryStackFrame<QueryStackFrameExtra> {
-        QueryStackFrame {
-            info: self.info.extract(),
-            dep_kind: self.dep_kind,
-            def_id: self.def_id,
-            def_id_for_ty_in_cycle: self.def_id_for_ty_in_cycle,
-        }
-    }
-}
-
-#[derive(Clone, Debug)]
-pub struct QueryStackFrameExtra {
-    pub description: String,
-    pub span: Option<Span>,
-    pub def_kind: Option<DefKind>,
-}
-
-impl QueryStackFrameExtra {
-    #[inline]
-    pub fn new(description: String, span: Option<Span>, def_kind: Option<DefKind>) -> Self {
-        Self { description, span, def_kind }
-    }
-
-    // FIXME(eddyb) Get more valid `Span`s on queries.
-    #[inline]
-    pub fn default_span(&self, span: Span) -> Span {
-        if !span.is_dummy() {
-            return span;
-        }
-        self.span.unwrap_or(span)
-    }
-}
-
-/// Track a 'side effect' for a particular query.
-/// This is used to hold a closure which can create `QueryStackFrameExtra`.
-#[derive(Clone)]
-pub struct QueryStackDeferred<'tcx> {
-    _dummy: PhantomData<&'tcx ()>,
-
-    // `extract` may contain references to 'tcx, but we can't tell drop checking that it won't
-    // access it in the destructor.
-    extract: Arc<dyn Fn() -> QueryStackFrameExtra + DynSync + DynSend>,
-}
-
-impl<'tcx> QueryStackDeferred<'tcx> {
-    pub fn new<C: Copy + DynSync + DynSend + 'tcx>(
-        context: C,
-        extract: fn(C) -> QueryStackFrameExtra,
-    ) -> Self {
-        let extract: Arc<dyn Fn() -> QueryStackFrameExtra + DynSync + DynSend + 'tcx> =
-            Arc::new(move || extract(context));
-        // SAFETY: The `extract` closure does not access 'tcx in its destructor as the only
-        // captured variable is `context` which is Copy and cannot have a destructor.
-        Self { _dummy: PhantomData, extract: unsafe { transmute(extract) } }
-    }
-
-    pub fn extract(&self) -> QueryStackFrameExtra {
-        (self.extract)()
-    }
-}
-
-impl<'tcx> Debug for QueryStackDeferred<'tcx> {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.write_str("QueryStackDeferred")
-    }
-}

compiler/rustc_query_impl/Cargo.toml

@@ -14,7 +14,6 @@ rustc_index = { path = "../rustc_index" }
 rustc_macros = { path = "../rustc_macros" }
 rustc_middle = { path = "../rustc_middle" }
 rustc_serialize = { path = "../rustc_serialize" }
-rustc_session = { path = "../rustc_session" }
 rustc_span = { path = "../rustc_span" }
 rustc_thread_pool = { path = "../rustc_thread_pool" }
 tracing = "0.1"

compiler/rustc_query_impl/src/execution.rs

@@ -48,7 +48,7 @@ pub fn collect_active_jobs_from_all_queries<'tcx>(
     let mut complete = true;
 
     for_each_query_vtable!(ALL, tcx, |query| {
-        let res = gather_active_jobs(query, tcx, require_complete, &mut job_map_out);
+        let res = gather_active_jobs(query, require_complete, &mut job_map_out);
         if res.is_none() {
             complete = false;
         }
@@ -66,7 +66,6 @@ pub fn collect_active_jobs_from_all_queries<'tcx>(
 /// grep for.)
 fn gather_active_jobs<'tcx, C>(
     query: &'tcx QueryVTable<'tcx, C>,
-    tcx: TyCtxt<'tcx>,
     require_complete: bool,
     job_map_out: &mut QueryJobMap<'tcx>, // Out-param; job info is gathered into this map
 ) -> Option<()>
@@ -113,7 +112,7 @@ where
     // Call `make_frame` while we're not holding a `state.active` lock as `make_frame` may call
     // queries leading to a deadlock.
     for (key, job) in active {
-        let frame = crate::plumbing::create_deferred_query_stack_frame(tcx, query, key);
+        let frame = crate::plumbing::create_query_stack_frame(query, key);
         job_map_out.insert(job.id, QueryJobInfo { frame, job });
     }
 
@@ -126,9 +125,9 @@ fn mk_cycle<'tcx, C: QueryCache>(
     query: &'tcx QueryVTable<'tcx, C>,
     tcx: TyCtxt<'tcx>,
     key: C::Key,
-    cycle_error: CycleError,
+    cycle_error: CycleError<'tcx>,
 ) -> C::Value {
-    let error = report_cycle(tcx.sess, &cycle_error);
+    let error = report_cycle(tcx, &cycle_error);
     match query.cycle_error_handling {
         CycleErrorHandling::Error => {
             let guar = error.emit();
@@ -231,7 +230,7 @@ fn cycle_error<'tcx, C: QueryCache>(
         .expect("failed to collect active queries");
 
     let error = find_cycle_in_stack(try_execute, job_map, &current_query_job(), span);
-    (mk_cycle(query, tcx, key, error.lift()), None)
+    (mk_cycle(query, tcx, key, error), None)
 }
 
 #[inline(always)]
@@ -275,7 +274,7 @@ fn wait_for_query<'tcx, C: QueryCache>(
 
             (v, Some(index))
         }
-        Err(cycle) => (mk_cycle(query, tcx, key, cycle.lift()), None),
+        Err(cycle) => (mk_cycle(query, tcx, key, cycle), None),
     }
 }