Rollup merge of #104206 - compiler-errors:ocx-more-2, r=lcnr

Remove `save_and_restore_in_snapshot_flag`, use `ObligationCtxt` more

r? ```@lcnr```

4 files changed, 128 insertions, 188 deletions

diff --git a/compiler/rustc_hir_typeck/src/fn_ctxt/_impl.rs b/compiler/rustc_hir_typeck/src/fn_ctxt/_impl.rs
index 6ed7a93d463..c6bd771fad2 100644
--- a/compiler/rustc_hir_typeck/src/fn_ctxt/_impl.rs
+++ b/compiler/rustc_hir_typeck/src/fn_ctxt/_impl.rs
@@ -18,6 +18,7 @@ use rustc_infer::infer::canonical::{Canonical, OriginalQueryValues, QueryRespons
 use rustc_infer::infer::error_reporting::TypeAnnotationNeeded::E0282;
 use rustc_infer::infer::{InferOk, InferResult};
 use rustc_middle::ty::adjustment::{Adjust, Adjustment, AutoBorrow, AutoBorrowMutability};
+use rustc_middle::ty::error::TypeError;
 use rustc_middle::ty::fold::TypeFoldable;
 use rustc_middle::ty::visit::TypeVisitable;
 use rustc_middle::ty::{
@@ -32,9 +33,7 @@ use rustc_span::symbol::{kw, sym, Ident};
 use rustc_span::{Span, DUMMY_SP};
 use rustc_trait_selection::infer::InferCtxtExt as _;
 use rustc_trait_selection::traits::error_reporting::TypeErrCtxtExt as _;
-use rustc_trait_selection::traits::{
-    self, ObligationCause, ObligationCauseCode, TraitEngine, TraitEngineExt,
-};
+use rustc_trait_selection::traits::{self, ObligationCause, ObligationCauseCode, ObligationCtxt};
 
 use std::collections::hash_map::Entry;
 use std::slice;
@@ -766,34 +765,16 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
 
         let expect_args = self
             .fudge_inference_if_ok(|| {
+                let ocx = ObligationCtxt::new_in_snapshot(self);
+
                 // Attempt to apply a subtyping relationship between the formal
                 // return type (likely containing type variables if the function
                 // is polymorphic) and the expected return type.
                 // No argument expectations are produced if unification fails.
                 let origin = self.misc(call_span);
-                let ures = self.at(&origin, self.param_env).sup(ret_ty, formal_ret);
-
-                // FIXME(#27336) can't use ? here, Try::from_error doesn't default
-                // to identity so the resulting type is not constrained.
-                match ures {
-                    Ok(ok) => {
-                        // Process any obligations locally as much as
-                        // we can.  We don't care if some things turn
-                        // out unconstrained or ambiguous, as we're
-                        // just trying to get hints here.
-                        let errors = self.save_and_restore_in_snapshot_flag(|_| {
-                            let mut fulfill = <dyn TraitEngine<'_>>::new(self.tcx);
-                            for obligation in ok.obligations {
-                                fulfill.register_predicate_obligation(self, obligation);
-                            }
-                            fulfill.select_where_possible(self)
-                        });
-
-                        if !errors.is_empty() {
-                            return Err(());
-                        }
-                    }
-                    Err(_) => return Err(()),
+                ocx.sup(&origin, self.param_env, ret_ty, formal_ret)?;
+                if !ocx.select_where_possible().is_empty() {
+                    return Err(TypeError::Mismatch);
                 }
 
                 // Record all the argument types, with the substitutions
diff --git a/compiler/rustc_infer/src/infer/mod.rs b/compiler/rustc_infer/src/infer/mod.rs
index ccba197dc80..fd3b3e4d59f 100644
--- a/compiler/rustc_infer/src/infer/mod.rs
+++ b/compiler/rustc_infer/src/infer/mod.rs
@@ -778,32 +778,6 @@ impl<'tcx> InferCtxt<'tcx> {
         }
     }
 
-    /// Clear the "currently in a snapshot" flag, invoke the closure,
-    /// then restore the flag to its original value. This flag is a
-    /// debugging measure designed to detect cases where we start a
-    /// snapshot, create type variables, and register obligations
-    /// which may involve those type variables in the fulfillment cx,
-    /// potentially leaving "dangling type variables" behind.
-    /// In such cases, an assertion will fail when attempting to
-    /// register obligations, within a snapshot. Very useful, much
-    /// better than grovelling through megabytes of `RUSTC_LOG` output.
-    ///
-    /// HOWEVER, in some cases the flag is unhelpful. In particular, we
-    /// sometimes create a "mini-fulfilment-cx" in which we enroll
-    /// obligations. As long as this fulfillment cx is fully drained
-    /// before we return, this is not a problem, as there won't be any
-    /// escaping obligations in the main cx. In those cases, you can
-    /// use this function.
-    pub fn save_and_restore_in_snapshot_flag<F, R>(&self, func: F) -> R
-    where
-        F: FnOnce(&Self) -> R,
-    {
-        let flag = self.in_snapshot.replace(false);
-        let result = func(self);
-        self.in_snapshot.set(flag);
-        result
-    }
-
     fn start_snapshot(&self) -> CombinedSnapshot<'tcx> {
         debug!("start_snapshot()");
 
diff --git a/compiler/rustc_trait_selection/src/traits/specialize/mod.rs b/compiler/rustc_trait_selection/src/traits/specialize/mod.rs
index 43819b3f490..231a18f86ea 100644
--- a/compiler/rustc_trait_selection/src/traits/specialize/mod.rs
+++ b/compiler/rustc_trait_selection/src/traits/specialize/mod.rs
@@ -10,10 +10,12 @@
 //! [rustc dev guide]: https://rustc-dev-guide.rust-lang.org/traits/specialization.html
 
 pub mod specialization_graph;
+use rustc_infer::traits::{TraitEngine, TraitEngineExt as _};
 use specialization_graph::GraphExt;
 
 use crate::errors::NegativePositiveConflict;
 use crate::infer::{InferCtxt, InferOk, TyCtxtInferExt};
+use crate::traits::engine::TraitEngineExt as _;
 use crate::traits::select::IntercrateAmbiguityCause;
 use crate::traits::{self, coherence, FutureCompatOverlapErrorKind, ObligationCause};
 use rustc_data_structures::fx::FxIndexSet;
@@ -200,36 +202,32 @@ fn fulfill_implication<'tcx>(
         return Err(());
     };
 
+    // Needs to be `in_snapshot` because this function is used to rebase
+    // substitutions, which may happen inside of a select within a probe.
+    let mut engine = <dyn TraitEngine<'tcx>>::new_in_snapshot(infcx.tcx);
     // attempt to prove all of the predicates for impl2 given those for impl1
     // (which are packed up in penv)
+    engine.register_predicate_obligations(infcx, obligations.chain(more_obligations));
 
-    infcx.save_and_restore_in_snapshot_flag(|infcx| {
-        let errors = traits::fully_solve_obligations(&infcx, obligations.chain(more_obligations));
-        match &errors[..] {
-            [] => {
-                debug!(
-                    "fulfill_implication: an impl for {:?} specializes {:?}",
-                    source_trait, target_trait
-                );
+    let errors = engine.select_all_or_error(infcx);
+    if !errors.is_empty() {
+        // no dice!
+        debug!(
+            "fulfill_implication: for impls on {:?} and {:?}, \
+                 could not fulfill: {:?} given {:?}",
+            source_trait,
+            target_trait,
+            errors,
+            param_env.caller_bounds()
+        );
+        return Err(());
+    }
 
-                // Now resolve the *substitution* we built for the target earlier, replacing
-                // the inference variables inside with whatever we got from fulfillment.
-                Ok(infcx.resolve_vars_if_possible(target_substs))
-            }
-            errors => {
-                // no dice!
-                debug!(
-                    "fulfill_implication: for impls on {:?} and {:?}, \
-                     could not fulfill: {:?} given {:?}",
-                    source_trait,
-                    target_trait,
-                    errors,
-                    param_env.caller_bounds()
-                );
-                Err(())
-            }
-        }
-    })
+    debug!("fulfill_implication: an impl for {:?} specializes {:?}", source_trait, target_trait);
+
+    // Now resolve the *substitution* we built for the target earlier, replacing
+    // the inference variables inside with whatever we got from fulfillment.
+    Ok(infcx.resolve_vars_if_possible(target_substs))
 }
 
 // Query provider for `specialization_graph_of`.
diff --git a/compiler/rustc_traits/src/dropck_outlives.rs b/compiler/rustc_traits/src/dropck_outlives.rs
index d5a8ca5ea78..7b4ad9fea13 100644
--- a/compiler/rustc_traits/src/dropck_outlives.rs
+++ b/compiler/rustc_traits/src/dropck_outlives.rs
@@ -2,20 +2,18 @@ use rustc_data_structures::fx::FxHashSet;
 use rustc_hir::def_id::DefId;
 use rustc_infer::infer::canonical::{Canonical, QueryResponse};
 use rustc_infer::infer::TyCtxtInferExt;
-use rustc_infer::traits::TraitEngineExt as _;
 use rustc_middle::ty::query::Providers;
 use rustc_middle::ty::InternalSubsts;
 use rustc_middle::ty::{self, EarlyBinder, ParamEnvAnd, Ty, TyCtxt};
 use rustc_span::source_map::{Span, DUMMY_SP};
+use rustc_trait_selection::infer::InferCtxtBuilderExt;
 use rustc_trait_selection::traits::query::dropck_outlives::trivial_dropck_outlives;
 use rustc_trait_selection::traits::query::dropck_outlives::{
     DropckConstraint, DropckOutlivesResult,
 };
 use rustc_trait_selection::traits::query::normalize::AtExt;
 use rustc_trait_selection::traits::query::{CanonicalTyGoal, NoSolution};
-use rustc_trait_selection::traits::{
-    Normalized, ObligationCause, TraitEngine, TraitEngineExt as _,
-};
+use rustc_trait_selection::traits::{Normalized, ObligationCause};
 
 pub(crate) fn provide(p: &mut Providers) {
     *p = Providers { dropck_outlives, adt_dtorck_constraint, ..*p };
@@ -27,120 +25,109 @@ fn dropck_outlives<'tcx>(
 ) -> Result<&'tcx Canonical<'tcx, QueryResponse<'tcx, DropckOutlivesResult<'tcx>>>, NoSolution> {
     debug!("dropck_outlives(goal={:#?})", canonical_goal);
 
-    let (ref infcx, goal, canonical_inference_vars) =
-        tcx.infer_ctxt().build_with_canonical(DUMMY_SP, &canonical_goal);
-    let tcx = infcx.tcx;
-    let ParamEnvAnd { param_env, value: for_ty } = goal;
-
-    let mut result = DropckOutlivesResult { kinds: vec![], overflows: vec![] };
-
-    // A stack of types left to process. Each round, we pop
-    // something from the stack and invoke
-    // `dtorck_constraint_for_ty`. This may produce new types that
-    // have to be pushed on the stack. This continues until we have explored
-    // all the reachable types from the type `for_ty`.
-    //
-    // Example: Imagine that we have the following code:
-    //
-    // ```rust
-    // struct A {
-    //     value: B,
-    //     children: Vec<A>,
-    // }
-    //
-    // struct B {
-    //     value: u32
-    // }
-    //
-    // fn f() {
-    //   let a: A = ...;
-    //   ..
-    // } // here, `a` is dropped
-    // ```
-    //
-    // at the point where `a` is dropped, we need to figure out
-    // which types inside of `a` contain region data that may be
-    // accessed by any destructors in `a`. We begin by pushing `A`
-    // onto the stack, as that is the type of `a`. We will then
-    // invoke `dtorck_constraint_for_ty` which will expand `A`
-    // into the types of its fields `(B, Vec<A>)`. These will get
-    // pushed onto the stack. Eventually, expanding `Vec<A>` will
-    // lead to us trying to push `A` a second time -- to prevent
-    // infinite recursion, we notice that `A` was already pushed
-    // once and stop.
-    let mut ty_stack = vec![(for_ty, 0)];
-
-    // Set used to detect infinite recursion.
-    let mut ty_set = FxHashSet::default();
-
-    let mut fulfill_cx = <dyn TraitEngine<'_>>::new(infcx.tcx);
-
-    let cause = ObligationCause::dummy();
-    let mut constraints = DropckConstraint::empty();
-    while let Some((ty, depth)) = ty_stack.pop() {
-        debug!(
-            "{} kinds, {} overflows, {} ty_stack",
-            result.kinds.len(),
-            result.overflows.len(),
-            ty_stack.len()
-        );
-        dtorck_constraint_for_ty(tcx, DUMMY_SP, for_ty, depth, ty, &mut constraints)?;
-
-        // "outlives" represent types/regions that may be touched
-        // by a destructor.
-        result.kinds.append(&mut constraints.outlives);
-        result.overflows.append(&mut constraints.overflows);
-
-        // If we have even one overflow, we should stop trying to evaluate further --
-        // chances are, the subsequent overflows for this evaluation won't provide useful
-        // information and will just decrease the speed at which we can emit these errors
-        // (since we'll be printing for just that much longer for the often enormous types
-        // that result here).
-        if !result.overflows.is_empty() {
-            break;
-        }
+    tcx.infer_ctxt().enter_canonical_trait_query(&canonical_goal, |ocx, goal| {
+        let tcx = ocx.infcx.tcx;
+        let ParamEnvAnd { param_env, value: for_ty } = goal;
+
+        let mut result = DropckOutlivesResult { kinds: vec![], overflows: vec![] };
+
+        // A stack of types left to process. Each round, we pop
+        // something from the stack and invoke
+        // `dtorck_constraint_for_ty`. This may produce new types that
+        // have to be pushed on the stack. This continues until we have explored
+        // all the reachable types from the type `for_ty`.
+        //
+        // Example: Imagine that we have the following code:
+        //
+        // ```rust
+        // struct A {
+        //     value: B,
+        //     children: Vec<A>,
+        // }
+        //
+        // struct B {
+        //     value: u32
+        // }
+        //
+        // fn f() {
+        //   let a: A = ...;
+        //   ..
+        // } // here, `a` is dropped
+        // ```
+        //
+        // at the point where `a` is dropped, we need to figure out
+        // which types inside of `a` contain region data that may be
+        // accessed by any destructors in `a`. We begin by pushing `A`
+        // onto the stack, as that is the type of `a`. We will then
+        // invoke `dtorck_constraint_for_ty` which will expand `A`
+        // into the types of its fields `(B, Vec<A>)`. These will get
+        // pushed onto the stack. Eventually, expanding `Vec<A>` will
+        // lead to us trying to push `A` a second time -- to prevent
+        // infinite recursion, we notice that `A` was already pushed
+        // once and stop.
+        let mut ty_stack = vec![(for_ty, 0)];
+
+        // Set used to detect infinite recursion.
+        let mut ty_set = FxHashSet::default();
+
+        let cause = ObligationCause::dummy();
+        let mut constraints = DropckConstraint::empty();
+        while let Some((ty, depth)) = ty_stack.pop() {
+            debug!(
+                "{} kinds, {} overflows, {} ty_stack",
+                result.kinds.len(),
+                result.overflows.len(),
+                ty_stack.len()
+            );
+            dtorck_constraint_for_ty(tcx, DUMMY_SP, for_ty, depth, ty, &mut constraints)?;
+
+            // "outlives" represent types/regions that may be touched
+            // by a destructor.
+            result.kinds.append(&mut constraints.outlives);
+            result.overflows.append(&mut constraints.overflows);
+
+            // If we have even one overflow, we should stop trying to evaluate further --
+            // chances are, the subsequent overflows for this evaluation won't provide useful
+            // information and will just decrease the speed at which we can emit these errors
+            // (since we'll be printing for just that much longer for the often enormous types
+            // that result here).
+            if !result.overflows.is_empty() {
+                break;
+            }
 
-        // dtorck types are "types that will get dropped but which
-        // do not themselves define a destructor", more or less. We have
-        // to push them onto the stack to be expanded.
-        for ty in constraints.dtorck_types.drain(..) {
-            match infcx.at(&cause, param_env).normalize(ty) {
-                Ok(Normalized { value: ty, obligations }) => {
-                    fulfill_cx.register_predicate_obligations(infcx, obligations);
-
-                    debug!("dropck_outlives: ty from dtorck_types = {:?}", ty);
-
-                    match ty.kind() {
-                        // All parameters live for the duration of the
-                        // function.
-                        ty::Param(..) => {}
-
-                        // A projection that we couldn't resolve - it
-                        // might have a destructor.
-                        ty::Projection(..) | ty::Opaque(..) => {
-                            result.kinds.push(ty.into());
-                        }
+            // dtorck types are "types that will get dropped but which
+            // do not themselves define a destructor", more or less. We have
+            // to push them onto the stack to be expanded.
+            for ty in constraints.dtorck_types.drain(..) {
+                let Normalized { value: ty, obligations } =
+                    ocx.infcx.at(&cause, param_env).normalize(ty)?;
+                ocx.register_obligations(obligations);
+
+                debug!("dropck_outlives: ty from dtorck_types = {:?}", ty);
+
+                match ty.kind() {
+                    // All parameters live for the duration of the
+                    // function.
+                    ty::Param(..) => {}
+
+                    // A projection that we couldn't resolve - it
+                    // might have a destructor.
+                    ty::Projection(..) | ty::Opaque(..) => {
+                        result.kinds.push(ty.into());
+                    }
 
-                        _ => {
-                            if ty_set.insert(ty) {
-                                ty_stack.push((ty, depth + 1));
-                            }
+                    _ => {
+                        if ty_set.insert(ty) {
+                            ty_stack.push((ty, depth + 1));
                         }
                     }
                 }
-
-                // We don't actually expect to fail to normalize.
-                // That implies a WF error somewhere else.
-                Err(NoSolution) => {
-                    return Err(NoSolution);
-                }
             }
         }
-    }
-
-    debug!("dropck_outlives: result = {:#?}", result);
 
-    infcx.make_canonicalized_query_response(canonical_inference_vars, result, &mut *fulfill_cx)
+        debug!("dropck_outlives: result = {:#?}", result);
+        Ok(result)
+    })
 }
 
 /// Returns a set of constraints that needs to be satisfied in