snapshot: avoid leaking inference vars

compiler/rustc_hir_typeck/src/coercion.rs

@@ -1076,7 +1076,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         let coerce = Coerce::new(self, cause, AllowTwoPhase::No);
         coerce
             .autoderef(rustc_span::DUMMY_SP, expr_ty)
-            .find_map(|(ty, steps)| self.probe(|_| coerce.unify(ty, target)).ok().map(|_| steps))
+            .find_map(|(ty, steps)| self.probe(|_| coerce.unify(ty, target).ok().map(|_| steps)))
     }
 
     /// Given a type, this function will calculate and return the type given

compiler/rustc_hir_typeck/src/method/mod.rs

@@ -16,6 +16,7 @@ use rustc_hir as hir;
 use rustc_hir::def::{CtorOf, DefKind, Namespace};
 use rustc_hir::def_id::DefId;
 use rustc_infer::infer::{self, InferOk};
+use rustc_infer::trivial_no_snapshot_leaks;
 use rustc_middle::query::Providers;
 use rustc_middle::traits::ObligationCause;
 use rustc_middle::ty::{self, GenericParamDefKind, Ty, TypeVisitableExt};
@@ -43,6 +44,8 @@ pub struct MethodCallee<'tcx> {
     pub sig: ty::FnSig<'tcx>,
 }
 
+// FIXME(#122188): This is wrong, as this type may leak inference vars.
+trivial_no_snapshot_leaks!('tcx, MethodError<'tcx>);
 #[derive(Debug)]
 pub enum MethodError<'tcx> {
     // Did not find an applicable method, but we did find various near-misses that may work.
@@ -79,8 +82,9 @@ pub struct NoMatchData<'tcx> {
     pub mode: probe::Mode,
 }
 
-// A pared down enum describing just the places from which a method
-// candidate can arise. Used for error reporting only.
+trivial_no_snapshot_leaks!('tcx, CandidateSource);
+/// A pared down enum describing just the places from which a method
+/// candidate can arise. Used for error reporting only.
 #[derive(Copy, Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
 pub enum CandidateSource {
     Impl(DefId),

compiler/rustc_hir_typeck/src/method/probe.rs

@@ -13,8 +13,10 @@ use rustc_hir_analysis::autoderef::{self, Autoderef};
 use rustc_infer::infer::canonical::OriginalQueryValues;
 use rustc_infer::infer::canonical::{Canonical, QueryResponse};
 use rustc_infer::infer::error_reporting::TypeAnnotationNeeded::E0282;
+use rustc_infer::infer::snapshot::NoSnapshotLeaks;
 use rustc_infer::infer::DefineOpaqueTypes;
 use rustc_infer::infer::{self, InferOk, TyCtxtInferExt};
+use rustc_infer::trivial_no_snapshot_leaks;
 use rustc_middle::middle::stability;
 use rustc_middle::query::Providers;
 use rustc_middle::ty::fast_reject::{simplify_type, TreatParams};
@@ -97,6 +99,8 @@ impl<'a, 'tcx> Deref for ProbeContext<'a, 'tcx> {
     }
 }
 
+// FIXME(#122188): This is wrong as this type may leak inference variables.
+trivial_no_snapshot_leaks!('tcx, Candidate<'tcx>);
 #[derive(Debug, Clone)]
 pub(crate) struct Candidate<'tcx> {
     // Candidates are (I'm not quite sure, but they are mostly) basically
@@ -152,6 +156,7 @@ pub(crate) enum CandidateKind<'tcx> {
     ),
 }
 
+trivial_no_snapshot_leaks!('tcx, ProbeResult);
 #[derive(Debug, PartialEq, Eq, Copy, Clone)]
 enum ProbeResult {
     NoMatch,
@@ -195,6 +200,8 @@ impl AutorefOrPtrAdjustment {
     }
 }
 
+// FIXME(#122188): This is wrong as this type may leak inference variables.
+trivial_no_snapshot_leaks!('tcx, Pick<'tcx>);
 #[derive(Debug, Clone)]
 pub struct Pick<'tcx> {
     pub item: ty::AssocItem,
@@ -368,6 +375,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         op: OP,
     ) -> Result<R, MethodError<'tcx>>
     where
+        R: NoSnapshotLeaks<'tcx>,
         OP: FnOnce(ProbeContext<'_, 'tcx>) -> Result<R, MethodError<'tcx>>,
     {
         let mut orig_values = OriginalQueryValues::default();

compiler/rustc_infer/src/infer/mod.rs

@@ -62,7 +62,7 @@ mod projection;
 pub mod region_constraints;
 mod relate;
 pub mod resolve;
-pub(crate) mod snapshot;
+pub mod snapshot;
 pub mod type_variable;
 
 #[must_use]

compiler/rustc_infer/src/infer/snapshot/check_leaks.rs

@@ -0,0 +1,140 @@
+use super::VariableLengths;
+use crate::infer::InferCtxt;
+use rustc_middle::ty::{self, Ty, TyCtxt};
+use rustc_middle::ty::{TypeSuperVisitable, TypeVisitor};
+use std::ops::ControlFlow;
+
+/// Check for leaking inference variables and placeholders
+/// from snapshot. This is only used if `debug_assertions`
+/// are enabled.
+pub struct HasSnapshotLeaksVisitor {
+    universe: ty::UniverseIndex,
+    variable_lengths: VariableLengths,
+}
+impl HasSnapshotLeaksVisitor {
+    pub fn new<'tcx>(infcx: &InferCtxt<'tcx>) -> Self {
+        HasSnapshotLeaksVisitor {
+            universe: infcx.universe(),
+            variable_lengths: infcx.variable_lengths(),
+        }
+    }
+}
+
+fn continue_if(b: bool) -> ControlFlow<()> {
+    if b { ControlFlow::Continue(()) } else { ControlFlow::Break(()) }
+}
+
+impl<'tcx> TypeVisitor<TyCtxt<'tcx>> for HasSnapshotLeaksVisitor {
+    type Result = ControlFlow<()>;
+
+    fn visit_region(&mut self, r: ty::Region<'tcx>) -> Self::Result {
+        match r.kind() {
+            ty::ReVar(var) => continue_if(var.as_usize() < self.variable_lengths.region_vars),
+            ty::RePlaceholder(p) => continue_if(self.universe.can_name(p.universe)),
+            ty::ReEarlyParam(_)
+            | ty::ReBound(_, _)
+            | ty::ReLateParam(_)
+            | ty::ReStatic
+            | ty::ReErased
+            | ty::ReError(_) => ControlFlow::Continue(()),
+        }
+    }
+    fn visit_ty(&mut self, t: Ty<'tcx>) -> Self::Result {
+        match t.kind() {
+            ty::Infer(ty::TyVar(var)) => {
+                continue_if(var.as_usize() < self.variable_lengths.type_vars)
+            }
+            ty::Infer(ty::IntVar(var)) => {
+                continue_if(var.as_usize() < self.variable_lengths.int_vars)
+            }
+            ty::Infer(ty::FloatVar(var)) => {
+                continue_if(var.as_usize() < self.variable_lengths.float_vars)
+            }
+            ty::Placeholder(p) => continue_if(self.universe.can_name(p.universe)),
+            ty::Infer(ty::FreshTy(..) | ty::FreshIntTy(..) | ty::FreshFloatTy(..))
+            | ty::Bool
+            | ty::Char
+            | ty::Int(_)
+            | ty::Uint(_)
+            | ty::Float(_)
+            | ty::Adt(_, _)
+            | ty::Foreign(_)
+            | ty::Str
+            | ty::Array(_, _)
+            | ty::Slice(_)
+            | ty::RawPtr(_)
+            | ty::Ref(_, _, _)
+            | ty::FnDef(_, _)
+            | ty::FnPtr(_)
+            | ty::Dynamic(_, _, _)
+            | ty::Closure(_, _)
+            | ty::CoroutineClosure(_, _)
+            | ty::Coroutine(_, _)
+            | ty::CoroutineWitness(_, _)
+            | ty::Never
+            | ty::Tuple(_)
+            | ty::Alias(_, _)
+            | ty::Param(_)
+            | ty::Bound(_, _)
+            | ty::Error(_) => t.super_visit_with(self),
+        }
+    }
+    fn visit_const(&mut self, c: ty::Const<'tcx>) -> Self::Result {
+        match c.kind() {
+            ty::ConstKind::Infer(ty::InferConst::Var(var)) => {
+                continue_if(var.as_usize() < self.variable_lengths.const_vars)
+            }
+            // FIXME(const_trait_impl): need to handle effect vars here and in `fudge_inference_if_ok`.
+            ty::ConstKind::Infer(ty::InferConst::EffectVar(_)) => ControlFlow::Continue(()),
+            ty::ConstKind::Placeholder(p) => continue_if(self.universe.can_name(p.universe)),
+            ty::ConstKind::Infer(ty::InferConst::Fresh(_))
+            | ty::ConstKind::Param(_)
+            | ty::ConstKind::Bound(_, _)
+            | ty::ConstKind::Unevaluated(_)
+            | ty::ConstKind::Value(_)
+            | ty::ConstKind::Expr(_)
+            | ty::ConstKind::Error(_) => c.super_visit_with(self),
+        }
+    }
+}
+
+#[macro_export]
+#[cfg(debug_assertions)]
+macro_rules! type_foldable_verify_no_snapshot_leaks {
+    ($tcx:lifetime, $t:ty) => {
+        const _: () = {
+            use rustc_middle::ty::TypeVisitable;
+            use $crate::infer::snapshot::check_leaks::HasSnapshotLeaksVisitor;
+            use $crate::infer::InferCtxt;
+            impl<$tcx> $crate::infer::snapshot::NoSnapshotLeaks<$tcx> for $t {
+                type StartData = HasSnapshotLeaksVisitor;
+                type EndData = ($t, HasSnapshotLeaksVisitor);
+                fn snapshot_start_data(infcx: &$crate::infer::InferCtxt<$tcx>) -> Self::StartData {
+                    HasSnapshotLeaksVisitor::new(infcx)
+                }
+                fn end_of_snapshot(
+                    _: &InferCtxt<'tcx>,
+                    value: $t,
+                    visitor: Self::StartData,
+                ) -> Self::EndData {
+                    (value, visitor)
+                }
+                fn avoid_leaks(_: &InferCtxt<$tcx>, (value, mut visitor): Self::EndData) -> Self {
+                    if value.visit_with(&mut visitor).is_break() {
+                        bug!("leaking vars from snapshot: {value:?}");
+                    }
+
+                    value
+                }
+            }
+        };
+    };
+}
+
+#[macro_export]
+#[cfg(not(debug_assertions))]
+macro_rules! type_foldable_verify_no_snapshot_leaks {
+    ($tcx:lifetime, $t:ty) => {
+        trivial_no_snapshot_leaks!($tcx, $t);
+    };
+}