Support leak check for new style region constraint

compiler/rustc_hir_analysis/src/collect.rs

@@ -36,7 +36,7 @@ use rustc_middle::query::Providers;
 use rustc_middle::ty::util::{Discr, IntTypeExt};
 use rustc_middle::ty::{
     self, AdtKind, Const, IsSuggestable, Ty, TyCtxt, TypeVisitableExt, TypingMode, Unnormalized,
-    fold_regions,
+    VisibleForLeakCheck, fold_regions,
 };
 use rustc_span::def_id::LocalModId;
 use rustc_span::{DUMMY_SP, Ident, Span, Symbol, bug, kw, span_bug, sym};
@@ -444,7 +444,10 @@ impl<'tcx> ItemCtxt<'tcx> {
                 let lhs = self.lowerer().lower_lifetime(lhs, RegionInferReason::RegionPredicate);
                 let rhs = self.lowerer().lower_lifetime(rhs, RegionInferReason::RegionPredicate);
                 SolverRegionConstraint::new_leaf(LeafRegionConstraint::RegionOutlives(
-                    lhs, rhs, span,
+                    lhs,
+                    rhs,
+                    VisibleForLeakCheck::Unreachable,
+                    span,
                 ))
             }
             hir::TestBinderConstraint::PlaceholderOutlives { lhs, rhs } => {

compiler/rustc_infer/src/infer/outlives/obligations.rs

@@ -295,7 +295,7 @@ impl<'tcx> InferCtxt<'tcx> {
                         )
                         .emit();
                 }
-                RegionOutlives(a, b, span) => {
+                RegionOutlives(a, b, _, span) => {
                     let origin = SubregionOrigin::SolverRegionConstraint(span);
                     let category = origin.to_constraint_category();
                     conversion.push_sub_region_constraint(

compiler/rustc_infer/src/infer/region_constraints/leak_check.rs

@@ -410,21 +410,58 @@ impl<'tcx> MiniGraph<'tcx> {
                     &AddCombination(..) | &AddVar(..) => {}
                     // FIXME(-Zassumptions-on-binders): we probably need to handle this in some manner
                     &OverwriteSolverRegionConstraint { .. } => {}
+                    &AddSolverLeafRegionConstraint(i) => {
+                        if let SolverLeafRegionConstraint::RegionOutlives(
+                            sup,
+                            sub,
+                            visible_for_leak_check,
+                            _,
+                        ) = region_constraints.solver_data().leafs()[i]
+                        {
+                            match visible_for_leak_check {
+                                ty::VisibleForLeakCheck::Yes => each_edge(sub, sup),
+                                ty::VisibleForLeakCheck::No => {}
+                                ty::VisibleForLeakCheck::Unreachable => unreachable!(),
+                            }
+                        }
+                    }
                 }
             }
         } else {
-            region_constraints
-                .legacy_data()
-                .constraints
-                .iter()
-                .flat_map(|(c, _)| c.iter_outlives())
-                .for_each(|Constraint { kind: _, sub, sup, visible_for_leak_check }| {
-                    match visible_for_leak_check {
-                        ty::VisibleForLeakCheck::Yes => each_edge(sub, sup),
-                        ty::VisibleForLeakCheck::No => {}
-                        ty::VisibleForLeakCheck::Unreachable => unreachable!(),
+            let assumptions_on_binders_enabled =
+                rustc_middle::ty::tls::with(|tcx| tcx.assumptions_on_binders());
+            if assumptions_on_binders_enabled {
+                region_constraints.solver_data().leafs().iter().for_each(|leaf| {
+                    if let SolverLeafRegionConstraint::RegionOutlives(
+                        sup,
+                        sub,
+                        visible_for_leak_check,
+                        _,
+                    ) = leaf
+                    {
+                        match visible_for_leak_check {
+                            ty::VisibleForLeakCheck::Yes => each_edge(*sub, *sup),
+                            ty::VisibleForLeakCheck::No => {}
+                            ty::VisibleForLeakCheck::Unreachable => unreachable!(),
+                        }
+                    } else {
+                        unreachable!()
                     }
                 })
+            } else {
+                region_constraints
+                    .legacy_data()
+                    .constraints
+                    .iter()
+                    .flat_map(|(c, _)| c.iter_outlives())
+                    .for_each(|Constraint { kind: _, sub, sup, visible_for_leak_check }| {
+                        match visible_for_leak_check {
+                            ty::VisibleForLeakCheck::Yes => each_edge(sub, sup),
+                            ty::VisibleForLeakCheck::No => {}
+                            ty::VisibleForLeakCheck::Unreachable => unreachable!(),
+                        }
+                    })
+            }
         }
     }
 

compiler/rustc_infer/src/infer/region_constraints/mod.rs

@@ -86,6 +86,9 @@ pub struct RegionConstraintData<'tcx> {
 pub type SolverRegionConstraint<'tcx> =
     rustc_type_ir::region_constraint::RegionConstraint<TyCtxt<'tcx>, Span>;
 
+pub type SolverLeafRegionConstraint<'tcx> =
+    rustc_type_ir::region_constraint::LeafRegionConstraint<TyCtxt<'tcx>, Span>;
+
 #[derive(Clone, Debug, Default)]
 pub struct SolverRegionConstraintData<'tcx>(Option<SolverRegionConstraint<'tcx>>);
 
@@ -101,6 +104,17 @@ impl<'tcx> SolverRegionConstraintData<'tcx> {
     pub(crate) fn overwrite(&mut self, constraint: SolverRegionConstraint<'tcx>) {
         self.0 = Some(constraint);
     }
+
+    pub(crate) fn leafs_len(&self) -> usize {
+        match &self.0 {
+            Some(v) => v.and_constraint.0.len(),
+            None => 0,
+        }
+    }
+
+    pub(crate) fn leafs(&self) -> &[SolverLeafRegionConstraint<'tcx>] {
+        &self.0.as_ref().unwrap().and_constraint.0
+    }
 }
 
 /// Represents a constraint that influences the inference process.
@@ -350,6 +364,9 @@ pub(crate) enum UndoLog<'tcx> {
 
     /// Overwrite the `-Zassumptions-on-binders` storage
     OverwriteSolverRegionConstraint { old_constraint: Box<SolverRegionConstraint<'tcx>> },
+
+    /// We added the given leaf solver constraint.
+    AddSolverLeafRegionConstraint(usize),
 }
 
 #[derive(Copy, Clone, PartialEq)]
@@ -477,10 +494,22 @@ impl<'tcx> RegionConstraintCollector<'_, 'tcx> {
             "solver region constraint has escaping bound vars, which is indicative of a bug in how constraints are handled: {constraint:?}",
         );
 
+        // FIXME(-Zassumptions-on-binders): We don't actually register leaf constraints in
+        // `or_constraint`. This is a problem for complete incremental undo log.
+        //
+        // Currently it's only used for snapshot based leak check. Therefore, we can get away with
+        // ignoring `Or` constraint since region constraints from type outlives aren't visible for
+        // leak check.
+        for leaf in constraint.and_constraint.0.clone().into_iter() {
+            debug!("RegionConstraintCollector: add_solver_leaf_region_constraint({:?})", leaf);
+            let index = self.storage.solver_data.leafs_len();
+            self.undo_log.push(AddSolverLeafRegionConstraint(index));
+        }
+
         let old_constraint = self.storage.solver_data.get_constraint();
         let new_constraint = rustc_type_ir::region_constraint::RegionConstraint::build_and(
-            constraint,
             old_constraint.clone(),
+            constraint,
         );
 
         // FIXME(-Zassumptions-on-binders): This is pretty bad for perf, we don't make incremental
@@ -943,6 +972,10 @@ impl<'tcx> Rollback<UndoLog<'tcx>> for RegionConstraintStorage<'tcx> {
             OverwriteSolverRegionConstraint { old_constraint } => {
                 self.solver_data.overwrite(*old_constraint);
             }
+            AddSolverLeafRegionConstraint(_) => {
+                // Nothing to do.
+                // The storage is restored via `OverwriteSolverRegionConstraint`
+            }
         }
     }
 }

compiler/rustc_next_trait_solver/src/solve/eval_ctxt/solver_region_constraints.rs

@@ -16,7 +16,7 @@ use rustc_type_ir::{
 use tracing::{debug, instrument};
 
 use crate::delegate::SolverDelegate;
-use crate::solve::{Certainty, EvalCtxt, Goal, NoSolution};
+use crate::solve::{Certainty, EvalCtxt, Goal, NoSolution, VisibleForLeakCheck};
 
 /// Logic for `-Zassumptions-on-binders` stuff
 impl<'a, D, I> EvalCtxt<'a, D>
@@ -168,7 +168,7 @@ where
         use Component::*;
         use LeafRegionConstraint::*;
         match c {
-            Region(c_r) => Or::new_leaf(RegionOutlives(*c_r, r, ())),
+            Region(c_r) => Or::new_leaf(RegionOutlives(*c_r, r, VisibleForLeakCheck::No, ())),
             Placeholder(p) => {
                 Or::new_leaf(PlaceholderTyOutlives(Ty::new_placeholder(self.cx(), *p), r, ()))
             }
@@ -194,7 +194,7 @@ where
 
         let item_bounds =
             rustc_type_ir::outlives::declared_bounds_from_definition(self.cx(), alias)
-                .map(|bound| And::new([RegionOutlives(bound, r, ())]));
+                .map(|bound| And::new([RegionOutlives(bound, r, VisibleForLeakCheck::No, ())]));
         let item_bound_outlives = Or::new(item_bounds);
 
         let where_clause_outlives =

compiler/rustc_next_trait_solver/src/solve/mod.rs

@@ -157,8 +157,12 @@ where
         if self.cx().assumptions_on_binders() {
             use rustc_type_ir::region_constraint::{LeafRegionConstraint, RegionConstraint};
 
-            let constraint =
-                RegionConstraint::new_leaf(LeafRegionConstraint::RegionOutlives(a, b, ()));
+            let constraint = RegionConstraint::new_leaf(LeafRegionConstraint::RegionOutlives(
+                a,
+                b,
+                VisibleForLeakCheck::Yes,
+                (),
+            ));
             self.register_solver_region_constraint(constraint);
         } else {
             self.register_region_outlives(a, b, VisibleForLeakCheck::Yes);