Auto merge of #146329 - lcnr:opaque-type-infer-alias-candidates, r=BoxyUwU

consider item bounds for non-yet-defined opaque types

Based on rust-lang/rust#140405.

fixes https://github.com/rust-lang/trait-system-refactor-initiative/issues/182
fixes https://github.com/rust-lang/trait-system-refactor-initiative/issues/196
fixes https://github.com/rust-lang/trait-system-refactor-initiative/issues/205

there's some jank here, see https://github.com/rust-lang/trait-system-refactor-initiative/issues/229

## Design

If the self type is an inference variable which has been sub-unified with am opaque type, we need to incompletely guide inference to avoid breakage.

In this case, we
- look at the item bounds of all sub-unified opaque types, and
- blanket impls which do not constrain the self type

Even if there are applicable candidates, we always force their certainty to be `Maybe`, so they will always have to be reproven once we've constrained the inference variable.

This is a bit iffy, see the added tests.

r? `@BoxyUwU`

8 files changed, 207 insertions, 38 deletions

diff --git a/compiler/rustc_middle/src/ty/context.rs b/compiler/rustc_middle/src/ty/context.rs
index 218ac2cfbc1..8ea767dccd3 100644
--- a/compiler/rustc_middle/src/ty/context.rs
+++ b/compiler/rustc_middle/src/ty/context.rs
@@ -651,7 +651,11 @@ impl<'tcx> Interner for TyCtxt<'tcx> {
             | ty::Bound(_, _) => bug!("unexpected self type: {self_ty}"),
         }
 
-        let trait_impls = tcx.trait_impls_of(trait_def_id);
+        #[allow(rustc::usage_of_type_ir_traits)]
+        self.for_each_blanket_impl(trait_def_id, f)
+    }
+    fn for_each_blanket_impl(self, trait_def_id: DefId, mut f: impl FnMut(DefId)) {
+        let trait_impls = self.trait_impls_of(trait_def_id);
         for &impl_def_id in trait_impls.blanket_impls() {
             f(impl_def_id);
         }
diff --git a/compiler/rustc_next_trait_solver/src/solve/assembly/mod.rs b/compiler/rustc_next_trait_solver/src/solve/assembly/mod.rs
index e3cf0330b14..fb777496e31 100644
--- a/compiler/rustc_next_trait_solver/src/solve/assembly/mod.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/assembly/mod.rs
@@ -11,8 +11,9 @@ use rustc_type_ir::lang_items::SolverTraitLangItem;
 use rustc_type_ir::search_graph::CandidateHeadUsages;
 use rustc_type_ir::solve::SizedTraitKind;
 use rustc_type_ir::{
-    self as ty, Interner, TypeFlags, TypeFoldable, TypeSuperVisitable, TypeVisitable,
-    TypeVisitableExt as _, TypeVisitor, TypingMode, Upcast as _, elaborate,
+    self as ty, Interner, TypeFlags, TypeFoldable, TypeFolder, TypeSuperFoldable,
+    TypeSuperVisitable, TypeVisitable, TypeVisitableExt, TypeVisitor, TypingMode, Upcast,
+    elaborate,
 };
 use tracing::{debug, instrument};
 
@@ -187,6 +188,7 @@ where
         ecx: &mut EvalCtxt<'_, D>,
         goal: Goal<I, Self>,
         impl_def_id: I::ImplId,
+        then: impl FnOnce(&mut EvalCtxt<'_, D>, Certainty) -> QueryResult<I>,
     ) -> Result<Candidate<I>, NoSolution>;
 
     /// If the predicate contained an error, we want to avoid emitting unnecessary trait
@@ -365,6 +367,15 @@ pub(super) enum AssembleCandidatesFrom {
     EnvAndBounds,
 }
 
+impl AssembleCandidatesFrom {
+    fn should_assemble_impl_candidates(&self) -> bool {
+        match self {
+            AssembleCandidatesFrom::All => true,
+            AssembleCandidatesFrom::EnvAndBounds => false,
+        }
+    }
+}
+
 /// This is currently used to track the [CandidateHeadUsages] of all failed `ParamEnv`
 /// candidates. This is then used to ignore their head usages in case there's another
 /// always applicable `ParamEnv` candidate. Look at how `param_env_head_usages` is
@@ -397,14 +408,15 @@ where
             return (candidates, failed_candidate_info);
         };
 
+        let goal: Goal<I, G> = goal
+            .with(self.cx(), goal.predicate.with_replaced_self_ty(self.cx(), normalized_self_ty));
+
         if normalized_self_ty.is_ty_var() {
             debug!("self type has been normalized to infer");
-            candidates.extend(self.forced_ambiguity(MaybeCause::Ambiguity));
+            self.try_assemble_bounds_via_registered_opaques(goal, assemble_from, &mut candidates);
             return (candidates, failed_candidate_info);
         }
 
-        let goal: Goal<I, G> = goal
-            .with(self.cx(), goal.predicate.with_replaced_self_ty(self.cx(), normalized_self_ty));
         // Vars that show up in the rest of the goal substs may have been constrained by
         // normalizing the self type as well, since type variables are not uniquified.
         let goal = self.resolve_vars_if_possible(goal);
@@ -484,8 +496,9 @@ where
                 if cx.impl_is_default(impl_def_id) {
                     return;
                 }
-
-                match G::consider_impl_candidate(self, goal, impl_def_id) {
+                match G::consider_impl_candidate(self, goal, impl_def_id, |ecx, certainty| {
+                    ecx.evaluate_added_goals_and_make_canonical_response(certainty)
+                }) {
                     Ok(candidate) => candidates.push(candidate),
                     Err(NoSolution) => (),
                 }
@@ -943,6 +956,116 @@ where
         }
     }
 
+    /// If the self type is the hidden type of an opaque, try to assemble
+    /// candidates for it by consider its item bounds and by using blanket
+    /// impls. This is used to incompletely guide type inference when handling
+    /// non-defining uses in the defining scope.
+    ///
+    /// We otherwise just fail fail with ambiguity. Even if we're using an
+    /// opaque type item bound or a blank impls, we still force its certainty
+    /// to be `Maybe` so that we properly prove this goal later.
+    ///
+    /// See <https://github.com/rust-lang/trait-system-refactor-initiative/issues/182>
+    /// for why this is necessary.
+    fn try_assemble_bounds_via_registered_opaques<G: GoalKind<D>>(
+        &mut self,
+        goal: Goal<I, G>,
+        assemble_from: AssembleCandidatesFrom,
+        candidates: &mut Vec<Candidate<I>>,
+    ) {
+        let self_ty = goal.predicate.self_ty();
+        // If the self type is sub unified with any opaque type, we
+        // also look at blanket impls for it.
+        let mut assemble_blanket_impls = false;
+        for alias_ty in self.opaques_with_sub_unified_hidden_type(self_ty) {
+            assemble_blanket_impls = true;
+            debug!("self ty is sub unified with {alias_ty:?}");
+
+            struct ReplaceOpaque<I: Interner> {
+                cx: I,
+                alias_ty: ty::AliasTy<I>,
+                self_ty: I::Ty,
+            }
+            impl<I: Interner> TypeFolder<I> for ReplaceOpaque<I> {
+                fn cx(&self) -> I {
+                    self.cx
+                }
+                fn fold_ty(&mut self, ty: I::Ty) -> I::Ty {
+                    if let ty::Alias(ty::Opaque, alias_ty) = ty.kind() {
+                        if alias_ty == self.alias_ty {
+                            return self.self_ty;
+                        }
+                    }
+                    ty.super_fold_with(self)
+                }
+            }
+
+            // We look at all item-bounds of the opaque, replacing the
+            // opaque with the current self type before considering
+            // them as a candidate. Imagine e've got `?x: Trait<?y>`
+            // and `?x` has been sub-unified with the hidden type of
+            // `impl Trait<u32>`, We take the item bound `opaque: Trait<u32>`
+            // and replace all occurrences of `opaque` with `?x`. This results
+            // in a `?x: Trait<u32>` alias-bound candidate.
+            for item_bound in self
+                .cx()
+                .item_self_bounds(alias_ty.def_id)
+                .iter_instantiated(self.cx(), alias_ty.args)
+            {
+                let assumption =
+                    item_bound.fold_with(&mut ReplaceOpaque { cx: self.cx(), alias_ty, self_ty });
+                candidates.extend(G::probe_and_match_goal_against_assumption(
+                    self,
+                    CandidateSource::AliasBound,
+                    goal,
+                    assumption,
+                    |ecx| {
+                        // We want to reprove this goal once we've inferred the
+                        // hidden type, so we force the certainty to `Maybe`.
+                        ecx.evaluate_added_goals_and_make_canonical_response(Certainty::AMBIGUOUS)
+                    },
+                ));
+            }
+        }
+
+        // We also need to consider blanket impls for not-yet-defined opaque types.
+        //
+        // See tests/ui/impl-trait/non-defining-uses/use-blanket-impl.rs for an example.
+        if assemble_blanket_impls && assemble_from.should_assemble_impl_candidates() {
+            let cx = self.cx();
+            cx.for_each_blanket_impl(goal.predicate.trait_def_id(cx), |impl_def_id| {
+                // For every `default impl`, there's always a non-default `impl`
+                // that will *also* apply. There's no reason to register a candidate
+                // for this impl, since it is *not* proof that the trait goal holds.
+                if cx.impl_is_default(impl_def_id) {
+                    return;
+                }
+
+                match G::consider_impl_candidate(self, goal, impl_def_id, |ecx, certainty| {
+                    if ecx.shallow_resolve(self_ty).is_ty_var() {
+                        // We force the certainty of impl candidates to be `Maybe`.
+                        let certainty = certainty.and(Certainty::AMBIGUOUS);
+                        ecx.evaluate_added_goals_and_make_canonical_response(certainty)
+                    } else {
+                        // We don't want to use impls if they constrain the opaque.
+                        //
+                        // FIXME(trait-system-refactor-initiative#229): This isn't
+                        // perfect yet as it still allows us to incorrectly constrain
+                        // other inference variables.
+                        Err(NoSolution)
+                    }
+                }) {
+                    Ok(candidate) => candidates.push(candidate),
+                    Err(NoSolution) => (),
+                }
+            });
+        }
+
+        if candidates.is_empty() {
+            candidates.extend(self.forced_ambiguity(MaybeCause::Ambiguity));
+        }
+    }
+
     /// Assemble and merge candidates for goals which are related to an underlying trait
     /// goal. Right now, this is normalizes-to and host effect goals.
     ///
diff --git a/compiler/rustc_next_trait_solver/src/solve/effect_goals.rs b/compiler/rustc_next_trait_solver/src/solve/effect_goals.rs
index 6646857a136..cb72c1cd92b 100644
--- a/compiler/rustc_next_trait_solver/src/solve/effect_goals.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/effect_goals.rs
@@ -124,6 +124,7 @@ where
         ecx: &mut EvalCtxt<'_, D>,
         goal: Goal<I, Self>,
         impl_def_id: I::ImplId,
+        then: impl FnOnce(&mut EvalCtxt<'_, D>, Certainty) -> QueryResult<I>,
     ) -> Result<Candidate<I>, NoSolution> {
         let cx = ecx.cx();
 
@@ -175,7 +176,7 @@ where
                 });
             ecx.add_goals(GoalSource::ImplWhereBound, const_conditions);
 
-            ecx.evaluate_added_goals_and_make_canonical_response(certainty)
+            then(ecx, certainty)
         })
     }
 
diff --git a/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs b/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs
index 5cd597d4fe6..3e3a5246f3d 100644
--- a/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs
@@ -407,20 +407,21 @@ where
         // If we have run this goal before, and it was stalled, check that any of the goal's
         // args have changed. Otherwise, we don't need to re-run the goal because it'll remain
         // stalled, since it'll canonicalize the same way and evaluation is pure.
-        if let Some(stalled_on) = stalled_on
-            && !stalled_on.stalled_vars.iter().any(|value| self.delegate.is_changed_arg(*value))
-            && !self
-                .delegate
-                .opaque_types_storage_num_entries()
-                .needs_reevaluation(stalled_on.num_opaques)
+        if let Some(GoalStalledOn { num_opaques, ref stalled_vars, ref sub_roots, stalled_cause }) =
+            stalled_on
+            && !stalled_vars.iter().any(|value| self.delegate.is_changed_arg(*value))
+            && !sub_roots
+                .iter()
+                .any(|&vid| self.delegate.sub_unification_table_root_var(vid) != vid)
+            && !self.delegate.opaque_types_storage_num_entries().needs_reevaluation(num_opaques)
         {
             return Ok((
                 NestedNormalizationGoals::empty(),
                 GoalEvaluation {
                     goal,
-                    certainty: Certainty::Maybe(stalled_on.stalled_cause),
+                    certainty: Certainty::Maybe(stalled_cause),
                     has_changed: HasChanged::No,
-                    stalled_on: Some(stalled_on),
+                    stalled_on,
                 },
             ));
         }
@@ -476,16 +477,6 @@ where
                 HasChanged::No => {
                     let mut stalled_vars = orig_values;
 
-                    // Remove the canonicalized universal vars, since we only care about stalled existentials.
-                    stalled_vars.retain(|arg| match arg.kind() {
-                        ty::GenericArgKind::Type(ty) => matches!(ty.kind(), ty::Infer(_)),
-                        ty::GenericArgKind::Const(ct) => {
-                            matches!(ct.kind(), ty::ConstKind::Infer(_))
-                        }
-                        // Lifetimes can never stall goals.
-                        ty::GenericArgKind::Lifetime(_) => false,
-                    });
-
                     // Remove the unconstrained RHS arg, which is expected to have changed.
                     if let Some(normalizes_to) = goal.predicate.as_normalizes_to() {
                         let normalizes_to = normalizes_to.skip_binder();
@@ -497,6 +488,27 @@ where
                         stalled_vars.swap_remove(idx);
                     }
 
+                    // Remove the canonicalized universal vars, since we only care about stalled existentials.
+                    let mut sub_roots = Vec::new();
+                    stalled_vars.retain(|arg| match arg.kind() {
+                        // Lifetimes can never stall goals.
+                        ty::GenericArgKind::Lifetime(_) => false,
+                        ty::GenericArgKind::Type(ty) => match ty.kind() {
+                            ty::Infer(ty::TyVar(vid)) => {
+                                sub_roots.push(self.delegate.sub_unification_table_root_var(vid));
+                                true
+                            }
+                            ty::Infer(_) => true,
+                            ty::Param(_) | ty::Placeholder(_) => false,
+                            _ => unreachable!("unexpected orig_value: {ty:?}"),
+                        },
+                        ty::GenericArgKind::Const(ct) => match ct.kind() {
+                            ty::ConstKind::Infer(_) => true,
+                            ty::ConstKind::Param(_) | ty::ConstKind::Placeholder(_) => false,
+                            _ => unreachable!("unexpected orig_value: {ct:?}"),
+                        },
+                    });
+
                     Some(GoalStalledOn {
                         num_opaques: canonical_goal
                             .canonical
@@ -505,6 +517,7 @@ where
                             .opaque_types
                             .len(),
                         stalled_vars,
+                        sub_roots,
                         stalled_cause,
                     })
                 }
@@ -1047,6 +1060,10 @@ where
         self.delegate.resolve_vars_if_possible(value)
     }
 
+    pub(super) fn shallow_resolve(&self, ty: I::Ty) -> I::Ty {
+        self.delegate.shallow_resolve(ty)
+    }
+
     pub(super) fn eager_resolve_region(&self, r: I::Region) -> I::Region {
         if let ty::ReVar(vid) = r.kind() {
             self.delegate.opportunistic_resolve_lt_var(vid)
@@ -1163,6 +1180,33 @@ where
     ) -> bool {
         may_use_unstable_feature(&**self.delegate, param_env, symbol)
     }
+
+    pub(crate) fn opaques_with_sub_unified_hidden_type(
+        &self,
+        self_ty: I::Ty,
+    ) -> impl Iterator<Item = ty::AliasTy<I>> + use<'a, D, I> {
+        let delegate = self.delegate;
+        delegate
+            .clone_opaque_types_lookup_table()
+            .into_iter()
+            .chain(delegate.clone_duplicate_opaque_types())
+            .filter_map(move |(key, hidden_ty)| {
+                if let ty::Infer(ty::TyVar(self_vid)) = self_ty.kind() {
+                    if let ty::Infer(ty::TyVar(hidden_vid)) = hidden_ty.kind() {
+                        if delegate.sub_unification_table_root_var(self_vid)
+                            == delegate.sub_unification_table_root_var(hidden_vid)
+                        {
+                            return Some(ty::AliasTy::new_from_args(
+                                delegate.cx(),
+                                key.def_id.into(),
+                                key.args,
+                            ));
+                        }
+                    }
+                }
+                None
+            })
+    }
 }
 
 /// Eagerly replace aliases with inference variables, emitting `AliasRelate`
diff --git a/compiler/rustc_next_trait_solver/src/solve/mod.rs b/compiler/rustc_next_trait_solver/src/solve/mod.rs
index db3460c4ff6..cd27c9c26c1 100644
--- a/compiler/rustc_next_trait_solver/src/solve/mod.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/mod.rs
@@ -24,7 +24,7 @@ mod trait_goals;
 use derive_where::derive_where;
 use rustc_type_ir::inherent::*;
 pub use rustc_type_ir::solve::*;
-use rustc_type_ir::{self as ty, Interner, TypingMode};
+use rustc_type_ir::{self as ty, Interner, TyVid, TypingMode};
 use tracing::instrument;
 
 pub use self::eval_ctxt::{
@@ -418,11 +418,6 @@ pub struct GoalEvaluation<I: Interner> {
     pub has_changed: HasChanged,
     /// If the [`Certainty`] was `Maybe`, then keep track of whether the goal has changed
     /// before rerunning it.
-    ///
-    /// We knowingly ignore the `sub_root` of our inference variables here. This means we
-    /// may not reevaluate a goal even though a change to the `sub_root` could cause a goal
-    /// to make progress. Tracking them adds additional complexity for an incredibly minor
-    /// type inference improvement. We could look into properly handling this in the future.
     pub stalled_on: Option<GoalStalledOn<I>>,
 }
 
@@ -431,6 +426,7 @@ pub struct GoalEvaluation<I: Interner> {
 pub struct GoalStalledOn<I: Interner> {
     pub num_opaques: usize,
     pub stalled_vars: Vec<I::GenericArg>,
+    pub sub_roots: Vec<TyVid>,
     /// The cause that will be returned on subsequent evaluations if this goal remains stalled.
     pub stalled_cause: MaybeCause,
 }
diff --git a/compiler/rustc_next_trait_solver/src/solve/normalizes_to/mod.rs b/compiler/rustc_next_trait_solver/src/solve/normalizes_to/mod.rs
index 6e6be829f52..54b92ebac1d 100644
--- a/compiler/rustc_next_trait_solver/src/solve/normalizes_to/mod.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/normalizes_to/mod.rs
@@ -194,6 +194,7 @@ where
         ecx: &mut EvalCtxt<'_, D>,
         goal: Goal<I, NormalizesTo<I>>,
         impl_def_id: I::ImplId,
+        then: impl FnOnce(&mut EvalCtxt<'_, D>, Certainty) -> QueryResult<I>,
     ) -> Result<Candidate<I>, NoSolution> {
         let cx = ecx.cx();
 
@@ -314,8 +315,7 @@ where
                         // nested goal for consistency.
                         ty::TypingMode::Coherence => {
                             ecx.add_goal(GoalSource::Misc, goal.with(cx, PredicateKind::Ambiguous));
-                            return ecx
-                                .evaluate_added_goals_and_make_canonical_response(Certainty::Yes);
+                            return then(ecx, Certainty::Yes);
                         }
                         ty::TypingMode::Analysis { .. }
                         | ty::TypingMode::Borrowck { .. }
@@ -325,8 +325,7 @@ where
                                 goal,
                                 goal.predicate.alias,
                             );
-                            return ecx
-                                .evaluate_added_goals_and_make_canonical_response(Certainty::Yes);
+                            return then(ecx, Certainty::Yes);
                         }
                     }
                 } else {
diff --git a/compiler/rustc_next_trait_solver/src/solve/trait_goals.rs b/compiler/rustc_next_trait_solver/src/solve/trait_goals.rs
index 3e720a47f32..a69e867289c 100644
--- a/compiler/rustc_next_trait_solver/src/solve/trait_goals.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/trait_goals.rs
@@ -55,6 +55,7 @@ where
         ecx: &mut EvalCtxt<'_, D>,
         goal: Goal<I, TraitPredicate<I>>,
         impl_def_id: I::ImplId,
+        then: impl FnOnce(&mut EvalCtxt<'_, D>, Certainty) -> QueryResult<I>,
     ) -> Result<Candidate<I>, NoSolution> {
         let cx = ecx.cx();
 
@@ -112,7 +113,7 @@ where
                     .map(|pred| goal.with(cx, pred)),
             );
 
-            ecx.evaluate_added_goals_and_make_canonical_response(maximal_certainty)
+            then(ecx, maximal_certainty)
         })
     }
 
diff --git a/compiler/rustc_type_ir/src/interner.rs b/compiler/rustc_type_ir/src/interner.rs
index 3a83b1739ff..cf58aec14a6 100644
--- a/compiler/rustc_type_ir/src/interner.rs
+++ b/compiler/rustc_type_ir/src/interner.rs
@@ -347,6 +347,7 @@ pub trait Interner:
         self_ty: Self::Ty,
         f: impl FnMut(Self::ImplId),
     );
+    fn for_each_blanket_impl(self, trait_def_id: Self::TraitId, f: impl FnMut(Self::ImplId));
 
     fn has_item_definition(self, def_id: Self::DefId) -> bool;