Only instantiate impl args

2 files changed, 70 insertions, 46 deletions

diff --git a/compiler/rustc_trait_selection/src/solve/inspect/analyse.rs b/compiler/rustc_trait_selection/src/solve/inspect/analyse.rs
index 1193a9059ca..e2282a5bfbc 100644
--- a/compiler/rustc_trait_selection/src/solve/inspect/analyse.rs
+++ b/compiler/rustc_trait_selection/src/solve/inspect/analyse.rs
@@ -133,45 +133,25 @@ 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`.
+    /// See [`Self::instantiate_impl_args`] if you need the impl args too.
     #[instrument(
         level = "debug",
         skip_all,
         fields(goal = ?self.goal.goal, steps = ?self.steps)
     )]
-    pub fn instantiate_nested_goals_and_opt_impl_args(
-        &self,
-        span: Span,
-    ) -> (Vec<InspectGoal<'a, 'tcx>>, Option<ty::GenericArgsRef<'tcx>>) {
+    pub fn instantiate_nested_goals(&self, span: Span) -> Vec<InspectGoal<'a, 'tcx>> {
         let infcx = self.goal.infcx;
         let param_env = self.goal.goal.param_env;
         let mut orig_values = self.goal.orig_values.to_vec();
 
         let mut instantiated_goals = vec![];
-        let mut opt_impl_args = None;
         for step in &self.steps {
             match **step {
                 inspect::ProbeStep::AddGoal(source, goal) => instantiated_goals.push((
                     source,
                     instantiate_canonical_state(infcx, span, param_env, &mut orig_values, goal),
                 )),
-                inspect::ProbeStep::RecordImplArgs { impl_args } => {
-                    opt_impl_args = Some(instantiate_canonical_state(
-                        infcx,
-                        span,
-                        param_env,
-                        &mut orig_values,
-                        impl_args,
-                    ));
-                }
+                inspect::ProbeStep::RecordImplArgs { .. } => {}
                 inspect::ProbeStep::MakeCanonicalResponse { .. }
                 | inspect::ProbeStep::NestedProbe(_) => unreachable!(),
             }
@@ -187,14 +167,59 @@ impl<'a, 'tcx> InspectCandidate<'a, 'tcx> {
             let _ = term_hack.constrain(infcx, span, param_env);
         }
 
-        let opt_impl_args = opt_impl_args.map(|impl_args| eager_resolve_vars(infcx, impl_args));
-
-        let goals = instantiated_goals
+        instantiated_goals
             .into_iter()
             .map(|(source, goal)| self.instantiate_proof_tree_for_nested_goal(source, goal, span))
-            .collect();
+            .collect()
+    }
+
+    /// Instantiate the args of an impl if this candidate came from a
+    /// `CandidateSource::Impl`. This function modifies the state of the
+    /// `infcx`.
+    #[instrument(
+        level = "debug",
+        skip_all,
+        fields(goal = ?self.goal.goal, steps = ?self.steps)
+    )]
+    pub fn instantiate_impl_args(&self, span: Span) -> 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();
+
+        for step in &self.steps {
+            match **step {
+                inspect::ProbeStep::RecordImplArgs { impl_args } => {
+                    let impl_args = instantiate_canonical_state(
+                        infcx,
+                        span,
+                        param_env,
+                        &mut orig_values,
+                        impl_args,
+                    );
+
+                    let () = instantiate_canonical_state(
+                        infcx,
+                        span,
+                        param_env,
+                        &mut orig_values,
+                        self.final_state,
+                    );
+
+                    // No reason we couldn't support this, but we don't need to for select.
+                    assert!(
+                        self.goal.normalizes_to_term_hack.is_none(),
+                        "cannot use `instantiate_impl_args` with a `NormalizesTo` goal"
+                    );
+
+                    return eager_resolve_vars(infcx, impl_args);
+                }
+                inspect::ProbeStep::AddGoal(..) => {}
+                inspect::ProbeStep::MakeCanonicalResponse { .. }
+                | inspect::ProbeStep::NestedProbe(_) => unreachable!(),
+            }
+        }
 
-        (goals, opt_impl_args)
+        bug!("expected impl args probe step for `instantiate_impl_args`");
     }
 
     pub fn instantiate_proof_tree_for_nested_goal(
diff --git a/compiler/rustc_trait_selection/src/solve/select.rs b/compiler/rustc_trait_selection/src/solve/select.rs
index 09d00136f53..fb1adc2fd2a 100644
--- a/compiler/rustc_trait_selection/src/solve/select.rs
+++ b/compiler/rustc_trait_selection/src/solve/select.rs
@@ -10,7 +10,7 @@ use rustc_infer::traits::{
 use rustc_macros::extension;
 use rustc_middle::{bug, span_bug};
 use rustc_span::Span;
-use thin_vec::ThinVec;
+use thin_vec::thin_vec;
 
 use crate::solve::inspect::{self, ProofTreeInferCtxtExt};
 
@@ -147,22 +147,21 @@ fn to_selection<'tcx>(
         return None;
     }
 
-    let (nested, impl_args) = cand.instantiate_nested_goals_and_opt_impl_args(span);
-    let mut nested: ThinVec<_> = nested
-        .into_iter()
-        .map(|nested| {
-            Obligation::new(
-                nested.infcx().tcx,
-                ObligationCause::dummy_with_span(span),
-                nested.goal().param_env,
-                nested.goal().predicate,
-            )
-        })
-        .collect();
-
-    if let Ok(Certainty::Yes) = cand.result() {
-        nested.clear();
-    }
+    let nested = match cand.result().expect("expected positive result") {
+        Certainty::Yes => thin_vec![],
+        Certainty::Maybe(_) => cand
+            .instantiate_nested_goals(span)
+            .into_iter()
+            .map(|nested| {
+                Obligation::new(
+                    nested.infcx().tcx,
+                    ObligationCause::dummy_with_span(span),
+                    nested.goal().param_env,
+                    nested.goal().predicate,
+                )
+            })
+            .collect(),
+    };
 
     Some(match cand.kind() {
         ProbeKind::TraitCandidate { source, result: _ } => match source {
@@ -171,7 +170,7 @@ fn to_selection<'tcx>(
                 // For impl candidates, we do the rematch manually to compute the args.
                 ImplSource::UserDefined(ImplSourceUserDefinedData {
                     impl_def_id,
-                    args: impl_args.expect("expected recorded impl args for impl candidate"),
+                    args: cand.instantiate_impl_args(span),
                     nested,
                 })
             }