Use ObligationCtxt in custom type ops

compiler/rustc_trait_selection/src/traits/query/type_op/custom.rs

@@ -1,32 +1,31 @@
 use crate::infer::canonical::query_response;
-use crate::infer::{InferCtxt, InferOk};
+use crate::infer::InferCtxt;
 use crate::traits::query::type_op::TypeOpOutput;
 use crate::traits::query::Fallible;
 use crate::traits::ObligationCtxt;
 use rustc_infer::infer::region_constraints::RegionConstraintData;
+use rustc_middle::traits::query::NoSolution;
 use rustc_span::source_map::DUMMY_SP;
 
 use std::fmt;
 
-pub struct CustomTypeOp<F, G> {
+pub struct CustomTypeOp<F> {
     closure: F,
-    description: G,
+    description: &'static str,
 }
 
-impl<F, G> CustomTypeOp<F, G> {
-    pub fn new<'tcx, R>(closure: F, description: G) -> Self
+impl<F> CustomTypeOp<F> {
+    pub fn new<'tcx, R>(closure: F, description: &'static str) -> Self
     where
-        F: FnOnce(&InferCtxt<'tcx>) -> Fallible<InferOk<'tcx, R>>,
-        G: Fn() -> String,
+        F: FnOnce(&ObligationCtxt<'_, 'tcx>) -> Fallible<R>,
     {
         CustomTypeOp { closure, description }
     }
 }
 
-impl<'tcx, F, R: fmt::Debug, G> super::TypeOp<'tcx> for CustomTypeOp<F, G>
+impl<'tcx, F, R: fmt::Debug> super::TypeOp<'tcx> for CustomTypeOp<F>
 where
-    F: for<'a, 'cx> FnOnce(&'a InferCtxt<'tcx>) -> Fallible<InferOk<'tcx, R>>,
-    G: Fn() -> String,
+    F: FnOnce(&ObligationCtxt<'_, 'tcx>) -> Fallible<R>,
 {
     type Output = R;
     /// We can't do any custom error reporting for `CustomTypeOp`, so
@@ -41,16 +40,13 @@ where
             info!("fully_perform({:?})", self);
         }
 
-        Ok(scrape_region_constraints(infcx, || (self.closure)(infcx))?.0)
+        Ok(scrape_region_constraints(infcx, self.closure)?.0)
     }
 }
 
-impl<F, G> fmt::Debug for CustomTypeOp<F, G>
-where
-    G: Fn() -> String,
-{
+impl<F> fmt::Debug for CustomTypeOp<F> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "{}", (self.description)())
+        self.description.fmt(f)
     }
 }
 
@@ -58,7 +54,7 @@ where
 /// constraints that result, creating query-region-constraints.
 pub fn scrape_region_constraints<'tcx, Op: super::TypeOp<'tcx, Output = R>, R>(
     infcx: &InferCtxt<'tcx>,
-    op: impl FnOnce() -> Fallible<InferOk<'tcx, R>>,
+    op: impl FnOnce(&ObligationCtxt<'_, 'tcx>) -> Fallible<R>,
 ) -> Fallible<(TypeOpOutput<'tcx, Op>, RegionConstraintData<'tcx>)> {
     // During NLL, we expect that nobody will register region
     // obligations **except** as part of a custom type op (and, at the
@@ -72,16 +68,20 @@ pub fn scrape_region_constraints<'tcx, Op: super::TypeOp<'tcx, Output = R>, R>(
         pre_obligations,
     );
 
-    let InferOk { value, obligations } = infcx.commit_if_ok(|_| op())?;
-    let ocx = ObligationCtxt::new(infcx);
-    ocx.register_obligations(obligations);
-    let errors = ocx.select_all_or_error();
-    if !errors.is_empty() {
-        infcx.tcx.sess.diagnostic().delay_span_bug(
-            DUMMY_SP,
-            format!("errors selecting obligation during MIR typeck: {:?}", errors),
-        );
-    }
+    let value = infcx.commit_if_ok(|_| {
+        let ocx = ObligationCtxt::new_in_snapshot(infcx);
+        let value = op(&ocx)?;
+        let errors = ocx.select_all_or_error();
+        if errors.is_empty() {
+            Ok(value)
+        } else {
+            infcx.tcx.sess.delay_span_bug(
+                DUMMY_SP,
+                format!("errors selecting obligation during MIR typeck: {:?}", errors),
+            );
+            Err(NoSolution)
+        }
+    })?;
 
     let region_obligations = infcx.take_registered_region_obligations();
     let region_constraint_data = infcx.take_and_reset_region_constraints();