Record impl args in the proof tree

Author: compiler-errors

compiler/rustc_middle/src/traits/solve.rs

@@ -279,8 +279,9 @@ pub enum GoalSource {
 
 /// Possible ways the given goal can be proven.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub enum CandidateSource {
-    /// A user written impl.
+pub enum CandidateSource<'tcx> {
+    /// A user written impl. Optionally, impl args are recorded if proof trees are
+    /// enabled, this came from a trait goal and the goal succeeded or was ambiguous.
     ///
     /// ## Examples
     ///
@@ -291,7 +292,7 @@ pub enum CandidateSource {
     ///     let y = x.clone();
     /// }
     /// ```
-    Impl(DefId),
+    Impl(DefId, Option<inspect::CanonicalState<'tcx, ty::GenericArgsRef<'tcx>>>),
     /// A builtin impl generated by the compiler. When adding a new special
     /// trait, try to use actual impls whenever possible. Builtin impls should
     /// only be used in cases where the impl cannot be manually be written.

compiler/rustc_middle/src/traits/solve/inspect.rs

@@ -143,7 +143,7 @@ pub enum ProbeKind<'tcx> {
     /// Probe entered when normalizing the self ty during candidate assembly
     NormalizedSelfTyAssembly,
     /// A candidate for proving a trait or alias-relate goal.
-    TraitCandidate { source: CandidateSource, result: QueryResult<'tcx> },
+    TraitCandidate { source: CandidateSource<'tcx>, result: QueryResult<'tcx> },
     /// Used in the probe that wraps normalizing the non-self type for the unsize
     /// trait, which is also structurally matched on.
     UnsizeAssembly,

compiler/rustc_trait_selection/src/solve/assembly/mod.rs

@@ -25,7 +25,7 @@ pub(super) mod structural_traits;
 /// and the `result` when using the given `source`.
 #[derive(Debug, Clone)]
 pub(super) struct Candidate<'tcx> {
-    pub(super) source: CandidateSource,
+    pub(super) source: CandidateSource<'tcx>,
     pub(super) result: CanonicalResponse<'tcx>,
 }
 
@@ -47,7 +47,7 @@ pub(super) trait GoalKind<'tcx>:
     /// [`EvalCtxt::evaluate_added_goals_and_make_canonical_response`]).
     fn probe_and_match_goal_against_assumption(
         ecx: &mut EvalCtxt<'_, 'tcx>,
-        source: CandidateSource,
+        source: CandidateSource<'tcx>,
         goal: Goal<'tcx, Self>,
         assumption: ty::Clause<'tcx>,
         then: impl FnOnce(&mut EvalCtxt<'_, 'tcx>) -> QueryResult<'tcx>,
@@ -58,7 +58,7 @@ pub(super) trait GoalKind<'tcx>:
     /// goal by equating it with the assumption.
     fn probe_and_consider_implied_clause(
         ecx: &mut EvalCtxt<'_, 'tcx>,
-        parent_source: CandidateSource,
+        parent_source: CandidateSource<'tcx>,
         goal: Goal<'tcx, Self>,
         assumption: ty::Clause<'tcx>,
         requirements: impl IntoIterator<Item = (GoalSource, Goal<'tcx, ty::Predicate<'tcx>>)>,
@@ -76,7 +76,7 @@ pub(super) trait GoalKind<'tcx>:
     /// since they're not implied by the well-formedness of the object type.
     fn probe_and_consider_object_bound_candidate(
         ecx: &mut EvalCtxt<'_, 'tcx>,
-        source: CandidateSource,
+        source: CandidateSource<'tcx>,
         goal: Goal<'tcx, Self>,
         assumption: ty::Clause<'tcx>,
     ) -> Result<Candidate<'tcx>, NoSolution> {
@@ -813,7 +813,7 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
                 // This feels dangerous.
                 Certainty::Yes => {
                     candidates.retain(|c| match c.source {
-                        CandidateSource::Impl(_) | CandidateSource::BuiltinImpl(_) => {
+                        CandidateSource::Impl(_, _) | CandidateSource::BuiltinImpl(_) => {
                             debug!(?c, "discard impl candidate");
                             false
                         }

compiler/rustc_trait_selection/src/solve/eval_ctxt/mod.rs

@@ -1090,4 +1090,11 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
             ControlFlow::Continue(())
         });
     }
+
+    pub(crate) fn canonicalize_impl_args(
+        &self,
+        impl_args: ty::GenericArgsRef<'tcx>,
+    ) -> Option<inspect::CanonicalState<'tcx, ty::GenericArgsRef<'tcx>>> {
+        self.inspect.canonicalize_impl_args(self.infcx, self.max_input_universe, impl_args)
+    }
 }

compiler/rustc_trait_selection/src/solve/eval_ctxt/probe.rs

@@ -51,7 +51,7 @@ where
 
 pub(in crate::solve) struct TraitProbeCtxt<'me, 'a, 'tcx, F> {
     cx: ProbeCtxt<'me, 'a, 'tcx, F, QueryResult<'tcx>>,
-    source: CandidateSource,
+    source: CandidateSource<'tcx>,
 }
 
 impl<'tcx, F> TraitProbeCtxt<'_, '_, 'tcx, F>
@@ -86,7 +86,7 @@ impl<'a, 'tcx> EvalCtxt<'a, 'tcx> {
 
     pub(in crate::solve) fn probe_trait_candidate(
         &mut self,
-        source: CandidateSource,
+        source: CandidateSource<'tcx>,
     ) -> TraitProbeCtxt<'_, 'a, 'tcx, impl FnOnce(&QueryResult<'tcx>) -> inspect::ProbeKind<'tcx>>
     {
         TraitProbeCtxt {

compiler/rustc_trait_selection/src/solve/eval_ctxt/select.rs

@@ -1,12 +1,11 @@
 use std::ops::ControlFlow;
 
-use rustc_hir::def_id::DefId;
-use rustc_infer::infer::{DefineOpaqueTypes, InferCtxt, InferOk};
+use rustc_infer::infer::InferCtxt;
 use rustc_infer::traits::solve::inspect::ProbeKind;
 use rustc_infer::traits::solve::{CandidateSource, Certainty, Goal};
 use rustc_infer::traits::{
     BuiltinImplSource, ImplSource, ImplSourceUserDefinedData, Obligation, ObligationCause,
-    PolyTraitObligation, PredicateObligation, Selection, SelectionError, SelectionResult,
+    PolyTraitObligation, Selection, SelectionError, SelectionResult,
 };
 use rustc_macros::extension;
 use rustc_span::Span;
@@ -112,12 +111,14 @@ fn candidate_should_be_dropped_in_favor_of<'tcx>(
             CandidateSource::BuiltinImpl(BuiltinImplSource::Object { .. }),
         ) => false,
         (
-            CandidateSource::Impl(_) | CandidateSource::ParamEnv(_) | CandidateSource::AliasBound,
+            CandidateSource::Impl(_, _)
+            | CandidateSource::ParamEnv(_)
+            | CandidateSource::AliasBound,
             CandidateSource::BuiltinImpl(BuiltinImplSource::Object { .. }),
         ) => true,
 
         // Prefer specializing candidates over specialized candidates.
-        (CandidateSource::Impl(victim_def_id), CandidateSource::Impl(other_def_id)) => {
+        (CandidateSource::Impl(victim_def_id, _), CandidateSource::Impl(other_def_id, _)) => {
             victim.goal().infcx().tcx.specializes((other_def_id, victim_def_id))
         }
 
@@ -133,32 +134,34 @@ fn to_selection<'tcx>(
         return None;
     }
 
-    let make_nested = || {
-        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()
-    };
+    let nested = 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 {
-            CandidateSource::Impl(impl_def_id) => {
-                // FIXME: Remove this in favor of storing this in the tree
-                // For impl candidates, we do the rematch manually to compute the args.
-                ImplSource::UserDefined(rematch_impl(cand.goal(), impl_def_id, span))
-            }
-            CandidateSource::BuiltinImpl(builtin) => ImplSource::Builtin(builtin, make_nested()),
-            CandidateSource::ParamEnv(_) => ImplSource::Param(make_nested()),
-            CandidateSource::AliasBound => {
-                ImplSource::Builtin(BuiltinImplSource::Misc, make_nested())
+            CandidateSource::Impl(impl_def_id, impl_args) => {
+                ImplSource::UserDefined(ImplSourceUserDefinedData {
+                    impl_def_id,
+                    args: cand.instantiate_impl_args(
+                        impl_args.expect("expected impl args to be recorded"),
+                        span,
+                    ),
+                    nested,
+                })
             }
+            CandidateSource::BuiltinImpl(builtin) => ImplSource::Builtin(builtin, nested),
+            CandidateSource::ParamEnv(_) => ImplSource::Param(nested),
+            CandidateSource::AliasBound => ImplSource::Builtin(BuiltinImplSource::Misc, nested),
             CandidateSource::CoherenceUnknowable => {
                 span_bug!(span, "didn't expect to select an unknowable candidate")
             }
@@ -173,40 +176,3 @@ fn to_selection<'tcx>(
         }
     })
 }
-
-fn rematch_impl<'tcx>(
-    goal: &inspect::InspectGoal<'_, 'tcx>,
-    impl_def_id: DefId,
-    span: Span,
-) -> ImplSourceUserDefinedData<'tcx, PredicateObligation<'tcx>> {
-    let infcx = goal.infcx();
-    let goal_trait_ref = infcx
-        .enter_forall_and_leak_universe(goal.goal().predicate.to_opt_poly_trait_pred().unwrap())
-        .trait_ref;
-
-    let args = infcx.fresh_args_for_item(span, impl_def_id);
-    let impl_trait_ref =
-        infcx.tcx.impl_trait_ref(impl_def_id).unwrap().instantiate(infcx.tcx, args);
-
-    let InferOk { value: (), obligations: mut nested } = infcx
-        .at(&ObligationCause::dummy_with_span(span), goal.goal().param_env)
-        .eq(DefineOpaqueTypes::Yes, goal_trait_ref, impl_trait_ref)
-        .expect("rematching impl failed");
-
-    // FIXME(-Znext-solver=coinductive): We need to add supertraits here eventually.
-
-    nested.extend(
-        infcx.tcx.predicates_of(impl_def_id).instantiate(infcx.tcx, args).into_iter().map(
-            |(clause, _)| {
-                Obligation::new(
-                    infcx.tcx,
-                    ObligationCause::dummy_with_span(span),
-                    goal.goal().param_env,
-                    clause,
-                )
-            },
-        ),
-    );
-
-    ImplSourceUserDefinedData { impl_def_id, nested, args }
-}

compiler/rustc_trait_selection/src/solve/fulfill.rs

@@ -324,8 +324,10 @@ impl<'tcx> ProofTreeVisitor<'tcx> for BestObligation<'tcx> {
         };
 
         // Don't walk into impls that have `do_not_recommend`.
-        if let ProbeKind::TraitCandidate { source: CandidateSource::Impl(impl_def_id), result: _ } =
-            candidate.kind()
+        if let ProbeKind::TraitCandidate {
+            source: CandidateSource::Impl(impl_def_id, _),
+            result: _,
+        } = candidate.kind()
             && goal.infcx().tcx.has_attr(impl_def_id, sym::do_not_recommend)
         {
             return ControlFlow::Break(self.obligation.clone());
@@ -388,7 +390,7 @@ fn derive_cause<'tcx>(
     parent_trait_pred: ty::PolyTraitPredicate<'tcx>,
 ) -> ObligationCause<'tcx> {
     match candidate_kind {
-        ProbeKind::TraitCandidate { source: CandidateSource::Impl(impl_def_id), result: _ } => {
+        ProbeKind::TraitCandidate { source: CandidateSource::Impl(impl_def_id, _), result: _ } => {
             if let Some((_, span)) =
                 tcx.predicates_of(impl_def_id).instantiate_identity(tcx).iter().nth(idx)
             {

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

@@ -216,6 +216,31 @@ impl<'a, 'tcx> InspectCandidate<'a, 'tcx> {
     pub fn visit_nested_in_probe<V: ProofTreeVisitor<'tcx>>(&self, visitor: &mut V) -> V::Result {
         self.goal.infcx.probe(|_| self.visit_nested_no_probe(visitor))
     }
+
+    pub(crate) fn instantiate_impl_args(
+        &self,
+        impl_args: inspect::CanonicalState<'tcx, ty::GenericArgsRef<'tcx>>,
+        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();
+        let args = canonical::instantiate_canonical_state(
+            infcx,
+            span,
+            param_env,
+            &mut orig_values,
+            impl_args,
+        );
+        let () = canonical::instantiate_canonical_state(
+            infcx,
+            span,
+            param_env,
+            &mut orig_values,
+            self.final_state,
+        );
+        infcx.resolve_vars_if_possible(args)
+    }
 }
 
 impl<'a, 'tcx> InspectGoal<'a, 'tcx> {

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

@@ -614,4 +614,23 @@ impl<'tcx> ProofTreeBuilder<'tcx> {
             }
         }
     }
+
+    /// Optionally canonicalize impl args. This will return `None` if proof trees are disabled.
+    pub(crate) fn canonicalize_impl_args(
+        &self,
+        infcx: &InferCtxt<'tcx>,
+        max_input_universe: ty::UniverseIndex,
+        impl_args: ty::GenericArgsRef<'tcx>,
+    ) -> Option<inspect::CanonicalState<'tcx, ty::GenericArgsRef<'tcx>>> {
+        match self.state.as_deref() {
+            None => None,
+            Some(DebugSolver::GoalEvaluationStep(state)) => Some(canonical::make_canonical_state(
+                infcx,
+                &state.var_values,
+                max_input_universe,
+                impl_args,
+            )),
+            _ => bug!(),
+        }
+    }
 }

compiler/rustc_trait_selection/src/solve/normalizes_to/mod.rs

@@ -118,7 +118,7 @@ impl<'tcx> assembly::GoalKind<'tcx> for NormalizesTo<'tcx> {
 
     fn probe_and_match_goal_against_assumption(
         ecx: &mut EvalCtxt<'_, 'tcx>,
-        source: CandidateSource,
+        source: CandidateSource<'tcx>,
         goal: Goal<'tcx, Self>,
         assumption: ty::Clause<'tcx>,
         then: impl FnOnce(&mut EvalCtxt<'_, 'tcx>) -> QueryResult<'tcx>,
@@ -182,7 +182,7 @@ impl<'tcx> assembly::GoalKind<'tcx> for NormalizesTo<'tcx> {
             ty::ImplPolarity::Positive => {}
         };
 
-        ecx.probe_trait_candidate(CandidateSource::Impl(impl_def_id)).enter(|ecx| {
+        ecx.probe_trait_candidate(CandidateSource::Impl(impl_def_id, None)).enter(|ecx| {
             let impl_args = ecx.fresh_args_for_item(impl_def_id);
             let impl_trait_ref = impl_trait_header.trait_ref.instantiate(tcx, impl_args);