From 998644f43ab2636836d5726babdba8a28142da53 Mon Sep 17 00:00:00 2001
From: gavinleroy <gavinleroy6@gmail.com>
Date: Tue, 16 Jan 2024 12:31:30 +0100
Subject: [PATCH] Store inference context alongside captured obligations.

---
 compiler/rustc_infer/src/infer/mod.rs         |  1 +
 compiler/rustc_infer/src/traits/mod.rs        |  5 ++--
 .../src/solve/fulfill.rs                      | 25 +++++++++++++------
 3 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/compiler/rustc_infer/src/infer/mod.rs b/compiler/rustc_infer/src/infer/mod.rs
index ba5d7fccb405c..895cd6be1d3f2 100644
--- a/compiler/rustc_infer/src/infer/mod.rs
+++ b/compiler/rustc_infer/src/infer/mod.rs
@@ -236,6 +236,7 @@ impl<'tcx> InferCtxtInner<'tcx> {
     }
 }
 
+#[derive(Clone)]
 pub struct InferCtxt<'tcx> {
     pub tcx: TyCtxt<'tcx>,
 
diff --git a/compiler/rustc_infer/src/traits/mod.rs b/compiler/rustc_infer/src/traits/mod.rs
index c321ef1fbf966..8cb8eee345b27 100644
--- a/compiler/rustc_infer/src/traits/mod.rs
+++ b/compiler/rustc_infer/src/traits/mod.rs
@@ -20,6 +20,7 @@ use rustc_span::Span;
 pub use self::FulfillmentErrorCode::*;
 pub use self::ImplSource::*;
 pub use self::SelectionError::*;
+use crate::infer::InferCtxt;
 
 pub use self::engine::{TraitEngine, TraitEngineExt};
 pub use self::project::MismatchedProjectionTypes;
@@ -110,8 +111,8 @@ impl<'tcx> PolyTraitObligation<'tcx> {
 }
 
 pub enum FulfilledObligation<'tcx> {
-    Success(PredicateObligation<'tcx>),
-    Failure(FulfillmentError<'tcx>),
+    Success { infcx: InferCtxt<'tcx>, data: PredicateObligation<'tcx> },
+    Failed { infcx: InferCtxt<'tcx>, data: Vec<FulfillmentError<'tcx>> },
 }
 
 // `PredicateObligation` is used a lot. Make sure it doesn't unintentionally get bigger.
diff --git a/compiler/rustc_trait_selection/src/solve/fulfill.rs b/compiler/rustc_trait_selection/src/solve/fulfill.rs
index c7514891ea681..7a3fd74605cc1 100644
--- a/compiler/rustc_trait_selection/src/solve/fulfill.rs
+++ b/compiler/rustc_trait_selection/src/solve/fulfill.rs
@@ -54,18 +54,27 @@ impl<'tcx> FulfillmentCtxt<'tcx> {
 
     fn track_fulfillment_errors<'b, 'a: 'b>(
         &'a self,
+        infcx: &InferCtxt<'tcx>,
         errors: impl IntoIterator<Item = &'b FulfillmentError<'tcx>>,
     ) {
         if let Some(tracked_obligations) = &self.tracked_obligations {
-            tracked_obligations.borrow_mut().extend(
-                errors.into_iter().map(|error| FulfilledObligation::Failure(error.clone())),
-            );
+            tracked_obligations.borrow_mut().push(FulfilledObligation::Failed {
+                infcx: infcx.clone(),
+                data: errors.into_iter().map(Clone::clone).collect::<Vec<_>>(),
+            });
         }
     }
 
-    fn track_fulfillment_success(&self, predicate: &PredicateObligation<'tcx>) {
+    fn track_fulfillment_success(
+        &self,
+        infcx: &InferCtxt<'tcx>,
+        predicate: &PredicateObligation<'tcx>,
+    ) {
         if let Some(tracked_obligations) = &self.tracked_obligations {
-            tracked_obligations.borrow_mut().push(FulfilledObligation::Success(predicate.clone()));
+            tracked_obligations.borrow_mut().push(FulfilledObligation::Success {
+                infcx: infcx.clone(),
+                data: predicate.clone(),
+            });
         }
     }
 }
@@ -114,7 +123,7 @@ impl<'tcx> TraitEngine<'tcx> for FulfillmentCtxt<'tcx> {
             })
             .collect();
 
-        self.track_fulfillment_errors(&errors);
+        self.track_fulfillment_errors(infcx, &errors);
 
         errors
     }
@@ -133,7 +142,7 @@ impl<'tcx> TraitEngine<'tcx> for FulfillmentCtxt<'tcx> {
                 let (changed, certainty, nested_goals) =
                     match infcx.evaluate_root_goal(goal, GenerateProofTree::IfEnabled).0 {
                         Ok(result) => {
-                            self.track_fulfillment_success(&obligation);
+                            self.track_fulfillment_success(infcx, &obligation);
                             result
                         }
                         Err(NoSolution) => {
@@ -214,7 +223,7 @@ impl<'tcx> TraitEngine<'tcx> for FulfillmentCtxt<'tcx> {
             }
         }
 
-        self.track_fulfillment_errors(&errors);
+        self.track_fulfillment_errors(infcx, &errors);
 
         errors
     }