Merge {get,ensure,force}_query into call_query.

Author: cjgillot

compiler/rustc_macros/src/query.rs

@@ -500,11 +500,11 @@ pub fn rustc_queries(input: TokenStream) -> TokenStream {
                 ::rustc_middle::dep_graph::DepKind::#name => {
                     if <#arg as DepNodeParams<TyCtxt<'_>>>::can_reconstruct_query_key() {
                         if let Some(key) = <#arg as DepNodeParams<TyCtxt<'_>>>::recover($tcx, $dep_node) {
-                            force_query::<crate::ty::query::queries::#name<'_>, _>(
+                            crate::ty::query::queries::#name::query(
                                 $tcx,
-                                key,
                                 DUMMY_SP,
-                                *$dep_node
+                                key,
+                                QueryCaller::Force(*$dep_node),
                             );
                             return true;
                         }

compiler/rustc_middle/src/ty/query/plumbing.rs

@@ -338,7 +338,7 @@ macro_rules! define_queries_inner {
         // HACK(eddyb) this is like the `impl QueryConfig for queries::$name`
         // below, but using type aliases instead of associated types, to bypass
         // the limitations around normalizing under HRTB - for example, this:
-        // `for<'tcx> fn(...) -> <queries::$name<'tcx> as QueryConfig<TyCtxt<'tcx>>>::Value`
+        // `for<'tcx> fn(...) -> <queries::$name<'tcx> as QueryConfig>::Value`
         // doesn't currently normalize to `for<'tcx> fn(...) -> query_values::$name<'tcx>`.
         // This is primarily used by the `provide!` macro in `rustc_metadata`.
         #[allow(nonstandard_style, unused_lifetimes)]
@@ -402,6 +402,19 @@ macro_rules! define_queries_inner {
             ) -> Self::Value {
                 handle_cycle_error!([$($modifiers)*][tcx, error])
             }
+        }
+
+        impl queries::$name<$tcx> {
+            $(#[$attr])*
+            #[inline(always)]
+            pub fn query(
+                tcx: TyCtxt<$tcx>,
+                span: Span,
+                key: query_keys::$name<$tcx>,
+                caller: QueryCaller<DepKind>,
+            ) -> Option<<queries::$name<$tcx> as QueryConfig>::Stored> {
+                call_query::<queries::$name<'_>, _>(tcx, span, key, caller)
+            }
         })*
 
         #[derive(Copy, Clone)]
@@ -413,7 +426,12 @@ macro_rules! define_queries_inner {
             $($(#[$attr])*
             #[inline(always)]
             pub fn $name(self, key: query_helper_param_ty!($($K)*)) {
-                ensure_query::<queries::$name<'_>, _>(self.tcx, key.into_query_param())
+                queries::$name::query(
+                    self.tcx,
+                    DUMMY_SP,
+                    key.into_query_param(),
+                    QueryCaller::Ensure,
+                );
             })*
         }
 
@@ -496,7 +514,13 @@ macro_rules! define_queries_inner {
             pub fn $name(self, key: query_helper_param_ty!($($K)*))
                 -> <queries::$name<$tcx> as QueryConfig>::Stored
             {
-                get_query::<queries::$name<'_>, _>(self.tcx, self.span, key.into_query_param())
+                let ret = queries::$name::query(
+                    self.tcx,
+                    self.span,
+                    key.into_query_param(),
+                    QueryCaller::Get,
+                );
+                ret.unwrap()
             })*
         }
 

compiler/rustc_query_system/src/query/plumbing.rs

@@ -729,34 +729,55 @@ fn force_query_impl<CTX, C>(
     );
 }
 
-#[inline(always)]
-pub fn get_query<Q, CTX>(tcx: CTX, span: Span, key: Q::Key) -> Q::Stored
-where
-    Q: QueryDescription<CTX>,
-    Q::Key: crate::dep_graph::DepNodeParams<CTX>,
-    CTX: QueryContext,
-{
-    debug!("ty::query::get_query<{}>(key={:?}, span={:?})", Q::NAME, key, span);
-
-    get_query_impl(tcx, Q::query_state(tcx), span, key, &Q::VTABLE)
+pub enum QueryCaller<DK> {
+    Ensure,
+    Get,
+    Force(DepNode<DK>),
 }
 
-#[inline(always)]
-pub fn ensure_query<Q, CTX>(tcx: CTX, key: Q::Key)
+#[inline(never)]
+fn call_query_impl<CTX, C>(
+    tcx: CTX,
+    state: &QueryState<CTX::DepKind, CTX::Query, C>,
+    span: Span,
+    key: C::Key,
+    caller: QueryCaller<CTX::DepKind>,
+    query: &QueryVtable<CTX, C::Key, C::Value>,
+) -> Option<C::Stored>
 where
-    Q: QueryDescription<CTX>,
-    Q::Key: crate::dep_graph::DepNodeParams<CTX>,
+    C: QueryCache,
+    C::Key: Eq + Clone + crate::dep_graph::DepNodeParams<CTX>,
+    C::Stored: Clone,
     CTX: QueryContext,
 {
-    ensure_query_impl(tcx, Q::query_state(tcx), key, &Q::VTABLE)
+    match caller {
+        QueryCaller::Ensure => {
+            ensure_query_impl(tcx, state, key, query);
+            None
+        }
+        QueryCaller::Get => {
+            let ret = get_query_impl(tcx, state, span, key, query);
+            Some(ret)
+        }
+        QueryCaller::Force(dep_node) => {
+            force_query_impl(tcx, state, key, span, dep_node, query);
+            None
+        }
+    }
 }
 
 #[inline(always)]
-pub fn force_query<Q, CTX>(tcx: CTX, key: Q::Key, span: Span, dep_node: DepNode<CTX::DepKind>)
+pub fn call_query<Q, CTX>(
+    tcx: CTX,
+    span: Span,
+    key: Q::Key,
+    caller: QueryCaller<CTX::DepKind>,
+) -> Option<Q::Stored>
 where
     Q: QueryDescription<CTX>,
     Q::Key: crate::dep_graph::DepNodeParams<CTX>,
     CTX: QueryContext,
 {
-    force_query_impl(tcx, Q::query_state(tcx), key, span, dep_node, &Q::VTABLE)
+    debug!("ty::query::get_query<{}>(key={:?}, span={:?})", Q::NAME, key, span);
+    call_query_impl(tcx, Q::query_state(tcx), span, key, caller, &Q::VTABLE)
 }