(almost) Always use ObligationCtxt when dealing with canonical queries

2 files changed, 29 insertions, 7 deletions

diff --git a/compiler/rustc_trait_selection/src/infer.rs b/compiler/rustc_trait_selection/src/infer.rs
index 1b58c9b864e..3763c871e48 100644
--- a/compiler/rustc_trait_selection/src/infer.rs
+++ b/compiler/rustc_trait_selection/src/infer.rs
@@ -1,5 +1,5 @@
 use crate::traits::query::evaluate_obligation::InferCtxtExt as _;
-use crate::traits::{self, TraitEngine, TraitEngineExt};
+use crate::traits::{self, ObligationCtxt};
 
 use rustc_hir::def_id::DefId;
 use rustc_hir::lang_items::LangItem;
@@ -142,7 +142,7 @@ pub trait InferCtxtBuilderExt<'tcx> {
     fn enter_canonical_trait_query<K, R>(
         &mut self,
         canonical_key: &Canonical<'tcx, K>,
-        operation: impl FnOnce(&InferCtxt<'tcx>, &mut dyn TraitEngine<'tcx>, K) -> Fallible<R>,
+        operation: impl FnOnce(&ObligationCtxt<'_, 'tcx>, K) -> Fallible<R>,
     ) -> Fallible<CanonicalizedQueryResponse<'tcx, R>>
     where
         K: TypeFoldable<'tcx>,
@@ -170,17 +170,17 @@ impl<'tcx> InferCtxtBuilderExt<'tcx> for InferCtxtBuilder<'tcx> {
     fn enter_canonical_trait_query<K, R>(
         &mut self,
         canonical_key: &Canonical<'tcx, K>,
-        operation: impl FnOnce(&InferCtxt<'tcx>, &mut dyn TraitEngine<'tcx>, K) -> Fallible<R>,
+        operation: impl FnOnce(&ObligationCtxt<'_, 'tcx>, K) -> Fallible<R>,
     ) -> Fallible<CanonicalizedQueryResponse<'tcx, R>>
     where
         K: TypeFoldable<'tcx>,
         R: Debug + TypeFoldable<'tcx>,
         Canonical<'tcx, QueryResponse<'tcx, R>>: ArenaAllocatable<'tcx>,
     {
-        let (ref infcx, key, canonical_inference_vars) =
+        let (infcx, key, canonical_inference_vars) =
             self.build_with_canonical(DUMMY_SP, canonical_key);
-        let mut fulfill_cx = <dyn TraitEngine<'_>>::new(infcx.tcx);
-        let value = operation(infcx, &mut *fulfill_cx, key)?;
-        infcx.make_canonicalized_query_response(canonical_inference_vars, value, &mut *fulfill_cx)
+        let ocx = ObligationCtxt::new(&infcx);
+        let value = operation(&ocx, key)?;
+        ocx.make_canonicalized_query_response(canonical_inference_vars, value)
     }
 }
diff --git a/compiler/rustc_trait_selection/src/traits/engine.rs b/compiler/rustc_trait_selection/src/traits/engine.rs
index e0c8deec91a..18ad99b4935 100644
--- a/compiler/rustc_trait_selection/src/traits/engine.rs
+++ b/compiler/rustc_trait_selection/src/traits/engine.rs
@@ -1,14 +1,20 @@
 use std::cell::RefCell;
+use std::fmt::Debug;
 
 use super::TraitEngine;
 use super::{ChalkFulfillmentContext, FulfillmentContext};
 use crate::infer::InferCtxtExt;
 use rustc_data_structures::fx::FxHashSet;
 use rustc_hir::def_id::{DefId, LocalDefId};
+use rustc_infer::infer::canonical::{
+    Canonical, CanonicalVarValues, CanonicalizedQueryResponse, QueryResponse,
+};
 use rustc_infer::infer::{InferCtxt, InferOk};
+use rustc_infer::traits::query::Fallible;
 use rustc_infer::traits::{
     FulfillmentError, Obligation, ObligationCause, PredicateObligation, TraitEngineExt as _,
 };
+use rustc_middle::arena::ArenaAllocatable;
 use rustc_middle::ty::error::TypeError;
 use rustc_middle::ty::ToPredicate;
 use rustc_middle::ty::TypeFoldable;
@@ -154,4 +160,20 @@ impl<'a, 'tcx> ObligationCtxt<'a, 'tcx> {
         }
         implied_bounds
     }
+
+    pub fn make_canonicalized_query_response<T>(
+        &self,
+        inference_vars: CanonicalVarValues<'tcx>,
+        answer: T,
+    ) -> Fallible<CanonicalizedQueryResponse<'tcx, T>>
+    where
+        T: Debug + TypeFoldable<'tcx>,
+        Canonical<'tcx, QueryResponse<'tcx, T>>: ArenaAllocatable<'tcx>,
+    {
+        self.infcx.make_canonicalized_query_response(
+            inference_vars,
+            answer,
+            &mut **self.engine.borrow_mut(),
+        )
+    }
 }