Overhaul rustc_middle::query

compiler/rustc_middle/src/queries.rs

@@ -97,8 +97,7 @@ use crate::mir::interpret::{
 use crate::mono::{
     CodegenUnit, CollectionMode, MonoItem, MonoItemPartitions, NormalizationErrorInMono,
 };
-use crate::query::describe_as_module;
-use crate::query::plumbing::{define_callbacks, maybe_into_query_key};
+use crate::query::query_api::{define_query_api, maybe_into_query_key};
 use crate::traits::query::{
     CanonicalAliasGoal, CanonicalDropckOutlivesGoal, CanonicalImpliedOutlivesBoundsGoal,
     CanonicalMethodAutoderefStepsGoal, CanonicalPredicateGoal, CanonicalTypeOpAscribeUserTypeGoal,
@@ -120,6 +119,15 @@ use crate::ty::{
 };
 use crate::{mir, thir};
 
+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))
+    }
+}
+
 // Each of these queries corresponds to a function pointer field in the
 // `Providers` struct for requesting a value of that type, and a method
 // on `tcx: TyCtxt` (and `tcx.at(span)`) for doing that request in a way
@@ -2835,4 +2843,4 @@ rustc_queries! {
     non_query Metadata
 }
 
-rustc_with_all_queries! { define_callbacks! }
+rustc_with_all_queries! { define_query_api! }

compiler/rustc_middle/src/query/inner.rs → compiler/rustc_middle/src/query/calls.rs

@@ -1,13 +1,122 @@
 //! Helper functions that serve as the immediate implementation of
 //! `tcx.$query(..)` and its variations.
 
+use std::ops::Deref;
+
+use rustc_hir::def_id::LocalDefId;
 use rustc_span::{DUMMY_SP, ErrorGuaranteed, Span};
 
 use crate::dep_graph;
 use crate::dep_graph::DepNodeKey;
 use crate::query::erase::{self, Erasable, Erased};
-use crate::query::{QueryCache, QueryMode, QueryVTable};
-use crate::ty::TyCtxt;
+use crate::query::{IntoQueryKey, QueryCache, QueryMode, QueryVTable};
+use crate::ty::{self, TyCtxt};
+
+#[derive(Copy, Clone)]
+pub struct TyCtxtAt<'tcx> {
+    pub tcx: TyCtxt<'tcx>,
+    pub span: Span,
+}
+
+impl<'tcx> Deref for TyCtxtAt<'tcx> {
+    type Target = TyCtxt<'tcx>;
+    #[inline(always)]
+    fn deref(&self) -> &Self::Target {
+        &self.tcx
+    }
+}
+
+#[derive(Copy, Clone)]
+#[must_use]
+pub struct TyCtxtEnsureOk<'tcx> {
+    pub tcx: TyCtxt<'tcx>,
+}
+
+#[derive(Copy, Clone)]
+#[must_use]
+pub struct TyCtxtEnsureResult<'tcx> {
+    pub tcx: TyCtxt<'tcx>,
+}
+
+#[derive(Copy, Clone)]
+#[must_use]
+pub struct TyCtxtEnsureDone<'tcx> {
+    pub tcx: TyCtxt<'tcx>,
+}
+
+impl<'tcx> TyCtxtEnsureOk<'tcx> {
+    pub fn typeck(self, def_id: impl IntoQueryKey<LocalDefId>) {
+        self.typeck_root(
+            self.tcx.typeck_root_def_id(def_id.into_query_key().to_def_id()).expect_local(),
+        )
+    }
+}
+
+impl<'tcx> TyCtxt<'tcx> {
+    pub fn typeck(self, def_id: impl IntoQueryKey<LocalDefId>) -> &'tcx ty::TypeckResults<'tcx> {
+        self.typeck_root(
+            self.typeck_root_def_id(def_id.into_query_key().to_def_id()).expect_local(),
+        )
+    }
+
+    /// Returns a transparent wrapper for `TyCtxt` which uses
+    /// `span` as the location of queries performed through it.
+    #[inline(always)]
+    pub fn at(self, span: Span) -> TyCtxtAt<'tcx> {
+        TyCtxtAt { tcx: self, span }
+    }
+
+    /// FIXME: `ensure_ok`'s effects are subtle. Is this comment fully accurate?
+    ///
+    /// Wrapper that calls queries in a special "ensure OK" mode, for callers
+    /// that don't need the return value and just want to invoke a query for
+    /// its potential side-effect of emitting fatal errors.
+    ///
+    /// This can be more efficient than a normal query call, because if the
+    /// query's inputs are all green, the call can return immediately without
+    /// needing to obtain a value (by decoding one from disk or by executing
+    /// the query).
+    ///
+    /// (As with all query calls, execution is also skipped if the query result
+    /// is already cached in memory.)
+    ///
+    /// ## WARNING
+    /// A subsequent normal call to the same query might still cause it to be
+    /// executed! This can occur when the inputs are all green, but the query's
+    /// result is not cached on disk, so the query must be executed to obtain a
+    /// return value.
+    ///
+    /// Therefore, this call mode is not appropriate for callers that want to
+    /// ensure that the query is _never_ executed in the future.
+    #[inline(always)]
+    pub fn ensure_ok(self) -> TyCtxtEnsureOk<'tcx> {
+        TyCtxtEnsureOk { tcx: self }
+    }
+
+    /// This is a variant of `ensure_ok` only usable with queries that return
+    /// `Result<_, ErrorGuaranteed>`. Queries calls through this function will
+    /// return `Result<(), ErrorGuaranteed>`. I.e. the error status is returned
+    /// but nothing else. As with `ensure_ok`, this can be more efficient than
+    /// a normal query call.
+    #[inline(always)]
+    pub fn ensure_result(self) -> TyCtxtEnsureResult<'tcx> {
+        TyCtxtEnsureResult { tcx: self }
+    }
+
+    /// Wrapper that calls queries where callers don't need the return value and
+    /// just want to guarantee that the query won't be executed in the future.
+    ///
+    /// This is useful for queries that read from a [`Steal`] value, to ensure
+    /// that they are executed before the query that will steal the value.
+    ///
+    /// Currently this causes the query to be executed normally, but this behavior may change.
+    ///
+    /// [`Steal`]: rustc_data_structures::steal::Steal
+    #[inline(always)]
+    pub fn ensure_done(self) -> TyCtxtEnsureDone<'tcx> {
+        TyCtxtEnsureDone { tcx: self }
+    }
+}
 
 /// Checks whether there is already a value for this key in the in-memory
 /// query cache, returning that value if present.

compiler/rustc_middle/src/query/job.rs

@@ -4,9 +4,11 @@ use std::num::NonZero;
 use std::sync::Arc;
 
 use parking_lot::{Condvar, Mutex};
+use rustc_data_structures::hash_table::HashTable;
+use rustc_data_structures::sharded::Sharded;
 use rustc_span::Span;
 
-use crate::query::Cycle;
+use crate::queries::TaggedQueryKey;
 
 /// A value uniquely identifying an active query job.
 #[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
@@ -50,12 +52,67 @@ impl<'tcx> QueryJob<'tcx> {
     }
 }
 
+/// For a particular query and key, tracks the status of a query evaluation
+/// that has started, but has not yet finished successfully.
+///
+/// (Successful query evaluation for a key is represented by an entry in the
+/// query's in-memory cache.)
+pub enum ActiveKeyStatus<'tcx> {
+    /// Some thread is already evaluating the query for this key.
+    ///
+    /// The enclosed [`QueryJob`] can be used to wait for it to finish.
+    Started(QueryJob<'tcx>),
+
+    /// The query panicked. Queries trying to wait on this will raise a fatal error which will
+    /// silently panic.
+    Poisoned,
+}
+
+/// For a particular query, keeps track of "active" keys, i.e. keys whose
+/// evaluation has started but has not yet finished successfully.
+///
+/// (Successful query evaluation for a key is represented by an entry in the
+/// query's in-memory cache.)
+pub struct QueryState<'tcx, K> {
+    pub active: Sharded<HashTable<(K, ActiveKeyStatus<'tcx>)>>,
+}
+
+impl<'tcx, K> Default for QueryState<'tcx, K> {
+    fn default() -> QueryState<'tcx, K> {
+        QueryState { active: Default::default() }
+    }
+}
+
+/// Description of a frame in the query stack.
+///
+/// This is mostly used in case of cycles for error reporting.
+#[derive(Debug)]
+pub struct QueryStackFrame<'tcx> {
+    pub span: Span,
+
+    /// The query and key of the query method call that this stack frame
+    /// corresponds to.
+    ///
+    /// Code that doesn't care about the specific key can still use this to
+    /// check which query it's for, or obtain the query's name.
+    pub tagged_key: TaggedQueryKey<'tcx>,
+}
+
+#[derive(Debug)]
+pub struct QueryCycle<'tcx> {
+    /// The query and related span that uses the cycle.
+    pub usage: Option<QueryStackFrame<'tcx>>,
+
+    /// The span here corresponds to the reason for which this query was required.
+    pub frames: Vec<QueryStackFrame<'tcx>>,
+}
+
 #[derive(Debug)]
 pub struct QueryWaiter<'tcx> {
     pub parent: Option<QueryJobId>,
     pub condvar: Condvar,
     pub span: Span,
-    pub cycle: Mutex<Option<Cycle<'tcx>>>,
+    pub cycle: Mutex<Option<QueryCycle<'tcx>>>,
 }
 
 #[derive(Clone, Debug)]
@@ -70,7 +127,7 @@ impl<'tcx> QueryLatch<'tcx> {
     }
 
     /// Awaits for the query job to complete.
-    pub fn wait_on(&self, query: Option<QueryJobId>, span: Span) -> Result<(), Cycle<'tcx>> {
+    pub fn wait_on(&self, query: Option<QueryJobId>, span: Span) -> Result<(), QueryCycle<'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

@@ -1,34 +1,22 @@
-use rustc_hir::def_id::LocalDefId;
-
 pub use self::caches::{DefIdCache, DefaultCache, QueryCache, SingleCache, VecCache};
+pub use self::calls::{TyCtxtAt, TyCtxtEnsureDone, TyCtxtEnsureOk, TyCtxtEnsureResult};
 pub use self::into_query_key::IntoQueryKey;
-pub use self::job::{QueryJob, QueryJobId, QueryLatch, QueryWaiter};
-pub use self::keys::{LocalCrate, QueryKey};
-pub use self::plumbing::{
-    ActiveKeyStatus, Cycle, QueryMode, QueryState, QuerySystem, QueryVTable, TyCtxtAt,
-    TyCtxtEnsureDone, TyCtxtEnsureOk, TyCtxtEnsureResult,
+pub use self::job::{
+    ActiveKeyStatus, QueryCycle, QueryJob, QueryJobId, QueryLatch, QueryStackFrame, QueryState,
+    QueryWaiter,
 };
-pub use self::stack::QueryStackFrame;
+pub use self::keys::{LocalCrate, QueryKey};
+pub use self::system::{QueryMode, QuerySystem, QueryVTable};
 pub use crate::queries::Providers;
-use crate::ty::TyCtxt;
 
 pub(crate) mod arena_cached;
 mod caches;
+pub(crate) mod calls;
 pub mod erase;
-pub(crate) mod inner;
 mod into_query_key;
 mod job;
 mod keys;
 pub(crate) mod modifiers;
 pub mod on_disk_cache;
-pub(crate) mod plumbing;
-mod stack;
-
-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))
-    }
-}
+pub(crate) mod query_api;
+mod system;

compiler/rustc_middle/src/query/on_disk_cache.rs

@@ -114,11 +114,11 @@ struct Footer {
 struct SourceFileIndex(u32);
 
 #[derive(Copy, Clone, Debug, Hash, Eq, PartialEq, Encodable, Decodable)]
-pub struct AbsoluteBytePos(u64);
+struct AbsoluteBytePos(u64);
 
 impl AbsoluteBytePos {
     #[inline]
-    pub fn new(pos: usize) -> AbsoluteBytePos {
+    fn new(pos: usize) -> AbsoluteBytePos {
         AbsoluteBytePos(pos.try_into().expect("Incremental cache file size overflowed u64."))
     }
 
@@ -335,13 +335,6 @@ impl OnDiskCache {
         side_effect
     }
 
-    /// Returns true if there is a disk-cached query return value for the given node.
-    #[inline]
-    pub fn loadable_from_disk(&self, dep_node_index: SerializedDepNodeIndex) -> bool {
-        self.query_values_index.contains_key(&dep_node_index)
-        // with_decoder is infallible, so we can stop here
-    }
-
     /// Returns the disk-cached query return value for the given node, if there is one.
     pub fn try_load_query_value<'tcx, T>(
         &self,