Add and use ObligationCtxt::new_in_snapshot

2 files changed, 22 insertions, 13 deletions

diff --git a/compiler/rustc_trait_selection/src/traits/engine.rs b/compiler/rustc_trait_selection/src/traits/engine.rs
index 136b9432145..f6df407118c 100644
--- a/compiler/rustc_trait_selection/src/traits/engine.rs
+++ b/compiler/rustc_trait_selection/src/traits/engine.rs
@@ -17,6 +17,7 @@ use rustc_span::Span;
 
 pub trait TraitEngineExt<'tcx> {
     fn new(tcx: TyCtxt<'tcx>) -> Box<Self>;
+    fn new_in_snapshot(tcx: TyCtxt<'tcx>) -> Box<Self>;
 }
 
 impl<'tcx> TraitEngineExt<'tcx> for dyn TraitEngine<'tcx> {
@@ -27,6 +28,14 @@ impl<'tcx> TraitEngineExt<'tcx> for dyn TraitEngine<'tcx> {
             Box::new(FulfillmentContext::new())
         }
     }
+
+    fn new_in_snapshot(tcx: TyCtxt<'tcx>) -> Box<Self> {
+        if tcx.sess.opts.unstable_opts.chalk {
+            Box::new(ChalkFulfillmentContext::new())
+        } else {
+            Box::new(FulfillmentContext::new_in_snapshot())
+        }
+    }
 }
 
 /// Used if you want to have pleasant experience when dealing
@@ -41,6 +50,10 @@ impl<'a, 'tcx> ObligationCtxt<'a, 'tcx> {
         Self { infcx, engine: RefCell::new(<dyn TraitEngine<'_>>::new(infcx.tcx)) }
     }
 
+    pub fn new_in_snapshot(infcx: &'a InferCtxt<'a, 'tcx>) -> Self {
+        Self { infcx, engine: RefCell::new(<dyn TraitEngine<'_>>::new_in_snapshot(infcx.tcx)) }
+    }
+
     pub fn register_obligation(&self, obligation: PredicateObligation<'tcx>) {
         self.engine.borrow_mut().register_predicate_obligation(self.infcx, obligation);
     }
diff --git a/compiler/rustc_typeck/src/check/inherited.rs b/compiler/rustc_typeck/src/check/inherited.rs
index 371fc26154c..1439baf5440 100644
--- a/compiler/rustc_typeck/src/check/inherited.rs
+++ b/compiler/rustc_typeck/src/check/inherited.rs
@@ -7,7 +7,6 @@ use rustc_hir::def_id::LocalDefId;
 use rustc_hir::HirIdMap;
 use rustc_infer::infer;
 use rustc_infer::infer::{InferCtxt, InferOk, TyCtxtInferExt};
-use rustc_infer::traits::TraitEngineExt as _;
 use rustc_middle::ty::fold::TypeFoldable;
 use rustc_middle::ty::visit::TypeVisitable;
 use rustc_middle::ty::{self, Ty, TyCtxt};
@@ -15,7 +14,7 @@ use rustc_span::def_id::LocalDefIdMap;
 use rustc_span::{self, Span};
 use rustc_trait_selection::infer::InferCtxtExt as _;
 use rustc_trait_selection::traits::{
-    self, FulfillmentContext, ObligationCause, TraitEngine, TraitEngineExt as _,
+    self, ObligationCause, ObligationCtxt, TraitEngine, TraitEngineExt as _,
 };
 
 use std::cell::RefCell;
@@ -94,17 +93,14 @@ impl<'tcx> Inherited<'_, 'tcx> {
                         return fn_sig;
                     }
                     infcx.probe(|_| {
-                        let traits::Normalized { value: normalized_fn_sig, obligations } =
-                            traits::normalize(
-                                &mut traits::SelectionContext::new(infcx),
-                                // FIXME(compiler-errors): This is probably not the right param-env...
-                                infcx.tcx.param_env(def_id),
-                                ObligationCause::dummy(),
-                                fn_sig,
-                            );
-                        let mut fulfillment_ctxt = FulfillmentContext::new_in_snapshot();
-                        fulfillment_ctxt.register_predicate_obligations(infcx, obligations);
-                        if fulfillment_ctxt.select_all_or_error(infcx).is_empty() {
+                        let ocx = ObligationCtxt::new_in_snapshot(infcx);
+                        let normalized_fn_sig = ocx.normalize(
+                            ObligationCause::dummy(),
+                            // FIXME(compiler-errors): This is probably not the right param-env...
+                            infcx.tcx.param_env(def_id),
+                            fn_sig,
+                        );
+                        if ocx.select_all_or_error().is_empty() {
                             let normalized_fn_sig =
                                 infcx.resolve_vars_if_possible(normalized_fn_sig);
                             if !normalized_fn_sig.needs_infer() {