From 6e3808e27407c967762f56a2ddffe2602a3079a3 Mon Sep 17 00:00:00 2001
From: Michael Goulet <michael@errs.io>
Date: Thu, 25 Apr 2024 18:41:08 -0400
Subject: [PATCH] Store goal source in InspectGoal

---
 .../src/solve/inspect/analyse.rs              | 39 +++++++++++++------
 1 file changed, 27 insertions(+), 12 deletions(-)

diff --git a/compiler/rustc_trait_selection/src/solve/inspect/analyse.rs b/compiler/rustc_trait_selection/src/solve/inspect/analyse.rs
index 7c6dd4d3feb..de6c2322e8d 100644
--- a/compiler/rustc_trait_selection/src/solve/inspect/analyse.rs
+++ b/compiler/rustc_trait_selection/src/solve/inspect/analyse.rs
@@ -42,6 +42,7 @@ pub struct InspectGoal<'a, 'tcx> {
     result: Result<Certainty, NoSolution>,
     evaluation_kind: inspect::CanonicalGoalEvaluationKind<'tcx>,
     normalizes_to_term_hack: Option<NormalizesToTermHack<'tcx>>,
+    source: GoalSource,
 }
 
 /// The expected term of a `NormalizesTo` goal gets replaced
@@ -92,7 +93,7 @@ impl<'tcx> NormalizesToTermHack<'tcx> {
 pub struct InspectCandidate<'a, 'tcx> {
     goal: &'a InspectGoal<'a, 'tcx>,
     kind: inspect::ProbeKind<'tcx>,
-    nested_goals: Vec<inspect::CanonicalState<'tcx, Goal<'tcx, ty::Predicate<'tcx>>>>,
+    nested_goals: Vec<(GoalSource, inspect::CanonicalState<'tcx, Goal<'tcx, ty::Predicate<'tcx>>>)>,
     final_state: inspect::CanonicalState<'tcx, ()>,
     result: QueryResult<'tcx>,
     shallow_certainty: Certainty,
@@ -145,13 +146,16 @@ impl<'a, 'tcx> InspectCandidate<'a, 'tcx> {
         let instantiated_goals: Vec<_> = self
             .nested_goals
             .iter()
-            .map(|goal| {
-                canonical::instantiate_canonical_state(
-                    infcx,
-                    span,
-                    param_env,
-                    &mut orig_values,
-                    *goal,
+            .map(|(source, goal)| {
+                (
+                    *source,
+                    canonical::instantiate_canonical_state(
+                        infcx,
+                        span,
+                        param_env,
+                        &mut orig_values,
+                        *goal,
+                    ),
                 )
             })
             .collect();
@@ -173,7 +177,7 @@ impl<'a, 'tcx> InspectCandidate<'a, 'tcx> {
 
         instantiated_goals
             .into_iter()
-            .map(|goal| match goal.predicate.kind().no_bound_vars() {
+            .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
@@ -197,6 +201,7 @@ impl<'a, 'tcx> InspectCandidate<'a, 'tcx> {
                         self.goal.depth + 1,
                         proof_tree.unwrap(),
                         Some(NormalizesToTermHack { term, unconstrained_term }),
+                        source,
                     )
                 }
                 _ => InspectGoal::new(
@@ -204,6 +209,7 @@ impl<'a, 'tcx> InspectCandidate<'a, 'tcx> {
                     self.goal.depth + 1,
                     infcx.evaluate_root_goal(goal, GenerateProofTree::Yes).1.unwrap(),
                     None,
+                    source,
                 ),
             })
             .collect()
@@ -229,16 +235,23 @@ impl<'a, 'tcx> InspectGoal<'a, 'tcx> {
         self.result
     }
 
+    pub fn source(&self) -> GoalSource {
+        self.source
+    }
+
     fn candidates_recur(
         &'a self,
         candidates: &mut Vec<InspectCandidate<'a, 'tcx>>,
-        nested_goals: &mut Vec<inspect::CanonicalState<'tcx, Goal<'tcx, ty::Predicate<'tcx>>>>,
+        nested_goals: &mut Vec<(
+            GoalSource,
+            inspect::CanonicalState<'tcx, Goal<'tcx, ty::Predicate<'tcx>>>,
+        )>,
         probe: &inspect::Probe<'tcx>,
     ) {
         let mut shallow_certainty = None;
         for step in &probe.steps {
             match step {
-                &inspect::ProbeStep::AddGoal(_source, goal) => nested_goals.push(goal),
+                &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
@@ -321,6 +334,7 @@ impl<'a, 'tcx> InspectGoal<'a, 'tcx> {
         depth: usize,
         root: inspect::GoalEvaluation<'tcx>,
         normalizes_to_term_hack: Option<NormalizesToTermHack<'tcx>>,
+        source: GoalSource,
     ) -> Self {
         let inspect::GoalEvaluation { uncanonicalized_goal, kind, evaluation } = root;
         let inspect::GoalEvaluationKind::Root { orig_values } = kind else { unreachable!() };
@@ -343,6 +357,7 @@ impl<'a, 'tcx> InspectGoal<'a, 'tcx> {
             result,
             evaluation_kind: evaluation.kind,
             normalizes_to_term_hack,
+            source,
         }
     }
 }
@@ -369,6 +384,6 @@ impl<'tcx> InferCtxt<'tcx> {
     ) -> V::Result {
         let (_, proof_tree) = self.evaluate_root_goal(goal, GenerateProofTree::Yes);
         let proof_tree = proof_tree.unwrap();
-        visitor.visit_goal(&InspectGoal::new(self, 0, proof_tree, None))
+        visitor.visit_goal(&InspectGoal::new(self, 0, proof_tree, None, GoalSource::Misc))
     }
 }