combine both approaches

Author: jyn514

compiler/rustc_query_impl/src/plumbing.rs

@@ -104,7 +104,7 @@ impl QueryContext for QueryCtxt<'_> {
         token: QueryJobId,
         depth_limit: bool,
         diagnostics: Option<&Lock<ThinVec<Diagnostic>>>,
-        compute: impl FnOnce() -> R,
+        compute: &mut dyn FnMut() -> R,
     ) -> R {
         // The `TyCtxt` stored in TLS has the same global interner lifetime
         // as `self`, so we use `with_related_context` to relate the 'tcx lifetimes

compiler/rustc_query_system/src/query/mod.rs

@@ -128,7 +128,7 @@ pub trait QueryContext: HasDepContext {
         token: QueryJobId,
         depth_limit: bool,
         diagnostics: Option<&Lock<ThinVec<Diagnostic>>>,
-        compute: impl FnOnce() -> R,
+        compute: &mut dyn FnMut() -> R,
     ) -> R;
 
     fn depth_limit_error(self, job: QueryJobId);

compiler/rustc_query_system/src/query/plumbing.rs

@@ -421,7 +421,8 @@ where
         }
 
         let prof_timer = qcx.dep_context().profiler().query_provider();
-        let result = qcx.start_query(job_id, query.depth_limit(), None, || query.compute(qcx, key));
+        let result =
+            qcx.start_query(job_id, query.depth_limit(), None, &mut || query.compute(qcx, key));
         let dep_node_index = dep_graph.next_virtual_depnode_index();
         prof_timer.finish_with_query_invocation_id(dep_node_index.into());
 
@@ -445,7 +446,7 @@ where
 
         // The diagnostics for this query will be promoted to the current session during
         // `try_mark_green()`, so we can ignore them here.
-        if let Some(ret) = qcx.start_query(job_id, false, None, || {
+        if let Some(ret) = qcx.start_query(job_id, false, None, &mut || {
             try_load_from_disk_and_cache_in_memory(query, qcx, &key, &dep_node)
         }) {
             return ret;
@@ -456,7 +457,7 @@ where
     let diagnostics = Lock::new(ThinVec::new());
 
     let (result, dep_node_index) =
-        qcx.start_query(job_id, query.depth_limit(), Some(&diagnostics), || {
+        qcx.start_query(job_id, query.depth_limit(), Some(&diagnostics), &mut || {
             if query.anon() {
                 return dep_graph.with_anon_task(*qcx.dep_context(), query.dep_kind(), || {
                     query.compute(qcx, key)