Record impl args in the InsepctCandiate rather than rematching during select

compiler/rustc_trait_selection/src/solve/inspect/analyse.rs

@@ -93,6 +93,7 @@ pub struct InspectCandidate<'a, 'tcx> {
     kind: inspect::ProbeKind<'tcx>,
     nested_goals: Vec<(GoalSource, inspect::CanonicalState<'tcx, Goal<'tcx, ty::Predicate<'tcx>>>)>,
     final_state: inspect::CanonicalState<'tcx, ()>,
+    impl_args: Option<inspect::CanonicalState<'tcx, ty::GenericArgsRef<'tcx>>>,
     result: QueryResult<'tcx>,
     shallow_certainty: Certainty,
 }
@@ -135,7 +136,20 @@ impl<'a, 'tcx> InspectCandidate<'a, 'tcx> {
 
     /// Instantiate the nested goals for the candidate without rolling back their
     /// inference constraints. This function modifies the state of the `infcx`.
+    ///
+    /// See [`Self::instantiate_nested_goals_and_opt_impl_args`] if you need the impl args too.
     pub fn instantiate_nested_goals(&self, span: Span) -> Vec<InspectGoal<'a, 'tcx>> {
+        self.instantiate_nested_goals_and_opt_impl_args(span).0
+    }
+
+    /// Instantiate the nested goals for the candidate without rolling back their
+    /// inference constraints, and optionally the args of an impl if this candidate
+    /// came from a `CandidateSource::Impl`. This function modifies the state of the
+    /// `infcx`.
+    pub fn instantiate_nested_goals_and_opt_impl_args(
+        &self,
+        span: Span,
+    ) -> (Vec<InspectGoal<'a, 'tcx>>, Option<ty::GenericArgsRef<'tcx>>) {
         let infcx = self.goal.infcx;
         let param_env = self.goal.goal.param_env;
         let mut orig_values = self.goal.orig_values.to_vec();
@@ -164,6 +178,17 @@ impl<'a, 'tcx> InspectCandidate<'a, 'tcx> {
             self.final_state,
         );
 
+        let impl_args = self.impl_args.map(|impl_args| {
+            canonical::instantiate_canonical_state(
+                infcx,
+                span,
+                param_env,
+                &mut orig_values,
+                impl_args,
+            )
+            .fold_with(&mut EagerResolver::new(infcx))
+        });
+
         if let Some(term_hack) = self.goal.normalizes_to_term_hack {
             // FIXME: We ignore the expected term of `NormalizesTo` goals
             // when computing the result of its candidates. This is
@@ -171,7 +196,7 @@ impl<'a, 'tcx> InspectCandidate<'a, 'tcx> {
             let _ = term_hack.constrain(infcx, span, param_env);
         }
 
-        instantiated_goals
+        let goals = instantiated_goals
             .into_iter()
             .map(|(source, goal)| match goal.predicate.kind().no_bound_vars() {
                 Some(ty::PredicateKind::NormalizesTo(ty::NormalizesTo { alias, term })) => {
@@ -208,7 +233,9 @@ impl<'a, 'tcx> InspectCandidate<'a, 'tcx> {
                     source,
                 ),
             })
-            .collect()
+            .collect();
+
+        (goals, impl_args)
     }
 
     /// Visit all nested goals of this candidate, rolling back
@@ -245,9 +272,10 @@ impl<'a, 'tcx> InspectGoal<'a, 'tcx> {
         probe: &inspect::Probe<'tcx>,
     ) {
         let mut shallow_certainty = None;
+        let mut impl_args = None;
         for step in &probe.steps {
-            match step {
-                &inspect::ProbeStep::AddGoal(source, goal) => nested_goals.push((source, goal)),
+            match *step {
+                inspect::ProbeStep::AddGoal(source, goal) => nested_goals.push((source, goal)),
                 inspect::ProbeStep::NestedProbe(ref probe) => {
                     // Nested probes have to prove goals added in their parent
                     // but do not leak them, so we truncate the added goals
@@ -257,7 +285,10 @@ impl<'a, 'tcx> InspectGoal<'a, 'tcx> {
                     nested_goals.truncate(num_goals);
                 }
                 inspect::ProbeStep::MakeCanonicalResponse { shallow_certainty: c } => {
-                    assert_eq!(shallow_certainty.replace(*c), None);
+                    assert_eq!(shallow_certainty.replace(c), None);
+                }
+                inspect::ProbeStep::RecordImplArgs { impl_args: i } => {
+                    assert_eq!(impl_args.replace(i), None);
                 }
                 inspect::ProbeStep::EvaluateGoals(_) => (),
             }
@@ -284,6 +315,7 @@ impl<'a, 'tcx> InspectGoal<'a, 'tcx> {
                         final_state: probe.final_state,
                         result,
                         shallow_certainty,
+                        impl_args,
                     });
                 }
             }

compiler/rustc_trait_selection/src/solve/inspect/build.rs

@@ -242,6 +242,7 @@ enum WipProbeStep<'tcx> {
     EvaluateGoals(WipAddedGoalsEvaluation<'tcx>),
     NestedProbe(WipProbe<'tcx>),
     MakeCanonicalResponse { shallow_certainty: Certainty },
+    RecordImplArgs { impl_args: inspect::CanonicalState<'tcx, ty::GenericArgsRef<'tcx>> },
 }
 
 impl<'tcx> WipProbeStep<'tcx> {
@@ -250,6 +251,9 @@ impl<'tcx> WipProbeStep<'tcx> {
             WipProbeStep::AddGoal(source, goal) => inspect::ProbeStep::AddGoal(source, goal),
             WipProbeStep::EvaluateGoals(eval) => inspect::ProbeStep::EvaluateGoals(eval.finalize()),
             WipProbeStep::NestedProbe(probe) => inspect::ProbeStep::NestedProbe(probe.finalize()),
+            WipProbeStep::RecordImplArgs { impl_args } => {
+                inspect::ProbeStep::RecordImplArgs { impl_args }
+            }
             WipProbeStep::MakeCanonicalResponse { shallow_certainty } => {
                 inspect::ProbeStep::MakeCanonicalResponse { shallow_certainty }
             }
@@ -534,6 +538,30 @@ impl<'tcx> ProofTreeBuilder<'tcx> {
         }
     }
 
+    pub(crate) fn record_impl_args(
+        &mut self,
+        infcx: &InferCtxt<'tcx>,
+        max_input_universe: ty::UniverseIndex,
+        impl_args: ty::GenericArgsRef<'tcx>,
+    ) {
+        match self.as_mut() {
+            Some(DebugSolver::GoalEvaluationStep(state)) => {
+                let impl_args = canonical::make_canonical_state(
+                    infcx,
+                    &state.var_values,
+                    max_input_universe,
+                    impl_args,
+                );
+                state
+                    .current_evaluation_scope()
+                    .steps
+                    .push(WipProbeStep::RecordImplArgs { impl_args });
+            }
+            None => {}
+            _ => bug!(),
+        }
+    }
+
     pub fn make_canonical_response(&mut self, shallow_certainty: Certainty) {
         match self.as_mut() {
             Some(DebugSolver::GoalEvaluationStep(state)) => {
@@ -543,7 +571,7 @@ impl<'tcx> ProofTreeBuilder<'tcx> {
                     .push(WipProbeStep::MakeCanonicalResponse { shallow_certainty });
             }
             None => {}
-            _ => {}
+            _ => bug!(),
         }
     }