split probe into 2 functions for better readability

lcnr · lcnr · commit f5438d658ff1 · 2023-06-20T12:40:43.000+02:00
diff --git a/compiler/rustc_trait_selection/src/solve/alias_relate.rs b/compiler/rustc_trait_selection/src/solve/alias_relate.rs
@@ -110,12 +110,11 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
         direction: ty::AliasRelationDirection,
         invert: Invert,
     ) -> QueryResult<'tcx> {
-        self.probe(
+        self.probe(|r| CandidateKind::Candidate { name: "normalizes-to".into(), result: *r }).enter(
             |ecx| {
                 ecx.normalizes_to_inner(param_env, alias, other, direction, invert)?;
                 ecx.evaluate_added_goals_and_make_canonical_response(Certainty::Yes)
             },
-            |r| CandidateKind::Candidate { name: "normalizes-to".into(), result: *r },
         )
     }
 
@@ -157,7 +156,7 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
         alias_rhs: ty::AliasTy<'tcx>,
         direction: ty::AliasRelationDirection,
     ) -> QueryResult<'tcx> {
-        self.probe(
+        self.probe(|r| CandidateKind::Candidate { name: "substs relate".into(), result: *r }).enter(
             |ecx| {
                 match direction {
                     ty::AliasRelationDirection::Equate => {
@@ -170,7 +169,6 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
 
                 ecx.evaluate_added_goals_and_make_canonical_response(Certainty::Yes)
             },
-            |r| CandidateKind::Candidate { name: "substs relate".into(), result: *r },
         )
     }
 
@@ -181,8 +179,8 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
         rhs: ty::Term<'tcx>,
         direction: ty::AliasRelationDirection,
     ) -> QueryResult<'tcx> {
-        self.probe(
-            |ecx| {
+        self.probe(|r| CandidateKind::Candidate { name: "bidir normalizes-to".into(), result: *r })
+            .enter(|ecx| {
                 ecx.normalizes_to_inner(
                     param_env,
                     lhs.to_alias_ty(ecx.tcx()).unwrap(),
@@ -198,8 +196,6 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
                     Invert::Yes,
                 )?;
                 ecx.evaluate_added_goals_and_make_canonical_response(Certainty::Yes)
-            },
-            |r| CandidateKind::Candidate { name: "bidir normalizes-to".into(), result: *r },
-        )
+            })
     }
 }
diff --git a/compiler/rustc_trait_selection/src/solve/assembly/mod.rs b/compiler/rustc_trait_selection/src/solve/assembly/mod.rs
@@ -337,8 +337,8 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
             return
         };
 
-        let normalized_self_candidates: Result<_, NoSolution> = self.probe(
-            |ecx| {
+        let normalized_self_candidates: Result<_, NoSolution> =
+            self.probe(|_| CandidateKind::NormalizedSelfTyAssembly).enter(|ecx| {
                 ecx.with_incremented_depth(
                     |ecx| {
                         let result = ecx.evaluate_added_goals_and_make_canonical_response(
@@ -368,9 +368,7 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
                         Ok(ecx.assemble_and_evaluate_candidates(goal))
                     },
                 )
-            },
-            |_| CandidateKind::NormalizedSelfTyAssembly,
-        );
+            });
 
         if let Ok(normalized_self_candidates) = normalized_self_candidates {
             candidates.extend(normalized_self_candidates);
diff --git a/compiler/rustc_trait_selection/src/solve/eval_ctxt.rs b/compiler/rustc_trait_selection/src/solve/eval_ctxt.rs
@@ -30,6 +30,7 @@ use super::SolverMode;
 use super::{search_graph::SearchGraph, Goal};
 
 mod canonical;
+mod probe;
 
 pub struct EvalCtxt<'a, 'tcx> {
     /// The inference context that backs (mostly) inference and placeholder terms
@@ -529,32 +530,6 @@ impl<'a, 'tcx> EvalCtxt<'a, 'tcx> {
 }
 
 impl<'tcx> EvalCtxt<'_, 'tcx> {
-    /// `probe_kind` is only called when proof tree building is enabled so it can be
-    /// as expensive as necessary to output the desired information.
-    pub(super) fn probe<T>(
-        &mut self,
-        f: impl FnOnce(&mut EvalCtxt<'_, 'tcx>) -> T,
-        probe_kind: impl FnOnce(&T) -> CandidateKind<'tcx>,
-    ) -> T {
-        let mut ecx = EvalCtxt {
-            infcx: self.infcx,
-            var_values: self.var_values,
-            predefined_opaques_in_body: self.predefined_opaques_in_body,
-            max_input_universe: self.max_input_universe,
-            search_graph: self.search_graph,
-            nested_goals: self.nested_goals.clone(),
-            tainted: self.tainted,
-            inspect: self.inspect.new_goal_candidate(),
-        };
-        let r = self.infcx.probe(|_| f(&mut ecx));
-        if !self.inspect.is_noop() {
-            let cand_kind = probe_kind(&r);
-            ecx.inspect.candidate_kind(cand_kind);
-            self.inspect.goal_candidate(ecx.inspect);
-        }
-        r
-    }
-
     pub(super) fn tcx(&self) -> TyCtxt<'tcx> {
         self.infcx.tcx
     }
@@ -866,17 +841,20 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
             if candidate_key.def_id != key.def_id {
                 continue;
             }
-            values.extend(self.probe(
-                |ecx| {
+            values.extend(
+                self.probe(|r| CandidateKind::Candidate {
+                    name: "opaque type storage".into(),
+                    result: *r,
+                })
+                .enter(|ecx| {
                     for (a, b) in std::iter::zip(candidate_key.substs, key.substs) {
                         ecx.eq(param_env, a, b)?;
                     }
                     ecx.eq(param_env, candidate_ty, ty)?;
                     ecx.add_item_bounds_for_hidden_type(candidate_key, param_env, candidate_ty);
                     ecx.evaluate_added_goals_and_make_canonical_response(Certainty::Yes)
-                },
-                |r| CandidateKind::Candidate { name: "opaque type storage".into(), result: *r },
-            ));
+                }),
+            );
         }
         values
     }
diff --git a/compiler/rustc_trait_selection/src/solve/eval_ctxt/probe.rs b/compiler/rustc_trait_selection/src/solve/eval_ctxt/probe.rs
@@ -0,0 +1,47 @@
+use super::EvalCtxt;
+use rustc_middle::traits::solve::inspect;
+use std::marker::PhantomData;
+
+pub(in crate::solve) struct ProbeCtxt<'me, 'a, 'tcx, F, T> {
+    ecx: &'me mut EvalCtxt<'a, 'tcx>,
+    probe_kind: F,
+    _result: PhantomData<T>,
+}
+
+impl<'tcx, F, T> ProbeCtxt<'_, '_, 'tcx, F, T>
+where
+    F: FnOnce(&T) -> inspect::CandidateKind<'tcx>,
+{
+    pub(in crate::solve) fn enter(self, f: impl FnOnce(&mut EvalCtxt<'_, 'tcx>) -> T) -> T {
+        let ProbeCtxt { ecx: outer_ecx, probe_kind, _result } = self;
+
+        let mut nested_ecx = EvalCtxt {
+            infcx: outer_ecx.infcx,
+            var_values: outer_ecx.var_values,
+            predefined_opaques_in_body: outer_ecx.predefined_opaques_in_body,
+            max_input_universe: outer_ecx.max_input_universe,
+            search_graph: outer_ecx.search_graph,
+            nested_goals: outer_ecx.nested_goals.clone(),
+            tainted: outer_ecx.tainted,
+            inspect: outer_ecx.inspect.new_goal_candidate(),
+        };
+        let r = nested_ecx.infcx.probe(|_| f(&mut nested_ecx));
+        if !outer_ecx.inspect.is_noop() {
+            let cand_kind = probe_kind(&r);
+            nested_ecx.inspect.candidate_kind(cand_kind);
+            outer_ecx.inspect.goal_candidate(nested_ecx.inspect);
+        }
+        r
+    }
+}
+
+impl<'a, 'tcx> EvalCtxt<'a, 'tcx> {
+    /// `probe_kind` is only called when proof tree building is enabled so it can be
+    /// as expensive as necessary to output the desired information.
+    pub(in crate::solve) fn probe<F, T>(&mut self, probe_kind: F) -> ProbeCtxt<'_, 'a, 'tcx, F, T>
+    where
+        F: FnOnce(&T) -> inspect::CandidateKind<'tcx>,
+    {
+        ProbeCtxt { ecx: self, probe_kind, _result: PhantomData }
+    }
+}
diff --git a/compiler/rustc_trait_selection/src/solve/project_goals.rs b/compiler/rustc_trait_selection/src/solve/project_goals.rs
@@ -112,8 +112,8 @@ impl<'tcx> assembly::GoalKind<'tcx> for ProjectionPredicate<'tcx> {
     ) -> QueryResult<'tcx> {
         if let Some(projection_pred) = assumption.as_projection_clause() {
             if projection_pred.projection_def_id() == goal.predicate.def_id() {
-                ecx.probe(
-                    |ecx| {
+                ecx.probe(|r| CandidateKind::Candidate { name: "assumption".into(), result: *r })
+                    .enter(|ecx| {
                         let assumption_projection_pred =
                             ecx.instantiate_binder_with_infer(projection_pred);
                         ecx.eq(
@@ -128,9 +128,7 @@ impl<'tcx> assembly::GoalKind<'tcx> for ProjectionPredicate<'tcx> {
                         )
                         .expect("expected goal term to be fully unconstrained");
                         then(ecx)
-                    },
-                    |r| CandidateKind::Candidate { name: "assumption".into(), result: *r },
-                )
+                    })
             } else {
                 Err(NoSolution)
             }
@@ -154,6 +152,7 @@ impl<'tcx> assembly::GoalKind<'tcx> for ProjectionPredicate<'tcx> {
         }
 
         ecx.probe(
+            |r| CandidateKind::Candidate { name: "impl".into(), result: *r }).enter(
             |ecx| {
                 let impl_substs = ecx.fresh_substs_for_item(impl_def_id);
                 let impl_trait_ref = impl_trait_ref.subst(tcx, impl_substs);
@@ -235,7 +234,6 @@ impl<'tcx> assembly::GoalKind<'tcx> for ProjectionPredicate<'tcx> {
                     .expect("expected goal term to be fully unconstrained");
                 ecx.evaluate_added_goals_and_make_canonical_response(Certainty::Yes)
             },
-            |r| CandidateKind::Candidate { name: "impl".into(), result: *r },
         )
     }
 
@@ -331,8 +329,8 @@ impl<'tcx> assembly::GoalKind<'tcx> for ProjectionPredicate<'tcx> {
         goal: Goal<'tcx, Self>,
     ) -> QueryResult<'tcx> {
         let tcx = ecx.tcx();
-        ecx.probe(
-            |ecx| {
+        ecx.probe(|r| CandidateKind::Candidate { name: "builtin pointee".into(), result: *r })
+            .enter(|ecx| {
                 let metadata_ty = match goal.predicate.self_ty().kind() {
                     ty::Bool
                     | ty::Char
@@ -415,9 +413,7 @@ impl<'tcx> assembly::GoalKind<'tcx> for ProjectionPredicate<'tcx> {
                 ecx.eq(goal.param_env, goal.predicate.term, metadata_ty.into())
                     .expect("expected goal term to be fully unconstrained");
                 ecx.evaluate_added_goals_and_make_canonical_response(Certainty::Yes)
-            },
-            |r| CandidateKind::Candidate { name: "builtin pointee".into(), result: *r },
-        )
+            })
     }
 
     fn consider_builtin_future_candidate(
@@ -552,14 +548,15 @@ impl<'tcx> assembly::GoalKind<'tcx> for ProjectionPredicate<'tcx> {
             ),
         };
 
-        ecx.probe(
-            |ecx| {
-                ecx.eq(goal.param_env, goal.predicate.term, discriminant_ty.into())
-                    .expect("expected goal term to be fully unconstrained");
-                ecx.evaluate_added_goals_and_make_canonical_response(Certainty::Yes)
-            },
-            |r| CandidateKind::Candidate { name: "builtin discriminant kind".into(), result: *r },
-        )
+        ecx.probe(|r| CandidateKind::Candidate {
+            name: "builtin discriminant kind".into(),
+            result: *r,
+        })
+        .enter(|ecx| {
+            ecx.eq(goal.param_env, goal.predicate.term, discriminant_ty.into())
+                .expect("expected goal term to be fully unconstrained");
+            ecx.evaluate_added_goals_and_make_canonical_response(Certainty::Yes)
+        })
     }
 
     fn consider_builtin_destruct_candidate(
diff --git a/compiler/rustc_trait_selection/src/solve/trait_goals.rs b/compiler/rustc_trait_selection/src/solve/trait_goals.rs
