trait solver: Capture binder region constraints while relating

compiler/rustc_next_trait_solver/src/canonical/mod.rs

@@ -14,6 +14,7 @@ use std::iter;
 use canonicalizer::Canonicalizer;
 use rustc_index::IndexVec;
 use rustc_type_ir::inherent::*;
+use rustc_type_ir::region_constraint::RegionConstraint as SolverRegionConstraint;
 use rustc_type_ir::relate::{
     self, Relate, RelateResult, TypeRelation, VarianceDiagInfo, relate_args_invariantly,
 };
@@ -380,7 +381,18 @@ where
 
     #[instrument(skip(self), level = "trace")]
     fn regions(&mut self, a: Region<I>, b: Region<I>) -> RelateResult<I, Region<I>> {
-        self.infcx.equate_regions(a, b, VisibleForLeakCheck::Yes, self.span);
+        if self.cx().assumptions_on_binders() {
+            if a != b {
+                self.infcx.register_solver_region_constraint(
+                    SolverRegionConstraint::RegionOutlives(a, b),
+                );
+                self.infcx.register_solver_region_constraint(
+                    SolverRegionConstraint::RegionOutlives(b, a),
+                );
+            }
+        } else {
+            self.infcx.equate_regions(a, b, VisibleForLeakCheck::Yes, self.span);
+        }
 
         Ok(a)
     }

compiler/rustc_type_ir/src/region_constraint.rs

@@ -460,6 +460,10 @@ pub fn eagerly_handle_placeholders_in_universe<Infcx: InferCtxtLike<Interner = I
 
     let assumptions = infcx.get_placeholder_assumptions(u);
 
+    // Do this before rewriting type outlives constraints: alias/env matching below needs to
+    // see placeholders equated with current-universe region variables in the same `And` branch.
+    let constraint = normalize_equated_region_vars(infcx, constraint, u);
+
     // 1. rewrite type outlives constraints involving things from `u` into either region constraints
     //     involving things from `u` or type outlives constraints not involving things from `u`
     //
@@ -496,6 +500,152 @@ pub fn eagerly_handle_placeholders_in_universe<Infcx: InferCtxtLike<Interner = I
     evaluate_solver_constraint(&constraint)
 }
 
+fn normalize_equated_region_vars<Infcx: InferCtxtLike<Interner = I>, I: Interner>(
+    infcx: &Infcx,
+    constraint: RegionConstraint<I>,
+    u: UniverseIndex,
+) -> RegionConstraint<I> {
+    use RegionConstraint::*;
+
+    match constraint {
+        Ambiguity | RegionOutlives(..) | PlaceholderTyOutlives(..) | AliasTyOutlivesViaEnv(..) => {
+            constraint
+        }
+        Or(constraints) => Or(constraints
+            .into_iter()
+            .map(|constraint| normalize_equated_region_vars(infcx, constraint, u))
+            .collect()),
+        And(constraints) => {
+            let constraint = And(constraints
+                .into_iter()
+                .map(|constraint| normalize_equated_region_vars(infcx, constraint, u))
+                .collect());
+
+            let mut region_outlives = vec![];
+            collect_conjunctive_region_outlives(&constraint, &mut region_outlives);
+            let replacements = compute_equated_region_var_replacements(infcx, &region_outlives, u);
+
+            if replacements.is_empty() {
+                constraint
+            } else {
+                constraint.fold_with(&mut EquatedRegionVarReplacer { cx: infcx.cx(), replacements })
+            }
+        }
+    }
+}
+
+fn compute_equated_region_var_replacements<Infcx: InferCtxtLike<Interner = I>, I: Interner>(
+    infcx: &Infcx,
+    region_outlives: &[(Region<I>, Region<I>)],
+    u: UniverseIndex,
+) -> Vec<(Region<I>, Region<I>)> {
+    compute_equated_region_var_replacements_from(
+        region_outlives,
+        |r| is_current_universe_region_var(infcx, r, u),
+        is_region_var::<I>,
+    )
+}
+
+fn compute_equated_region_var_replacements_from<R>(
+    region_outlives: &[(R, R)],
+    mut is_current_universe_region_var: impl FnMut(R) -> bool,
+    mut is_region_var: impl FnMut(R) -> bool,
+) -> Vec<(R, R)>
+where
+    R: Copy + Eq + std::hash::Hash,
+{
+    let mut equated_regions_builder = TransitiveRelationBuilder::default();
+    let mut has_equated_regions = false;
+    for (r1, r2) in region_outlives.iter().copied() {
+        // Paired outlives constraints represent region equality. Build a transitive relation so
+        // current-universe variables equated through other variables still find a non-var partner.
+        if has_reverse_region_outlives_edge(region_outlives, r1, r2) {
+            equated_regions_builder.add(r1, r2);
+            equated_regions_builder.add(r2, r1);
+            has_equated_regions = true;
+        }
+    }
+
+    if !has_equated_regions {
+        return vec![];
+    }
+
+    let equated_regions = equated_regions_builder.freeze();
+    let mut candidates = IndexSet::new();
+    for (r1, r2) in region_outlives.iter().copied() {
+        if is_current_universe_region_var(r1) {
+            candidates.insert(r1);
+        }
+
+        if is_current_universe_region_var(r2) {
+            candidates.insert(r2);
+        }
+    }
+
+    candidates
+        .into_iter()
+        .filter_map(|candidate| {
+            std::iter::once(candidate)
+                .chain(equated_regions.reachable_from(candidate))
+                .find(|r| !is_region_var(*r))
+                .map(|partner| (candidate, partner))
+        })
+        .collect()
+}
+
+fn has_reverse_region_outlives_edge<R: Eq>(region_outlives: &[(R, R)], r1: R, r2: R) -> bool {
+    region_outlives.iter().any(|(outlives, outlived)| outlives == &r2 && outlived == &r1)
+}
+
+fn collect_conjunctive_region_outlives<I: Interner>(
+    constraint: &RegionConstraint<I>,
+    out: &mut Vec<(Region<I>, Region<I>)>,
+) {
+    use RegionConstraint::*;
+
+    match constraint {
+        RegionOutlives(r1, r2) => out.push((*r1, *r2)),
+        And(constraints) => {
+            for constraint in constraints.iter() {
+                collect_conjunctive_region_outlives(constraint, out);
+            }
+        }
+        Ambiguity | PlaceholderTyOutlives(..) | AliasTyOutlivesViaEnv(..) | Or(..) => {}
+    }
+}
+
+fn is_current_universe_region_var<Infcx: InferCtxtLike<Interner = I>, I: Interner>(
+    infcx: &Infcx,
+    region: Region<I>,
+    u: UniverseIndex,
+) -> bool {
+    is_region_var::<I>(region) && max_universe(infcx, region) == u
+}
+
+fn is_region_var<I: Interner>(region: Region<I>) -> bool {
+    matches!(region.kind(), RegionKind::ReVar(_))
+}
+
+struct EquatedRegionVarReplacer<I: Interner> {
+    cx: I,
+    replacements: Vec<(Region<I>, Region<I>)>,
+}
+
+impl<I: Interner> TypeFolder<I> for EquatedRegionVarReplacer<I> {
+    fn cx(&self) -> I {
+        self.cx
+    }
+
+    fn fold_region(&mut self, r: Region<I>) -> Region<I> {
+        // If a region variable has multiple non-var partners, the remaining folded
+        // constraints still relate those partners, so first-match only affects representation.
+        self.replacements.iter().find_map(|(from, to)| (*from == r).then_some(*to)).unwrap_or(r)
+    }
+}
+
+#[cfg(test)]
+mod tests;
+
 /// Filter our region constraints to not include constraints between region variables from `u` and
 /// other regions as those are always satisfied. This requires some care to handle correctly for example:
 /// `'!a_u1: '?x_u1: '!b_u1` should result in us requiring `'!a_u1: '!b_u1` rather than dropping the two
@@ -657,6 +807,11 @@ fn pull_region_outlives_constraints_out_of_universe<
             constraint
         }
         RegionOutlives(region_1, region_2) => {
+            if region_1 == region_2 {
+                // Reflexive constraints are always satisfied, even if the region is from `u`.
+                return RegionConstraint::new_true();
+            }
+
             let region_1_u = max_universe(infcx, region_1);
             let region_2_u = max_universe(infcx, region_2);
 

compiler/rustc_type_ir/src/region_constraint/tests.rs

@@ -0,0 +1,19 @@
+use super::compute_equated_region_var_replacements_from;
+
+#[test]
+fn equated_region_var_replacements_follow_transitive_region_var_chains() {
+    const REVAR_1: u8 = 1;
+    const REVAR_2: u8 = 2;
+    const PLACEHOLDER: u8 = 3;
+
+    let region_outlives =
+        [(REVAR_1, REVAR_2), (REVAR_2, REVAR_1), (REVAR_2, PLACEHOLDER), (PLACEHOLDER, REVAR_2)];
+
+    let replacements = compute_equated_region_var_replacements_from(
+        &region_outlives,
+        |r| matches!(r, REVAR_1 | REVAR_2),
+        |r| matches!(r, REVAR_1 | REVAR_2),
+    );
+
+    assert_eq!(replacements, vec![(REVAR_1, PLACEHOLDER), (REVAR_2, PLACEHOLDER)]);
+}

compiler/rustc_type_ir/src/relate/solver_relating.rs

@@ -2,6 +2,7 @@ use tracing::{debug, instrument};
 
 use self::combine::{PredicateEmittingRelation, super_combine_consts, super_combine_tys};
 use crate::data_structures::DelayedSet;
+use crate::region_constraint::RegionConstraint;
 use crate::relate::combine::combine_ty_args;
 pub use crate::relate::*;
 use crate::solve::{Goal, VisibleForLeakCheck};
@@ -218,14 +219,40 @@ where
 
     #[instrument(skip(self), level = "trace")]
     fn regions(&mut self, a: Region<I>, b: Region<I>) -> RelateResult<I, Region<I>> {
-        match self.ambient_variance {
-            // Subtype(&'a u8, &'b u8) => Outlives('a: 'b) => SubRegion('b, 'a)
-            ty::Covariant => self.infcx.sub_regions(b, a, VisibleForLeakCheck::Yes, self.span),
-            // Suptype(&'a u8, &'b u8) => Outlives('b: 'a) => SubRegion('a, 'b)
-            ty::Contravariant => self.infcx.sub_regions(a, b, VisibleForLeakCheck::Yes, self.span),
-            ty::Invariant => self.infcx.equate_regions(a, b, VisibleForLeakCheck::Yes, self.span),
-            ty::Bivariant => {
-                unreachable!("Expected bivariance to be handled in relate_with_variance")
+        if self.cx().assumptions_on_binders() {
+            if a == b {
+                return Ok(a);
+            }
+
+            match self.ambient_variance {
+                ty::Covariant => self
+                    .infcx
+                    .register_solver_region_constraint(RegionConstraint::RegionOutlives(a, b)),
+                ty::Contravariant => self
+                    .infcx
+                    .register_solver_region_constraint(RegionConstraint::RegionOutlives(b, a)),
+                ty::Invariant => {
+                    self.infcx
+                        .register_solver_region_constraint(RegionConstraint::RegionOutlives(a, b));
+                    self.infcx
+                        .register_solver_region_constraint(RegionConstraint::RegionOutlives(b, a));
+                }
+                ty::Bivariant => {
+                    unreachable!("Expected bivariance to be handled in relate_with_variance")
+                }
+            }
+        } else {
+            match self.ambient_variance {
+                ty::Covariant => self.infcx.sub_regions(b, a, VisibleForLeakCheck::Yes, self.span),
+                ty::Contravariant => {
+                    self.infcx.sub_regions(a, b, VisibleForLeakCheck::Yes, self.span)
+                }
+                ty::Invariant => {
+                    self.infcx.equate_regions(a, b, VisibleForLeakCheck::Yes, self.span)
+                }
+                ty::Bivariant => {
+                    unreachable!("Expected bivariance to be handled in relate_with_variance")
+                }
             }
         }
 

tests/ui/assumptions_on_binders/principal-upcast-region-eq-issue-157859.rs

@@ -0,0 +1,19 @@
+//@compile-flags: -Zassumptions-on-binders  -Znext-solver=globally
+//@ dont-require-annotations: ERROR
+
+trait Super<U> {
+    fn a(&self) {
+        let a: &dyn Sub = &();
+        let b: &dyn Super<for<'a> fn(&'a ())> = a;
+    }
+}
+
+impl<T> Super<T> for () {}
+
+trait Sub: Super<fn(&'static ())> {}
+
+impl Sub for () {}
+
+fn main() {
+    let a: &dyn Sub = &();
+}

tests/ui/assumptions_on_binders/principal-upcast-region-eq-issue-157859.stderr

@@ -0,0 +1,13 @@
+error[E0277]: the trait bound `&dyn Sub: CoerceUnsized<&dyn Super<for<'a> fn(&'a ())>>` is not satisfied
+  --> $DIR/principal-upcast-region-eq-issue-157859.rs:7:49
+   |
+LL |         let b: &dyn Super<for<'a> fn(&'a ())> = a;
+   |                                                 ^ the nightly-only, unstable trait `Unsize<dyn Super<for<'a> fn(&'a ())>>` is not implemented for `dyn Sub`
+   |
+   = note: all implementations of `Unsize` are provided automatically by the compiler, see <https://doc.rust-lang.org/stable/std/marker/trait.Unsize.html> for more information
+   = note: required for `&dyn Sub` to implement `CoerceUnsized<&dyn Super<for<'a> fn(&'a ())>>`
+   = note: required for the cast from `&dyn Sub` to `&dyn Super<for<'a> fn(&'a ())>`
+
+error: aborting due to 1 previous error
+
+For more information about this error, try `rustc --explain E0277`.

tests/ui/assumptions_on_binders/trait-upcast-projection-region-eq.rs

@@ -0,0 +1,23 @@
+//@ compile-flags: -Zassumptions-on-binders -Znext-solver=globally
+
+trait Super {
+    type Assoc;
+
+    fn a(&self) {
+        let a: &dyn Sub = &();
+        let b: &dyn Super<Assoc = for<'a> fn(&'a ())> = a;
+        //~^ ERROR the trait bound `&dyn Sub: CoerceUnsized<&dyn Super<Assoc = for<'a> fn(&'a ())>>` is not satisfied
+    }
+}
+
+impl Super for () {
+    type Assoc = fn(&'static ());
+}
+
+trait Sub: Super<Assoc = fn(&'static ())> {}
+
+impl Sub for () {}
+
+fn main() {
+    let a: &dyn Sub = &();
+}

tests/ui/assumptions_on_binders/trait-upcast-projection-region-eq.stderr

@@ -0,0 +1,13 @@
+error[E0277]: the trait bound `&dyn Sub: CoerceUnsized<&dyn Super<Assoc = for<'a> fn(&'a ())>>` is not satisfied
+  --> $DIR/trait-upcast-projection-region-eq.rs:8:57
+   |
+LL |         let b: &dyn Super<Assoc = for<'a> fn(&'a ())> = a;
+   |                                                         ^ the nightly-only, unstable trait `Unsize<dyn Super<Assoc = for<'a> fn(&'a ())>>` is not implemented for `dyn Sub`
+   |
+   = note: all implementations of `Unsize` are provided automatically by the compiler, see <https://doc.rust-lang.org/stable/std/marker/trait.Unsize.html> for more information
+   = note: required for `&dyn Sub` to implement `CoerceUnsized<&dyn Super<Assoc = for<'a> fn(&'a ())>>`
+   = note: required for the cast from `&dyn Sub` to `&dyn Super<Assoc = for<'a> fn(&'a ())>`
+
+error: aborting due to 1 previous error
+
+For more information about this error, try `rustc --explain E0277`.