From 1b141b6d73896e168abf3df8046a940c8928880e Mon Sep 17 00:00:00 2001
From: lcnr <rust@lcnr.de>
Date: Thu, 14 Sep 2023 10:41:36 +0200
Subject: [PATCH] inspect: explicitly store added goals

---
 .../rustc_middle/src/traits/solve/inspect.rs    |  1 +
 .../src/traits/solve/inspect/format.rs          |  1 +
 .../src/solve/eval_ctxt.rs                      |  2 +-
 .../rustc_trait_selection/src/solve/inspect.rs  | 17 +++++++++++++++++
 compiler/rustc_trait_selection/src/solve/mod.rs |  1 +
 5 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/compiler/rustc_middle/src/traits/solve/inspect.rs b/compiler/rustc_middle/src/traits/solve/inspect.rs
index a014e8769d8..9698d2e4790 100644
--- a/compiler/rustc_middle/src/traits/solve/inspect.rs
+++ b/compiler/rustc_middle/src/traits/solve/inspect.rs
@@ -72,6 +72,7 @@ impl Debug for Probe<'_> {
 
 #[derive(Eq, PartialEq)]
 pub enum ProbeStep<'tcx> {
+    AddGoal(Goal<'tcx, ty::Predicate<'tcx>>),
     EvaluateGoals(AddedGoalsEvaluation<'tcx>),
     NestedProbe(Probe<'tcx>),
 }
diff --git a/compiler/rustc_middle/src/traits/solve/inspect/format.rs b/compiler/rustc_middle/src/traits/solve/inspect/format.rs
index faaf791bf89..3741b933a6b 100644
--- a/compiler/rustc_middle/src/traits/solve/inspect/format.rs
+++ b/compiler/rustc_middle/src/traits/solve/inspect/format.rs
@@ -120,6 +120,7 @@ impl<'a, 'b> ProofTreeFormatter<'a, 'b> {
         self.nested(|this| {
             for step in &probe.steps {
                 match step {
+                    ProbeStep::AddGoal(goal) => writeln!(this.f, "ADDED GOAL: {goal:?}")?,
                     ProbeStep::EvaluateGoals(eval) => this.format_added_goals_evaluation(eval)?,
                     ProbeStep::NestedProbe(probe) => this.format_probe(probe)?,
                 }
diff --git a/compiler/rustc_trait_selection/src/solve/eval_ctxt.rs b/compiler/rustc_trait_selection/src/solve/eval_ctxt.rs
index 307c0516f70..8079e40698a 100644
--- a/compiler/rustc_trait_selection/src/solve/eval_ctxt.rs
+++ b/compiler/rustc_trait_selection/src/solve/eval_ctxt.rs
@@ -85,7 +85,7 @@ pub struct EvalCtxt<'a, 'tcx> {
     // evaluation code.
     tainted: Result<(), NoSolution>,
 
-    inspect: ProofTreeBuilder<'tcx>,
+    pub(super) inspect: ProofTreeBuilder<'tcx>,
 }
 
 #[derive(Debug, Clone)]
diff --git a/compiler/rustc_trait_selection/src/solve/inspect.rs b/compiler/rustc_trait_selection/src/solve/inspect.rs
index ca922b64e28..86608dae2be 100644
--- a/compiler/rustc_trait_selection/src/solve/inspect.rs
+++ b/compiler/rustc_trait_selection/src/solve/inspect.rs
@@ -118,6 +118,7 @@ impl<'tcx> WipProbe<'tcx> {
 
 #[derive(Eq, PartialEq, Debug)]
 pub enum WipProbeStep<'tcx> {
+    AddGoal(Goal<'tcx, ty::Predicate<'tcx>>),
     EvaluateGoals(WipAddedGoalsEvaluation<'tcx>),
     NestedProbe(WipProbe<'tcx>),
 }
@@ -125,6 +126,7 @@ pub enum WipProbeStep<'tcx> {
 impl<'tcx> WipProbeStep<'tcx> {
     pub fn finalize(self) -> inspect::ProbeStep<'tcx> {
         match self {
+            WipProbeStep::AddGoal(goal) => inspect::ProbeStep::AddGoal(goal),
             WipProbeStep::EvaluateGoals(eval) => inspect::ProbeStep::EvaluateGoals(eval.finalize()),
             WipProbeStep::NestedProbe(probe) => inspect::ProbeStep::NestedProbe(probe.finalize()),
         }
@@ -370,6 +372,21 @@ impl<'tcx> ProofTreeBuilder<'tcx> {
         }
     }
 
+    pub fn add_goal(&mut self, goal: Goal<'tcx, ty::Predicate<'tcx>>) {
+        if let Some(this) = self.as_mut() {
+            match this {
+                DebugSolver::GoalEvaluationStep(WipGoalEvaluationStep {
+                    evaluation: WipProbe { steps, .. },
+                    ..
+                })
+                | DebugSolver::Probe(WipProbe { steps, .. }) => {
+                    steps.push(WipProbeStep::AddGoal(goal))
+                }
+                _ => unreachable!(),
+            }
+        }
+    }
+
     pub fn finish_probe(&mut self, probe: ProofTreeBuilder<'tcx>) {
         if let Some(this) = self.as_mut() {
             match (this, probe.state.unwrap().tree) {
diff --git a/compiler/rustc_trait_selection/src/solve/mod.rs b/compiler/rustc_trait_selection/src/solve/mod.rs
index c492408bc76..558e23a5d36 100644
--- a/compiler/rustc_trait_selection/src/solve/mod.rs
+++ b/compiler/rustc_trait_selection/src/solve/mod.rs
@@ -228,6 +228,7 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
 
     #[instrument(level = "debug", skip(self))]
     fn add_goal(&mut self, goal: Goal<'tcx, ty::Predicate<'tcx>>) {
+        self.inspect.add_goal(goal);
         self.nested_goals.goals.push(goal);
     }