snapshot: avoid leaking inference vars

Author: lcnr

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/fn_ctxt/mod.rs

@@ -13,6 +13,7 @@ use rustc_hir_analysis::astconv::AstConv;
 use rustc_infer::infer;
 use rustc_infer::infer::error_reporting::sub_relations::SubRelations;
 use rustc_infer::infer::error_reporting::TypeErrCtxt;
+use rustc_infer::infer::snapshot::NoLeaksUnchecked;
 use rustc_infer::infer::type_variable::{TypeVariableOrigin, TypeVariableOriginKind};
 use rustc_middle::infer::unify_key::{ConstVariableOrigin, ConstVariableOriginKind};
 use rustc_middle::ty::{self, Const, Ty, TyCtxt, TypeVisitableExt};
@@ -177,11 +178,12 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                     if ocx.select_all_or_error().is_empty() {
                         let normalized_fn_sig = self.resolve_vars_if_possible(normalized_fn_sig);
                         if !normalized_fn_sig.has_infer() {
-                            return normalized_fn_sig;
+                            return NoLeaksUnchecked { value: normalized_fn_sig };
                         }
                     }
-                    fn_sig
+                    NoLeaksUnchecked { value: fn_sig }
                 })
+                .value
             }),
             autoderef_steps: Box::new(|ty| {
                 let mut autoderef = self.autoderef(DUMMY_SP, ty).silence_errors();

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: This is wrong, method error 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: 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: 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/fudge.rs

@@ -12,6 +12,8 @@ use ut::UnifyKey;
 
 use std::ops::Range;
 
+use super::{NoSnapshotLeaks, VariableLengths};
+
 fn vars_since_snapshot<'tcx, T>(
     table: &mut UnificationTable<'_, 'tcx, T>,
     snapshot_var_len: usize,
@@ -88,82 +90,108 @@ impl<'tcx> InferCtxt<'tcx> {
     where
         F: FnOnce() -> Result<T, E>,
         T: TypeFoldable<TyCtxt<'tcx>>,
+        E: NoSnapshotLeaks<'tcx>,
     {
-        let variable_lengths = self.variable_lengths();
-        let (mut fudger, value) = self.probe_unchecked(|_| {
-            match f() {
-                Ok(value) => {
-                    let value = self.resolve_vars_if_possible(value);
-
-                    // At this point, `value` could in principle refer
-                    // to inference variables that have been created during
-                    // the snapshot. Once we exit `probe()`, those are
-                    // going to be popped, so we will have to
-                    // eliminate any references to them.
-
-                    let mut inner = self.inner.borrow_mut();
-                    let type_vars =
-                        inner.type_variables().vars_since_snapshot(variable_lengths.type_vars);
-                    let int_vars = vars_since_snapshot(
-                        &mut inner.int_unification_table(),
-                        variable_lengths.int_vars,
-                    );
-                    let float_vars = vars_since_snapshot(
-                        &mut inner.float_unification_table(),
-                        variable_lengths.float_vars,
-                    );
-                    let region_vars = inner
-                        .unwrap_region_constraints()
-                        .vars_since_snapshot(variable_lengths.region_vars);
-                    let const_vars = const_vars_since_snapshot(
-                        &mut inner.const_unification_table(),
-                        variable_lengths.const_vars,
-                    );
-
-                    let fudger = InferenceFudger {
-                        infcx: self,
-                        type_vars,
-                        int_vars,
-                        float_vars,
-                        region_vars,
-                        const_vars,
-                    };
+        self.probe(|_| f().map(|value| FudgeInference(self.resolve_vars_if_possible(value))))
+            .map(|FudgeInference(value)| value)
+    }
+}
 
-                    Ok((fudger, value))
+#[macro_export]
+macro_rules! fudge_vars_no_snapshot_leaks {
+    ($tcx:lifetime, $t:ty) => {
+        const _: () = {
+            use $crate::infer::snapshot::fudge::InferenceFudgeData;
+            impl<$tcx> $crate::infer::snapshot::NoSnapshotLeaks<$tcx> for $t {
+                type DataStart = $crate::infer::snapshot::VariableLengths;
+                type DataEnd = InferenceFudgeData;
+                fn mk_data_snapshot_start(infcx: &InferCtxt<$tcx>) -> Self::DataStart {
+                    infcx.variable_lengths()
+                }
+                fn mk_data_snapshot_end(
+                    infcx: &InferCtxt<$tcx>,
+                    variable_lengths: Self::DataStart,
+                ) -> Self::DataEnd {
+                    InferenceFudgeData::new(infcx, variable_lengths)
+                }
+                fn avoid_leaks(self, infcx: &InferCtxt<'tcx>, fudge_data: Self::DataEnd) -> Self {
+                    fudge_data.fudge_inference(infcx, self)
                 }
-                Err(e) => Err(e),
             }
-        })?;
-
-        // At this point, we need to replace any of the now-popped
-        // type/region variables that appear in `value` with a fresh
-        // variable of the appropriate kind. We can't do this during
-        // the probe because they would just get popped then too. =)
-
-        // Micro-optimization: if no variables have been created, then
-        // `value` can't refer to any of them. =) So we can just return it.
-        if fudger.type_vars.0.is_empty()
-            && fudger.int_vars.is_empty()
-            && fudger.float_vars.is_empty()
-            && fudger.region_vars.0.is_empty()
-            && fudger.const_vars.0.is_empty()
-        {
-            Ok(value)
-        } else {
-            Ok(value.fold_with(&mut fudger))
-        }
+        };
+    };
+}
+
+struct FudgeInference<T>(T);
+impl<'tcx, T: TypeFoldable<TyCtxt<'tcx>>> NoSnapshotLeaks<'tcx> for FudgeInference<T> {
+    type DataStart = VariableLengths;
+    type DataEnd = InferenceFudgeData;
+    fn mk_data_snapshot_start(infcx: &InferCtxt<'tcx>) -> Self::DataStart {
+        infcx.variable_lengths()
+    }
+    fn mk_data_snapshot_end(
+        infcx: &InferCtxt<'tcx>,
+        variable_lengths: Self::DataStart,
+    ) -> Self::DataEnd {
+        InferenceFudgeData::new(infcx, variable_lengths)
+    }
+    fn avoid_leaks(self, infcx: &InferCtxt<'tcx>, fudge_data: Self::DataEnd) -> Self {
+        FudgeInference(fudge_data.fudge_inference(infcx, self.0))
     }
 }
 
-pub struct InferenceFudger<'a, 'tcx> {
-    infcx: &'a InferCtxt<'tcx>,
+pub struct InferenceFudgeData {
     type_vars: (Range<TyVid>, Vec<TypeVariableOrigin>),
     int_vars: Range<IntVid>,
     float_vars: Range<FloatVid>,
     region_vars: (Range<RegionVid>, Vec<RegionVariableOrigin>),
     const_vars: (Range<ConstVid>, Vec<ConstVariableOrigin>),
 }
 
+impl InferenceFudgeData {
+    pub fn new<'tcx>(
+        infcx: &InferCtxt<'tcx>,
+        variable_lengths: VariableLengths,
+    ) -> InferenceFudgeData {
+        let mut inner = infcx.inner.borrow_mut();
+        let type_vars = inner.type_variables().vars_since_snapshot(variable_lengths.type_vars);
+        let int_vars =
+            vars_since_snapshot(&mut inner.int_unification_table(), variable_lengths.int_vars);
+        let float_vars =
+            vars_since_snapshot(&mut inner.float_unification_table(), variable_lengths.float_vars);
+        let region_vars =
+            inner.unwrap_region_constraints().vars_since_snapshot(variable_lengths.region_vars);
+        let const_vars = const_vars_since_snapshot(
+            &mut inner.const_unification_table(),
+            variable_lengths.const_vars,
+        );
+
+        InferenceFudgeData { type_vars, int_vars, float_vars, region_vars, const_vars }
+    }
+
+    pub fn fudge_inference<'tcx, T: TypeFoldable<TyCtxt<'tcx>>>(
+        self,
+        infcx: &InferCtxt<'tcx>,
+        value: T,
+    ) -> T {
+        if self.type_vars.0.is_empty()
+            && self.int_vars.is_empty()
+            && self.float_vars.is_empty()
+            && self.region_vars.0.is_empty()
+            && self.const_vars.0.is_empty()
+        {
+            value
+        } else {
+            value.fold_with(&mut InferenceFudger { infcx, data: self })
+        }
+    }
+}
+
+struct InferenceFudger<'a, 'tcx> {
+    infcx: &'a InferCtxt<'tcx>,
+    data: InferenceFudgeData,
+}
+
 impl<'a, 'tcx> TypeFolder<TyCtxt<'tcx>> for InferenceFudger<'a, 'tcx> {
     fn interner(&self) -> TyCtxt<'tcx> {
         self.infcx.tcx
@@ -172,11 +200,11 @@ impl<'a, 'tcx> TypeFolder<TyCtxt<'tcx>> for InferenceFudger<'a, 'tcx> {
     fn fold_ty(&mut self, ty: Ty<'tcx>) -> Ty<'tcx> {
         match *ty.kind() {
             ty::Infer(ty::InferTy::TyVar(vid)) => {
-                if self.type_vars.0.contains(&vid) {
+                if self.data.type_vars.0.contains(&vid) {
                     // This variable was created during the fudging.
                     // Recreate it with a fresh variable here.
-                    let idx = vid.as_usize() - self.type_vars.0.start.as_usize();
-                    let origin = self.type_vars.1[idx];
+                    let idx = vid.as_usize() - self.data.type_vars.0.start.as_usize();
+                    let origin = self.data.type_vars.1[idx];
                     self.infcx.next_ty_var(origin)
                 } else {
                     // This variable was created before the
@@ -191,14 +219,14 @@ impl<'a, 'tcx> TypeFolder<TyCtxt<'tcx>> for InferenceFudger<'a, 'tcx> {
                 }
             }
             ty::Infer(ty::InferTy::IntVar(vid)) => {
-                if self.int_vars.contains(&vid) {
+                if self.data.int_vars.contains(&vid) {
                     self.infcx.next_int_var()
                 } else {
                     ty
                 }
             }
             ty::Infer(ty::InferTy::FloatVar(vid)) => {
-                if self.float_vars.contains(&vid) {
+                if self.data.float_vars.contains(&vid) {
                     self.infcx.next_float_var()
                 } else {
                     ty
@@ -210,22 +238,22 @@ impl<'a, 'tcx> TypeFolder<TyCtxt<'tcx>> for InferenceFudger<'a, 'tcx> {
 
     fn fold_region(&mut self, r: ty::Region<'tcx>) -> ty::Region<'tcx> {
         if let ty::ReVar(vid) = *r
-            && self.region_vars.0.contains(&vid)
+            && self.data.region_vars.0.contains(&vid)
         {
-            let idx = vid.index() - self.region_vars.0.start.index();
-            let origin = self.region_vars.1[idx];
+            let idx = vid.index() - self.data.region_vars.0.start.index();
+            let origin = self.data.region_vars.1[idx];
             return self.infcx.next_region_var(origin);
         }
         r
     }
 
     fn fold_const(&mut self, ct: ty::Const<'tcx>) -> ty::Const<'tcx> {
         if let ty::ConstKind::Infer(ty::InferConst::Var(vid)) = ct.kind() {
-            if self.const_vars.0.contains(&vid) {
+            if self.data.const_vars.0.contains(&vid) {
                 // This variable was created during the fudging.
                 // Recreate it with a fresh variable here.
-                let idx = vid.index() - self.const_vars.0.start.index();
-                let origin = self.const_vars.1[idx];
+                let idx = vid.index() - self.data.const_vars.0.start.index();
+                let origin = self.data.const_vars.1[idx];
                 self.infcx.next_const_var(ct.ty(), origin)
             } else {
                 ct