Don't use queries when printing the query stack

Zoxc · Zoxc · commit c7aa46822b74 · 2023-08-21T20:24:04.000+02:00
diff --git a/compiler/rustc_interface/src/util.rs b/compiler/rustc_interface/src/util.rs
@@ -193,7 +193,7 @@ pub(crate) fn run_in_thread_pool_with_globals<F: FnOnce() -> R + Send, R: Send>(
             // locals to it. The new thread runs the deadlock handler.
             let query_map = tls::with(|tcx| {
                 QueryCtxt::new(tcx)
-                    .try_collect_active_jobs()
+                    .try_collect_active_jobs(true)
                     .expect("active jobs shouldn't be locked in deadlock handler")
             });
             let registry = rayon_core::Registry::current();
diff --git a/compiler/rustc_macros/src/query.rs b/compiler/rustc_macros/src/query.rs
@@ -273,11 +273,15 @@ fn add_query_desc_cached_impl(
 
     let desc = quote! {
         #[allow(unused_variables)]
-        pub fn #name<'tcx>(tcx: TyCtxt<'tcx>, key: crate::query::queries::#name::Key<'tcx>) -> String {
-            let (#tcx, #key) = (tcx, key);
-            ::rustc_middle::ty::print::with_no_trimmed_paths!(
+        pub fn #name<'tcx>(
+            tcx: TyCtxt<'tcx>,
+            key: crate::query::queries::#name::Key<'tcx>,
+            use_queries: bool
+        ) -> String {
+            ::rustc_middle::query::print::run(tcx, use_queries, |ctx| {
+                let (#tcx, #key) = (ctx, ctx.arg(key));
                 format!(#desc)
-            )
+            })
         }
     };
 
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
diff --git a/compiler/rustc_middle/src/query/print.rs b/compiler/rustc_middle/src/query/print.rs
@@ -0,0 +1,181 @@
+use crate::mir::interpret::GlobalId;
+use crate::query::IntoQueryParam;
+use crate::query::TyCtxt;
+use crate::ty;
+use rustc_hir::def_id::CrateNum;
+use rustc_hir::def_id::DefId;
+use rustc_span::Symbol;
+use std::fmt;
+use std::fmt::Debug;
+use std::fmt::Display;
+
+/// This is used to print query keys without giving them access to `TyCtxt`.
+/// This enables more reliable printing when printing the query stack on panics.
+#[derive(Copy, Clone)]
+pub struct PrintContext<'tcx> {
+    tcx: TyCtxt<'tcx>,
+    use_queries: bool,
+}
+
+impl<'tcx> PrintContext<'tcx> {
+    pub fn arg<T>(&self, arg: T) -> PrintArg<'tcx, T> {
+        PrintArg { arg: Some(arg), ctx: *self }
+    }
+
+    pub fn def_path_str(&self, def_id: PrintArg<'tcx, impl IntoQueryParam<DefId>>) -> String {
+        def_id.print(|arg| {
+            let def_id = arg.into_query_param();
+            if self.use_queries {
+                self.tcx.def_path_str(def_id)
+            } else {
+                self.tcx.safe_def_path_str_untracked(def_id)
+            }
+        })
+    }
+}
+
+/// This runs some code in a printing context. If `use_queries` is false this function should
+/// ensure that no queries are run.
+pub fn run<'tcx, R>(
+    tcx: TyCtxt<'tcx>,
+    use_queries: bool,
+    print: impl FnOnce(PrintContext<'tcx>) -> R,
+) -> R {
+    let ctx = PrintContext { tcx, use_queries };
+    if use_queries { ty::print::with_no_trimmed_paths!(print(ctx)) } else { print(ctx) }
+}
+
+const INACCESSIBLE: &'static str = "<inaccessible>";
+
+/// This wraps a value that we want to print without giving access to the regular types
+/// and their Display and Debug impls. `map` and `map_with_tcx` gives access to the inner
+/// type, but erases the value for the no query path.
+#[derive(Copy, Clone)]
+pub struct PrintArg<'tcx, T> {
+    arg: Option<T>,
+    ctx: PrintContext<'tcx>,
+}
+
+impl<'tcx, T> PrintArg<'tcx, T> {
+    fn print(self, f: impl FnOnce(T) -> String) -> String {
+        self.arg.map(f).unwrap_or_else(|| INACCESSIBLE.to_owned())
+    }
+
+    fn fmt_map(
+        &self,
+        f: &mut fmt::Formatter<'_>,
+        map: impl FnOnce(&mut fmt::Formatter<'_>, &T) -> fmt::Result,
+    ) -> fmt::Result {
+        match &self.arg {
+            Some(arg) => map(f, arg),
+            _ => write!(f, "{}", INACCESSIBLE),
+        }
+    }
+
+    fn fmt_with_queries(
+        &self,
+        f: &mut fmt::Formatter<'_>,
+        map: impl FnOnce(&mut fmt::Formatter<'_>, &T) -> fmt::Result,
+    ) -> fmt::Result {
+        match &self.arg {
+            Some(arg) if self.ctx.use_queries => map(f, arg),
+            _ => write!(f, "{}", INACCESSIBLE),
+        }
+    }
+
+    pub fn ctx(&self) -> PrintContext<'tcx> {
+        self.ctx
+    }
+
+    /// Maps the argument where `f` is known to not call queries.
+    fn map_unchecked<R>(self, f: impl FnOnce(T) -> R) -> PrintArg<'tcx, R> {
+        PrintArg { arg: self.arg.map(f), ctx: self.ctx }
+    }
+
+    pub fn map<R>(self, f: impl FnOnce(T) -> R) -> PrintArg<'tcx, R> {
+        self.map_with_tcx(|_, arg| f(arg))
+    }
+
+    pub fn map_with_tcx<R>(self, f: impl FnOnce(TyCtxt<'tcx>, T) -> R) -> PrintArg<'tcx, R> {
+        PrintArg {
+            arg: if self.ctx.use_queries { self.arg.map(|arg| f(self.ctx.tcx, arg)) } else { None },
+            ctx: self.ctx,
+        }
+    }
+
+    pub fn describe_as_module(&self) -> String
+    where
+        T: IntoQueryParam<DefId> + Copy,
+    {
+        self.print(|arg| {
+            let def_id: DefId = arg.into_query_param();
+            if def_id.is_top_level_module() {
+                "top-level module".to_string()
+            } else {
+                format!("module `{}`", self.ctx.def_path_str(*self))
+            }
+        })
+    }
+}
+
+impl<'tcx, T0, T1> PrintArg<'tcx, (T0, T1)> {
+    pub fn i_0(self) -> PrintArg<'tcx, T0> {
+        self.map_unchecked(|arg| arg.0)
+    }
+    pub fn i_1(self) -> PrintArg<'tcx, T1> {
+        self.map_unchecked(|arg| arg.1)
+    }
+}
+
+impl<'tcx, T> PrintArg<'tcx, ty::ParamEnvAnd<'tcx, T>> {
+    pub fn value(self) -> PrintArg<'tcx, T> {
+        self.map_unchecked(|arg| arg.value)
+    }
+}
+
+impl<'tcx> PrintArg<'tcx, GlobalId<'tcx>> {
+    pub fn display(self) -> String {
+        self.print(|arg| {
+            if self.ctx.use_queries {
+                arg.display(self.ctx.tcx)
+            } else {
+                let instance_name = self.ctx.def_path_str(self.ctx.arg(arg.instance.def.def_id()));
+                if let Some(promoted) = arg.promoted {
+                    format!("{instance_name}::{promoted:?}")
+                } else {
+                    instance_name
+                }
+            }
+        })
+    }
+}
+
+impl<T: Debug> Debug for PrintArg<'_, T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.fmt_with_queries(f, |f, arg| arg.fmt(f))
+    }
+}
+
+impl<T: Display> Display for PrintArg<'_, T> {
+    default fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.fmt_with_queries(f, |f, arg| arg.fmt(f))
+    }
+}
+
+impl Display for PrintArg<'_, CrateNum> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.fmt_map(f, |f, arg| {
+            if self.ctx.use_queries {
+                write!(f, "`{}`", arg)
+            } else {
+                write!(f, "`{}`", self.ctx.tcx.safe_crate_str_untracked(*arg))
+            }
+        })
+    }
+}
+
+impl Display for PrintArg<'_, Symbol> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.fmt_map(f, |f, arg| fmt::Display::fmt(arg, f))
+    }
+}
diff --git a/compiler/rustc_middle/src/ty/context.rs b/compiler/rustc_middle/src/ty/context.rs
@@ -888,6 +888,31 @@ impl<'tcx> TyCtxt<'tcx> {
         }
     }
 
+    /// Prints a CrateNum without using queries.
+    pub fn safe_crate_str_untracked(self, krate: CrateNum) -> String {
+        if krate == LOCAL_CRATE {
+            "crate".to_owned()
+        } else {
+            let cstore = &*self.cstore_untracked();
+            cstore.crate_name(krate).as_str().to_owned()
+        }
+    }
+
+    /// Prints a DefId without using queries.
+    pub fn safe_def_path_str_untracked(self, def_id: DefId) -> String {
+        if def_id.is_local() {
+            let format = self.def_path(def_id).to_string_no_crate_verbose();
+            // Strip `::` from the formatted def path if present
+            format.strip_prefix("::").map(|stripped| stripped.to_owned()).unwrap_or(format)
+        } else {
+            format!(
+                "{}{}",
+                self.safe_crate_str_untracked(def_id.krate),
+                self.def_path(def_id).to_string_no_crate_verbose()
+            )
+        }
+    }
+
     pub fn def_path_debug_str(self, def_id: DefId) -> String {
         // We are explicitly not going through queries here in order to get
         // crate name and stable crate id since this code is called from debug!()
diff --git a/compiler/rustc_middle/src/ty/print/mod.rs b/compiler/rustc_middle/src/ty/print/mod.rs
@@ -3,7 +3,7 @@ use crate::ty::{self, Ty, TyCtxt};
 
 use rustc_data_structures::fx::FxHashSet;
 use rustc_data_structures::sso::SsoHashSet;
-use rustc_hir::def_id::{CrateNum, DefId, LocalDefId};
+use rustc_hir::def_id::{CrateNum, DefId};
 use rustc_hir::definitions::{DefPathData, DisambiguatedDefPathData};
 
 // `pretty` is a separate module only for organization.
@@ -327,13 +327,3 @@ impl<'tcx, P: Printer<'tcx>> Print<'tcx, P> for ty::Const<'tcx> {
         cx.print_const(*self)
     }
 }
-
-// This is only used by query descriptions
-pub fn describe_as_module(def_id: impl Into<LocalDefId>, tcx: TyCtxt<'_>) -> String {
-    let def_id = def_id.into();
-    if def_id.is_top_level_module() {
-        "top-level module".to_string()
-    } else {
-        format!("module `{}`", tcx.def_path_str(def_id))
-    }
-}
diff --git a/compiler/rustc_query_impl/src/plumbing.rs b/compiler/rustc_query_impl/src/plumbing.rs
@@ -78,11 +78,11 @@ impl QueryContext for QueryCtxt<'_> {
         tls::with_related_context(self.tcx, |icx| icx.query)
     }
 
-    fn try_collect_active_jobs(self) -> Option<QueryMap<DepKind>> {
+    fn try_collect_active_jobs(self, can_call_queries: bool) -> Option<QueryMap<DepKind>> {
         let mut jobs = QueryMap::default();
 
         for collect in super::TRY_COLLECT_ACTIVE_JOBS.iter() {
-            collect(self.tcx, &mut jobs);
+            collect(self.tcx, &mut jobs, can_call_queries);
         }
 
         Some(jobs)
@@ -153,7 +153,7 @@ impl QueryContext for QueryCtxt<'_> {
     fn depth_limit_error(self, job: QueryJobId) {
         let mut span = None;
         let mut layout_of_depth = None;
-        if let Some(map) = self.try_collect_active_jobs() {
+        if let Some(map) = self.try_collect_active_jobs(true) {
             if let Some((info, depth)) = job.try_find_layout_root(map) {
                 span = Some(info.job.span);
                 layout_of_depth = Some(LayoutOfDepth { desc: info.query.description, depth });
@@ -296,18 +296,25 @@ pub(crate) fn create_query_frame<
     K: Copy + Key + for<'a> HashStable<StableHashingContext<'a>>,
 >(
     tcx: TyCtxt<'tcx>,
-    do_describe: fn(TyCtxt<'tcx>, K) -> String,
+    do_describe: fn(TyCtxt<'tcx>, K, bool) -> String,
     key: K,
     kind: DepKind,
     name: &'static str,
+    can_call_queries: bool,
 ) -> QueryStackFrame<DepKind> {
+    if !can_call_queries {
+        // Return a minimal query frame if we can't call queries
+        let description = do_describe(tcx, key, false);
+        return QueryStackFrame::new(description, None, None, None, kind, None, || Hash64::ZERO);
+    }
+
     // Avoid calling queries while formatting the description
     let description = ty::print::with_no_queries!(
         // Disable visible paths printing for performance reasons.
         // Showing visible path instead of any path is not that important in production.
         ty::print::with_no_visible_paths!(
             // Force filename-line mode to avoid invoking `type_of` query.
-            ty::print::with_forced_impl_filename_line!(do_describe(tcx, key))
+            ty::print::with_forced_impl_filename_line!(do_describe(tcx, key, true))
         )
     );
     let description =
@@ -650,16 +657,28 @@ macro_rules! define_queries {
                 }
             }
 
-            pub fn try_collect_active_jobs<'tcx>(tcx: TyCtxt<'tcx>, qmap: &mut QueryMap<DepKind>) {
-                let make_query = |tcx, key| {
+            pub fn try_collect_active_jobs<'tcx>(
+                tcx: TyCtxt<'tcx>,
+                qmap: &mut QueryMap<DepKind>,
+                can_call_queries: bool,
+            ) {
+                let make_query = |tcx, key, can_call_queries| {
                     let kind = rustc_middle::dep_graph::DepKind::$name;
                     let name = stringify!($name);
-                    $crate::plumbing::create_query_frame(tcx, rustc_middle::query::descs::$name, key, kind, name)
+                    $crate::plumbing::create_query_frame(
+                        tcx,
+                        rustc_middle::query::descs::$name,
+                        key,
+                        kind,
+                        name,
+                        can_call_queries,
+                    )
                 };
                 tcx.query_system.states.$name.try_collect_active_jobs(
                     tcx,
                     make_query,
                     qmap,
+                    can_call_queries,
                 ).unwrap();
             }
 
@@ -710,7 +729,7 @@ macro_rules! define_queries {
 
         // These arrays are used for iteration and can't be indexed by `DepKind`.
 
-        const TRY_COLLECT_ACTIVE_JOBS: &[for<'tcx> fn(TyCtxt<'tcx>, &mut QueryMap<DepKind>)] =
+        const TRY_COLLECT_ACTIVE_JOBS: &[for<'tcx> fn(TyCtxt<'tcx>, &mut QueryMap<DepKind>, can_call_queries: bool)] =
             &[$(query_impl::$name::try_collect_active_jobs),*];
 
         const ALLOC_SELF_PROFILE_QUERY_STRINGS: &[
diff --git a/compiler/rustc_query_system/src/query/job.rs b/compiler/rustc_query_system/src/query/job.rs
@@ -628,7 +628,7 @@ pub fn print_query_stack<Qcx: QueryContext>(
     // state if it was responsible for triggering the panic.
     let mut count_printed = 0;
     let mut count_total = 0;
-    let query_map = qcx.try_collect_active_jobs();
+    let query_map = qcx.try_collect_active_jobs(false);
 
     if let Some(ref mut file) = file {
         let _ = writeln!(file, "\n\nquery stack during panic:");
diff --git a/compiler/rustc_query_system/src/query/mod.rs b/compiler/rustc_query_system/src/query/mod.rs
@@ -106,7 +106,7 @@ pub trait QueryContext: HasDepContext {
     /// Get the query information from the TLS context.
     fn current_query_job(self) -> Option<QueryJobId>;
 
-    fn try_collect_active_jobs(self) -> Option<QueryMap<Self::DepKind>>;
+    fn try_collect_active_jobs(self, can_call_queries: bool) -> Option<QueryMap<Self::DepKind>>;
 
     /// Load side effects associated to the node in the previous session.
     fn load_side_effects(self, prev_dep_node_index: SerializedDepNodeIndex) -> QuerySideEffects;
diff --git a/compiler/rustc_query_system/src/query/plumbing.rs b/compiler/rustc_query_system/src/query/plumbing.rs
