From 4440e8196aee718cbd3aafa41a0919f432c06330 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?John=20K=C3=A5re=20Alsaker?= <john.kare.alsaker@gmail.com>
Date: Sat, 25 Mar 2023 11:47:23 +0100
Subject: [PATCH] Add query accessor functions

---
 compiler/rustc_middle/src/ty/query.rs         | 87 ++++++++++++-------
 .../rustc_query_system/src/query/plumbing.rs  |  2 +-
 2 files changed, 56 insertions(+), 33 deletions(-)

diff --git a/compiler/rustc_middle/src/ty/query.rs b/compiler/rustc_middle/src/ty/query.rs
index 86e80a7a01d..07d47cae5ee 100644
--- a/compiler/rustc_middle/src/ty/query.rs
+++ b/compiler/rustc_middle/src/ty/query.rs
@@ -202,6 +202,40 @@ impl<'tcx> TyCtxt<'tcx> {
     }
 }
 
+#[inline]
+fn query_get_at<'tcx, Cache>(
+    tcx: TyCtxt<'tcx>,
+    execute_query: fn(TyCtxt<'tcx>, Span, Cache::Key, QueryMode) -> Option<Cache::Value>,
+    query_cache: &Cache,
+    span: Span,
+    key: Cache::Key,
+) -> Cache::Value
+where
+    Cache: QueryCache,
+{
+    let key = key.into_query_param();
+    match try_get_cached(tcx, query_cache, &key) {
+        Some(value) => value,
+        None => execute_query(tcx, span, key, QueryMode::Get).unwrap(),
+    }
+}
+
+#[inline]
+fn query_ensure<'tcx, Cache>(
+    tcx: TyCtxt<'tcx>,
+    execute_query: fn(TyCtxt<'tcx>, Span, Cache::Key, QueryMode) -> Option<Cache::Value>,
+    query_cache: &Cache,
+    key: Cache::Key,
+    check_cache: bool,
+) where
+    Cache: QueryCache,
+{
+    let key = key.into_query_param();
+    if try_get_cached(tcx, query_cache, &key).is_none() {
+        execute_query(tcx, DUMMY_SP, key, QueryMode::Ensure { check_cache });
+    }
+}
+
 macro_rules! query_helper_param_ty {
     (DefId) => { impl IntoQueryParam<DefId> };
     (LocalDefId) => { impl IntoQueryParam<LocalDefId> };
@@ -407,17 +441,13 @@ macro_rules! define_callbacks {
             $($(#[$attr])*
             #[inline(always)]
             pub fn $name(self, key: query_helper_param_ty!($($K)*)) {
-                let key = key.into_query_param();
-
-                match try_get_cached(self.tcx, &self.tcx.query_system.caches.$name, &key) {
-                    Some(_) => return,
-                    None => (self.tcx.query_system.fns.engine.$name)(
-                        self.tcx,
-                        DUMMY_SP,
-                        key,
-                        QueryMode::Ensure { check_cache: false },
-                    ),
-                };
+                query_ensure(
+                    self.tcx,
+                    self.tcx.query_system.fns.engine.$name,
+                    &self.tcx.query_system.caches.$name,
+                    key.into_query_param(),
+                    false,
+                );
             })*
         }
 
@@ -425,17 +455,13 @@ macro_rules! define_callbacks {
             $($(#[$attr])*
             #[inline(always)]
             pub fn $name(self, key: query_helper_param_ty!($($K)*)) {
-                let key = key.into_query_param();
-
-                match try_get_cached(self.tcx, &self.tcx.query_system.caches.$name, &key) {
-                    Some(_) => return,
-                    None => (self.tcx.query_system.fns.engine.$name)(
-                        self.tcx,
-                        DUMMY_SP,
-                        key,
-                        QueryMode::Ensure { check_cache: true },
-                    ),
-                };
+                query_ensure(
+                    self.tcx,
+                    self.tcx.query_system.fns.engine.$name,
+                    &self.tcx.query_system.caches.$name,
+                    key.into_query_param(),
+                    true,
+                );
             })*
         }
 
@@ -454,16 +480,13 @@ macro_rules! define_callbacks {
             #[inline(always)]
             pub fn $name(self, key: query_helper_param_ty!($($K)*)) -> $V
             {
-                let key = key.into_query_param();
-
-                restore::<$V>(match try_get_cached(self.tcx, &self.tcx.query_system.caches.$name, &key) {
-                    Some(value) => value,
-                    None => (self.tcx.query_system.fns.engine.$name)(
-                        self.tcx,
-                        self.span,
-                        key, QueryMode::Get
-                    ).unwrap(),
-                })
+                restore::<$V>(query_get_at(
+                    self.tcx,
+                    self.tcx.query_system.fns.engine.$name,
+                    &self.tcx.query_system.caches.$name,
+                    self.span,
+                    key.into_query_param(),
+                ))
             })*
         }
 
diff --git a/compiler/rustc_query_system/src/query/plumbing.rs b/compiler/rustc_query_system/src/query/plumbing.rs
index bce01debc53..3b17c665fb7 100644
--- a/compiler/rustc_query_system/src/query/plumbing.rs
+++ b/compiler/rustc_query_system/src/query/plumbing.rs
@@ -236,7 +236,7 @@ pub(crate) struct CycleError<D: DepKind> {
 /// It returns the shard index and a lock guard to the shard,
 /// which will be used if the query is not in the cache and we need
 /// to compute it.
-#[inline]
+#[inline(always)]
 pub fn try_get_cached<Tcx, C>(tcx: Tcx, cache: &C, key: &C::Key) -> Option<C::Value>
 where
     C: QueryCache,