Manually walk into WF obligations in BestObligation proof tree visitor

5 files changed, 105 insertions, 62 deletions

diff --git a/compiler/rustc_trait_selection/src/solve/delegate.rs b/compiler/rustc_trait_selection/src/solve/delegate.rs
index acd00d9f74f..abb79493432 100644
--- a/compiler/rustc_trait_selection/src/solve/delegate.rs
+++ b/compiler/rustc_trait_selection/src/solve/delegate.rs
@@ -1,7 +1,7 @@
 use std::ops::Deref;
 
 use rustc_data_structures::fx::FxHashSet;
-use rustc_hir::def_id::DefId;
+use rustc_hir::def_id::{CRATE_DEF_ID, DefId};
 use rustc_infer::infer::canonical::query_response::make_query_region_constraints;
 use rustc_infer::infer::canonical::{
     Canonical, CanonicalExt as _, CanonicalQueryInput, CanonicalVarInfo, CanonicalVarValues,
@@ -98,9 +98,10 @@ impl<'tcx> rustc_next_trait_solver::delegate::SolverDelegate for SolverDelegate<
         param_env: ty::ParamEnv<'tcx>,
         arg: ty::GenericArg<'tcx>,
     ) -> Option<Vec<Goal<'tcx, ty::Predicate<'tcx>>>> {
-        crate::traits::wf::unnormalized_obligations(&self.0, param_env, arg).map(|obligations| {
-            obligations.into_iter().map(|obligation| obligation.into()).collect()
-        })
+        crate::traits::wf::unnormalized_obligations(&self.0, param_env, arg, DUMMY_SP, CRATE_DEF_ID)
+            .map(|obligations| {
+                obligations.into_iter().map(|obligation| obligation.into()).collect()
+            })
     }
 
     fn clone_opaque_types_for_query_response(&self) -> Vec<(ty::OpaqueTypeKey<'tcx>, Ty<'tcx>)> {
diff --git a/compiler/rustc_trait_selection/src/solve/fulfill/derive_errors.rs b/compiler/rustc_trait_selection/src/solve/fulfill/derive_errors.rs
index 2c55f13a0df..c64bc19835b 100644
--- a/compiler/rustc_trait_selection/src/solve/fulfill/derive_errors.rs
+++ b/compiler/rustc_trait_selection/src/solve/fulfill/derive_errors.rs
@@ -10,13 +10,13 @@ use rustc_middle::ty::error::{ExpectedFound, TypeError};
 use rustc_middle::ty::{self, TyCtxt};
 use rustc_middle::{bug, span_bug};
 use rustc_next_trait_solver::solve::{GenerateProofTree, SolverDelegateEvalExt as _};
-use rustc_type_ir::solve::NoSolution;
+use rustc_type_ir::solve::{Goal, NoSolution};
 use tracing::{instrument, trace};
 
 use crate::solve::Certainty;
 use crate::solve::delegate::SolverDelegate;
 use crate::solve::inspect::{self, ProofTreeInferCtxtExt, ProofTreeVisitor};
-use crate::traits::{FulfillmentError, FulfillmentErrorCode};
+use crate::traits::{FulfillmentError, FulfillmentErrorCode, wf};
 
 pub(super) fn fulfillment_error_for_no_solution<'tcx>(
     infcx: &InferCtxt<'tcx>,
@@ -207,14 +207,10 @@ impl<'tcx> BestObligation<'tcx> {
                                             | GoalSource::AliasBoundConstCondition
                                             | GoalSource::InstantiateHigherRanked
                                             | GoalSource::AliasWellFormed
-                                    ) && match self.consider_ambiguities {
-                                        true => {
-                                            matches!(
-                                                nested_goal.result(),
-                                                Ok(Certainty::Maybe(MaybeCause::Ambiguity))
-                                            )
-                                        }
-                                        false => matches!(nested_goal.result(), Err(_)),
+                                    ) && match (self.consider_ambiguities, nested_goal.result()) {
+                                        (true, Ok(Certainty::Maybe(MaybeCause::Ambiguity)))
+                                        | (false, Err(_)) => true,
+                                        _ => false,
                                     }
                                 },
                             )
@@ -233,6 +229,39 @@ impl<'tcx> BestObligation<'tcx> {
 
         candidates
     }
+
+    /// HACK: We walk the nested obligations for a well-formed arg manually,
+    /// since there's nontrivial logic in `wf.rs` to set up an obligation cause.
+    /// Ideally we'd be able to track this better.
+    fn visit_well_formed_goal(
+        &mut self,
+        candidate: &inspect::InspectCandidate<'_, 'tcx>,
+        arg: ty::GenericArg<'tcx>,
+    ) -> ControlFlow<PredicateObligation<'tcx>> {
+        let infcx = candidate.goal().infcx();
+        let param_env = candidate.goal().goal().param_env;
+        let body_id = self.obligation.cause.body_id;
+
+        for obligation in wf::unnormalized_obligations(infcx, param_env, arg, self.span(), body_id)
+            .into_iter()
+            .flatten()
+        {
+            let nested_goal = candidate.instantiate_proof_tree_for_nested_goal(
+                GoalSource::Misc,
+                Goal::new(infcx.tcx, obligation.param_env, obligation.predicate),
+                self.span(),
+            );
+            // Skip nested goals that aren't the *reason* for our goal's failure.
+            match (self.consider_ambiguities, nested_goal.result()) {
+                (true, Ok(Certainty::Maybe(MaybeCause::Ambiguity))) | (false, Err(_)) => {}
+                _ => continue,
+            }
+
+            self.with_derived_obligation(obligation, |this| nested_goal.visit_with(this))?;
+        }
+
+        ControlFlow::Break(self.obligation.clone())
+    }
 }
 
 impl<'tcx> ProofTreeVisitor<'tcx> for BestObligation<'tcx> {
@@ -283,6 +312,9 @@ impl<'tcx> ProofTreeVisitor<'tcx> for BestObligation<'tcx> {
                     polarity: ty::PredicatePolarity::Positive,
                 }))
             }
+            ty::PredicateKind::Clause(ty::ClauseKind::WellFormed(arg)) => {
+                return self.visit_well_formed_goal(candidate, arg);
+            }
             _ => ChildMode::PassThrough,
         };
 
@@ -355,12 +387,8 @@ impl<'tcx> ProofTreeVisitor<'tcx> for BestObligation<'tcx> {
             }
 
             // Skip nested goals that aren't the *reason* for our goal's failure.
-            match self.consider_ambiguities {
-                true if matches!(
-                    nested_goal.result(),
-                    Ok(Certainty::Maybe(MaybeCause::Ambiguity))
-                ) => {}
-                false if matches!(nested_goal.result(), Err(_)) => {}
+            match (self.consider_ambiguities, nested_goal.result()) {
+                (true, Ok(Certainty::Maybe(MaybeCause::Ambiguity))) | (false, Err(_)) => {}
                 _ => continue,
             }
 
diff --git a/compiler/rustc_trait_selection/src/solve/inspect/analyse.rs b/compiler/rustc_trait_selection/src/solve/inspect/analyse.rs
index e735020a63e..9ba48cd588f 100644
--- a/compiler/rustc_trait_selection/src/solve/inspect/analyse.rs
+++ b/compiler/rustc_trait_selection/src/solve/inspect/analyse.rs
@@ -194,47 +194,57 @@ impl<'a, 'tcx> InspectCandidate<'a, 'tcx> {
 
         let goals = instantiated_goals
             .into_iter()
-            .map(|(source, goal)| match goal.predicate.kind().no_bound_vars() {
-                Some(ty::PredicateKind::NormalizesTo(ty::NormalizesTo { alias, term })) => {
-                    let unconstrained_term = match term.unpack() {
-                        ty::TermKind::Ty(_) => infcx.next_ty_var(span).into(),
-                        ty::TermKind::Const(_) => infcx.next_const_var(span).into(),
-                    };
-                    let goal =
-                        goal.with(infcx.tcx, ty::NormalizesTo { alias, term: unconstrained_term });
-                    // We have to use a `probe` here as evaluating a `NormalizesTo` can constrain the
-                    // expected term. This means that candidates which only fail due to nested goals
-                    // and which normalize to a different term then the final result could ICE: when
-                    // building their proof tree, the expected term was unconstrained, but when
-                    // instantiating the candidate it is already constrained to the result of another
-                    // candidate.
-                    let proof_tree = infcx
-                        .probe(|_| infcx.evaluate_root_goal_raw(goal, GenerateProofTree::Yes).1);
-                    InspectGoal::new(
-                        infcx,
-                        self.goal.depth + 1,
-                        proof_tree.unwrap(),
-                        Some(NormalizesToTermHack { term, unconstrained_term }),
-                        source,
-                    )
-                }
-                _ => {
-                    // We're using a probe here as evaluating a goal could constrain
-                    // inference variables by choosing one candidate. If we then recurse
-                    // into another candidate who ends up with different inference
-                    // constraints, we get an ICE if we already applied the constraints
-                    // from the chosen candidate.
-                    let proof_tree = infcx
-                        .probe(|_| infcx.evaluate_root_goal(goal, GenerateProofTree::Yes).1)
-                        .unwrap();
-                    InspectGoal::new(infcx, self.goal.depth + 1, proof_tree, None, source)
-                }
-            })
+            .map(|(source, goal)| self.instantiate_proof_tree_for_nested_goal(source, goal, span))
             .collect();
 
         (goals, opt_impl_args)
     }
 
+    pub fn instantiate_proof_tree_for_nested_goal(
+        &self,
+        source: GoalSource,
+        goal: Goal<'tcx, ty::Predicate<'tcx>>,
+        span: Span,
+    ) -> InspectGoal<'a, 'tcx> {
+        let infcx = self.goal.infcx;
+        match goal.predicate.kind().no_bound_vars() {
+            Some(ty::PredicateKind::NormalizesTo(ty::NormalizesTo { alias, term })) => {
+                let unconstrained_term = match term.unpack() {
+                    ty::TermKind::Ty(_) => infcx.next_ty_var(span).into(),
+                    ty::TermKind::Const(_) => infcx.next_const_var(span).into(),
+                };
+                let goal =
+                    goal.with(infcx.tcx, ty::NormalizesTo { alias, term: unconstrained_term });
+                // We have to use a `probe` here as evaluating a `NormalizesTo` can constrain the
+                // expected term. This means that candidates which only fail due to nested goals
+                // and which normalize to a different term then the final result could ICE: when
+                // building their proof tree, the expected term was unconstrained, but when
+                // instantiating the candidate it is already constrained to the result of another
+                // candidate.
+                let proof_tree =
+                    infcx.probe(|_| infcx.evaluate_root_goal_raw(goal, GenerateProofTree::Yes).1);
+                InspectGoal::new(
+                    infcx,
+                    self.goal.depth + 1,
+                    proof_tree.unwrap(),
+                    Some(NormalizesToTermHack { term, unconstrained_term }),
+                    source,
+                )
+            }
+            _ => {
+                // We're using a probe here as evaluating a goal could constrain
+                // inference variables by choosing one candidate. If we then recurse
+                // into another candidate who ends up with different inference
+                // constraints, we get an ICE if we already applied the constraints
+                // from the chosen candidate.
+                let proof_tree = infcx
+                    .probe(|_| infcx.evaluate_root_goal(goal, GenerateProofTree::Yes).1)
+                    .unwrap();
+                InspectGoal::new(infcx, self.goal.depth + 1, proof_tree, None, source)
+            }
+        }
+    }
+
     /// Visit all nested goals of this candidate, rolling back
     /// all inference constraints.
     pub fn visit_nested_in_probe<V: ProofTreeVisitor<'tcx>>(&self, visitor: &mut V) -> V::Result {
diff --git a/compiler/rustc_trait_selection/src/traits/query/type_op/implied_outlives_bounds.rs b/compiler/rustc_trait_selection/src/traits/query/type_op/implied_outlives_bounds.rs
index 1339739ce7f..ec0b7903396 100644
--- a/compiler/rustc_trait_selection/src/traits/query/type_op/implied_outlives_bounds.rs
+++ b/compiler/rustc_trait_selection/src/traits/query/type_op/implied_outlives_bounds.rs
@@ -5,8 +5,8 @@ use rustc_infer::traits::query::type_op::ImpliedOutlivesBounds;
 use rustc_middle::infer::canonical::CanonicalQueryResponse;
 use rustc_middle::traits::ObligationCause;
 use rustc_middle::ty::{self, ParamEnvAnd, Ty, TyCtxt, TypeFolder, TypeVisitableExt};
-use rustc_span::Span;
 use rustc_span::def_id::CRATE_DEF_ID;
+use rustc_span::{DUMMY_SP, Span};
 use rustc_type_ir::outlives::{Component, push_outlives_components};
 use smallvec::{SmallVec, smallvec};
 use tracing::debug;
@@ -92,7 +92,9 @@ pub fn compute_implied_outlives_bounds_inner<'tcx>(
 
         // From the full set of obligations, just filter down to the region relationships.
         for obligation in
-            wf::unnormalized_obligations(ocx.infcx, param_env, arg).into_iter().flatten()
+            wf::unnormalized_obligations(ocx.infcx, param_env, arg, DUMMY_SP, CRATE_DEF_ID)
+                .into_iter()
+                .flatten()
         {
             assert!(!obligation.has_escaping_bound_vars());
             let Some(pred) = obligation.predicate.kind().no_bound_vars() else {
diff --git a/compiler/rustc_trait_selection/src/traits/wf.rs b/compiler/rustc_trait_selection/src/traits/wf.rs
index 20b675bcb76..37d49cc2f55 100644
--- a/compiler/rustc_trait_selection/src/traits/wf.rs
+++ b/compiler/rustc_trait_selection/src/traits/wf.rs
@@ -8,8 +8,8 @@ use rustc_middle::ty::{
     self, GenericArg, GenericArgKind, GenericArgsRef, Ty, TyCtxt, TypeSuperVisitable,
     TypeVisitable, TypeVisitableExt, TypeVisitor,
 };
-use rustc_span::def_id::{CRATE_DEF_ID, DefId, LocalDefId};
-use rustc_span::{DUMMY_SP, Span};
+use rustc_span::Span;
+use rustc_span::def_id::{DefId, LocalDefId};
 use tracing::{debug, instrument, trace};
 
 use crate::infer::InferCtxt;
@@ -89,6 +89,8 @@ pub fn unnormalized_obligations<'tcx>(
     infcx: &InferCtxt<'tcx>,
     param_env: ty::ParamEnv<'tcx>,
     arg: GenericArg<'tcx>,
+    span: Span,
+    body_id: LocalDefId,
 ) -> Option<PredicateObligations<'tcx>> {
     debug_assert_eq!(arg, infcx.resolve_vars_if_possible(arg));
 
@@ -106,8 +108,8 @@ pub fn unnormalized_obligations<'tcx>(
     let mut wf = WfPredicates {
         infcx,
         param_env,
-        body_id: CRATE_DEF_ID,
-        span: DUMMY_SP,
+        body_id,
+        span,
         out: PredicateObligations::new(),
         recursion_depth: 0,
         item: None,