yyeet

lcnr · lcnr · commit 4f864db77912 · 2024-03-08T10:16:12.000+01:00
diff --git a/compiler/rustc_infer/src/infer/snapshot/fudge.rs b/compiler/rustc_infer/src/infer/snapshot/fudge.rs
@@ -43,26 +43,7 @@ fn const_vars_since_snapshot<'tcx>(
     )
 }
 
-struct VariableLengths {
-    type_var_len: usize,
-    const_var_len: usize,
-    int_var_len: usize,
-    float_var_len: usize,
-    region_constraints_len: usize,
-}
-
 impl<'tcx> InferCtxt<'tcx> {
-    fn variable_lengths(&self) -> VariableLengths {
-        let mut inner = self.inner.borrow_mut();
-        VariableLengths {
-            type_var_len: inner.type_variables().num_vars(),
-            const_var_len: inner.const_unification_table().len(),
-            int_var_len: inner.int_unification_table().len(),
-            float_var_len: inner.float_unification_table().len(),
-            region_constraints_len: inner.unwrap_region_constraints().num_region_vars(),
-        }
-    }
-
     /// This rather funky routine is used while processing expected
     /// types. What happens here is that we want to propagate a
     /// coercion through the return type of a fn to its
@@ -109,7 +90,7 @@ impl<'tcx> InferCtxt<'tcx> {
         T: TypeFoldable<TyCtxt<'tcx>>,
     {
         let variable_lengths = self.variable_lengths();
-        let (mut fudger, value) = self.probe(|_| {
+        let (mut fudger, value) = self.probe_unchecked(|_| {
             match f() {
                 Ok(value) => {
                     let value = self.resolve_vars_if_possible(value);
@@ -122,21 +103,21 @@ impl<'tcx> InferCtxt<'tcx> {
 
                     let mut inner = self.inner.borrow_mut();
                     let type_vars =
-                        inner.type_variables().vars_since_snapshot(variable_lengths.type_var_len);
+                        inner.type_variables().vars_since_snapshot(variable_lengths.type_vars);
                     let int_vars = vars_since_snapshot(
                         &mut inner.int_unification_table(),
-                        variable_lengths.int_var_len,
+                        variable_lengths.int_vars,
                     );
                     let float_vars = vars_since_snapshot(
                         &mut inner.float_unification_table(),
-                        variable_lengths.float_var_len,
+                        variable_lengths.float_vars,
                     );
                     let region_vars = inner
                         .unwrap_region_constraints()
-                        .vars_since_snapshot(variable_lengths.region_constraints_len);
+                        .vars_since_snapshot(variable_lengths.region_vars);
                     let const_vars = const_vars_since_snapshot(
                         &mut inner.const_unification_table(),
-                        variable_lengths.const_var_len,
+                        variable_lengths.const_vars,
                     );
 
                     let fudger = InferenceFudger {
diff --git a/compiler/rustc_infer/src/infer/snapshot/mod.rs b/compiler/rustc_infer/src/infer/snapshot/mod.rs
@@ -1,7 +1,10 @@
+use std::ops::ControlFlow;
+
 use super::region_constraints::RegionSnapshot;
 use super::InferCtxt;
 use rustc_data_structures::undo_log::UndoLogs;
-use rustc_middle::ty;
+use rustc_middle::ty::{TypeFoldable, TypeSuperVisitable};
+use rustc_middle::ty::{self, Ty, TyCtxt, TypeVisitor};
 
 mod fudge;
 pub(crate) mod undo_log;
@@ -57,27 +60,42 @@ impl<'tcx> InferCtxt<'tcx> {
 
     /// Execute `f` and commit the bindings if closure `f` returns `Ok(_)`.
     #[instrument(skip(self, f), level = "debug")]
-    pub fn commit_if_ok<T, E, F>(&self, f: F) -> Result<T, E>
+    pub fn commit_if_ok<T, E, F, CX>(&self, f: F) -> Result<T, E>
     where
         F: FnOnce(&CombinedSnapshot<'tcx>) -> Result<T, E>,
+        E: NoSnapshotLeaks<'tcx, CX>
     {
+        let no_leaks_data = E::mk_data(self);
         let snapshot = self.start_snapshot();
         let r = f(&snapshot);
         debug!("commit_if_ok() -- r.is_ok() = {}", r.is_ok());
         match r {
-            Ok(_) => {
+            Ok(value) => {
                 self.commit_from(snapshot);
+                Ok(value)
             }
-            Err(_) => {
+            Err(err) => {
                 self.rollback_to(snapshot);
+                Err(E::avoid_leaks(err, no_leaks_data))
             }
         }
-        r
     }
 
     /// Execute `f` then unroll any bindings it creates.
     #[instrument(skip(self, f), level = "debug")]
-    pub fn probe<R, F>(&self, f: F) -> R
+    pub fn probe<R, F, CX>(&self, f: F) -> R
+    where
+        F: FnOnce(&CombinedSnapshot<'tcx>) -> R,
+        R: NoSnapshotLeaks<'tcx, CX>,
+    {
+        let no_leaks_data = R::mk_data(self);
+        let snapshot = self.start_snapshot();
+        let r = f(&snapshot);
+        self.rollback_to(snapshot);
+        R::avoid_leaks(r, no_leaks_data)
+    }
+
+    pub fn probe_unchecked<R, F>(&self, f: F) -> R
     where
         F: FnOnce(&CombinedSnapshot<'tcx>) -> R,
     {
@@ -99,4 +117,128 @@ impl<'tcx> InferCtxt<'tcx> {
     pub fn opaque_types_added_in_snapshot(&self, snapshot: &CombinedSnapshot<'tcx>) -> bool {
         self.inner.borrow().undo_log.opaque_types_in_snapshot(&snapshot.undo_snapshot)
     }
+
+    fn variable_lengths(&self) -> VariableLengths {
+        let mut inner = self.inner.borrow_mut();
+        VariableLengths {
+            type_vars: inner.type_variables().num_vars(),
+            const_vars: inner.const_unification_table().len(),
+            int_vars: inner.int_unification_table().len(),
+            float_vars: inner.float_unification_table().len(),
+            region_vars: inner.unwrap_region_constraints().num_region_vars(),
+        }
+    }
+}
+
+trait NoSnapshotLeaks<'tcx, CX> {
+    fn mk_data(infcx: &InferCtxt<'tcx>) -> CX;
+    fn avoid_leaks(self, data: CX) -> Self;
+}
+
+pub struct CheckLeaks(HasSnapshotLeaksVisitor);
+impl<'tcx, T: TypeFoldable<TyCtxt<'tcx>>> NoSnapshotLeaks<'tcx, CheckLeaks> for T {
+    fn mk_data(infcx: &InferCtxt<'tcx>) -> CheckLeaks {
+        CheckLeaks(HasSnapshotLeaksVisitor {
+            universe: infcx.universe(),
+            variable_lengths: infcx.variable_lengths(),
+        })
+    }
+    fn avoid_leaks(self, CheckLeaks(mut visitor): CheckLeaks) -> Self {
+        if cfg!(debug_assertions) && self.visit_with(&mut visitor).is_break() {
+            bug!("leaking vars from snapshot: {:?}", self)
+        } else {
+            self
+        }
+    }
+}
+
+struct VariableLengths {
+    region_vars: usize,
+    type_vars: usize,
+    int_vars: usize,
+    float_vars: usize,
+    const_vars: usize,
+}
+
+struct HasSnapshotLeaksVisitor {
+    universe: ty::UniverseIndex,
+    variable_lengths: VariableLengths,
+}
+
+fn continue_if(b: bool) -> ControlFlow<()> {
+    if b { ControlFlow::Continue(()) } else { ControlFlow::Continue(()) }
+}
+
+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),
+        }
+    }
 }
diff --git a/compiler/rustc_infer/src/traits/mod.rs b/compiler/rustc_infer/src/traits/mod.rs
@@ -124,6 +124,7 @@ pub type Selection<'tcx> = ImplSource<'tcx, PredicateObligation<'tcx>>;
 pub type ObligationInspector<'tcx> =
     fn(&InferCtxt<'tcx>, &PredicateObligation<'tcx>, Result<Certainty, NoSolution>);
 
+#[derive(Clone, TypeVisitable)]
 pub struct FulfillmentError<'tcx> {
     pub obligation: PredicateObligation<'tcx>,
     pub code: FulfillmentErrorCode<'tcx>,
@@ -133,7 +134,7 @@ pub struct FulfillmentError<'tcx> {
     pub root_obligation: PredicateObligation<'tcx>,
 }
 
-#[derive(Clone)]
+#[derive(Clone, TypeVisitable)]
 pub enum FulfillmentErrorCode<'tcx> {
     /// Inherently impossible to fulfill; this trait is implemented if and only
     /// if it is already implemented.
diff --git a/compiler/rustc_infer/src/traits/project.rs b/compiler/rustc_infer/src/traits/project.rs
@@ -15,7 +15,7 @@ pub use rustc_middle::traits::{EvaluationResult, Reveal};
 pub(crate) type UndoLog<'tcx> =
     snapshot_map::UndoLog<ProjectionCacheKey<'tcx>, ProjectionCacheEntry<'tcx>>;
 
-#[derive(Clone)]
+#[derive(Clone, TypeVisitable)]
 pub struct MismatchedProjectionTypes<'tcx> {
     pub err: ty::error::TypeError<'tcx>,
 }

Original file line number	Diff line number	Diff line change
`@@ -15,7 +15,7 @@ pub use rustc_middle::traits::{EvaluationResult, Reveal};`
`15`	`15`	`pub(crate) type UndoLog<'tcx> =`
`16`	`16`	`snapshot_map::UndoLog<ProjectionCacheKey<'tcx>, ProjectionCacheEntry<'tcx>>;`
`17`	`17`
`18`		`-#[derive(Clone)]`
	`18`	`+#[derive(Clone, TypeVisitable)]`
`19`	`19`	`pub struct MismatchedProjectionTypes<'tcx> {`
`20`	`20`	`pub err: ty::error::TypeError<'tcx>,`
`21`	`21`	`}`