hir typeck: look into nested goals

uses a `ProofTreeVisitor` to look into nested
goals when looking at the pending obligations
during hir typeck. Used by closure signature
inference, coercion, and for async functions.

Author: lcnr

compiler/rustc_hir_typeck/src/closure.rs

@@ -342,7 +342,9 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             ty::Infer(ty::TyVar(vid)) => self.deduce_closure_signature_from_predicates(
                 Ty::new_var(self.tcx, self.root_var(vid)),
                 closure_kind,
-                self.obligations_for_self_ty(vid).map(|obl| (obl.predicate, obl.cause.span)),
+                self.obligations_for_self_ty(vid)
+                    .into_iter()
+                    .map(|obl| (obl.predicate, obl.cause.span)),
             ),
             ty::FnPtr(sig) => match closure_kind {
                 hir::ClosureKind::Closure => {
@@ -889,7 +891,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
 
         let output_ty = match *ret_ty.kind() {
             ty::Infer(ty::TyVar(ret_vid)) => {
-                self.obligations_for_self_ty(ret_vid).find_map(|obligation| {
+                self.obligations_for_self_ty(ret_vid).into_iter().find_map(|obligation| {
                     get_future_output(obligation.predicate, obligation.cause.span)
                 })?
             }

compiler/rustc_hir_typeck/src/fn_ctxt/_impl.rs

@@ -637,7 +637,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
 
     pub(crate) fn type_var_is_sized(&self, self_ty: ty::TyVid) -> bool {
         let sized_did = self.tcx.lang_items().sized_trait();
-        self.obligations_for_self_ty(self_ty).any(|obligation| {
+        self.obligations_for_self_ty(self_ty).into_iter().any(|obligation| {
             match obligation.predicate.kind().skip_binder() {
                 ty::PredicateKind::Clause(ty::ClauseKind::Trait(data)) => {
                     Some(data.def_id()) == sized_did

compiler/rustc_hir_typeck/src/fn_ctxt/inspect_obligations.rs

@@ -1,64 +1,140 @@
+//! A utility module to inspect currently ambiguous obligations in the current context.
+use crate::rustc_middle::ty::TypeVisitableExt;
 use crate::FnCtxt;
+use rustc_infer::traits::solve::Goal;
+use rustc_infer::traits::{self, ObligationCause};
 use rustc_middle::ty::{self, Ty};
-use rustc_infer::traits;
-use rustc_data_structures::captures::Captures;
+use rustc_span::Span;
+use rustc_trait_selection::solve::inspect::ProofTreeInferCtxtExt;
+use rustc_trait_selection::solve::inspect::{InspectConfig, InspectGoal, ProofTreeVisitor};
 
 impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
+    /// Returns a list of all obligations whose self type has been unified
+    /// with the unconstrained type `self_ty`.
     #[instrument(skip(self), level = "debug")]
-    pub(crate) fn obligations_for_self_ty<'b>(
-        &'b self,
+    pub(crate) fn obligations_for_self_ty(
+        &self,
         self_ty: ty::TyVid,
-    ) -> impl DoubleEndedIterator<Item = traits::PredicateObligation<'tcx>> + Captures<'tcx> + 'b
-    {
-        let ty_var_root = self.root_var(self_ty);
-        trace!("pending_obligations = {:#?}", self.fulfillment_cx.borrow().pending_obligations());
-
-        self.fulfillment_cx.borrow().pending_obligations().into_iter().filter_map(
-            move |obligation| match &obligation.predicate.kind().skip_binder() {
-                ty::PredicateKind::Clause(ty::ClauseKind::Projection(data))
-                    if self.self_type_matches_expected_vid(
-                        data.projection_ty.self_ty(),
-                        ty_var_root,
-                    ) =>
-                {
-                    Some(obligation)
-                }
-                ty::PredicateKind::Clause(ty::ClauseKind::Trait(data))
-                    if self.self_type_matches_expected_vid(data.self_ty(), ty_var_root) =>
-                {
-                    Some(obligation)
-                }
+    ) -> Vec<traits::PredicateObligation<'tcx>> {
+        if self.next_trait_solver() {
+            self.obligations_for_self_ty_next(self_ty)
+        } else {
+            let ty_var_root = self.root_var(self_ty);
+            let mut obligations = self.fulfillment_cx.borrow().pending_obligations();
+            trace!("pending_obligations = {:#?}", obligations);
+            obligations
+                .retain(|obligation| self.predicate_has_self_ty(obligation.predicate, ty_var_root));
+            obligations
+        }
+    }
 
-                ty::PredicateKind::Clause(ty::ClauseKind::Trait(..))
-                | ty::PredicateKind::Clause(ty::ClauseKind::Projection(..))
-                | ty::PredicateKind::Clause(ty::ClauseKind::ConstArgHasType(..))
-                | ty::PredicateKind::Subtype(..)
-                | ty::PredicateKind::Coerce(..)
-                | ty::PredicateKind::Clause(ty::ClauseKind::RegionOutlives(..))
-                | ty::PredicateKind::Clause(ty::ClauseKind::TypeOutlives(..))
-                | ty::PredicateKind::Clause(ty::ClauseKind::WellFormed(..))
-                | ty::PredicateKind::ObjectSafe(..)
-                | ty::PredicateKind::NormalizesTo(..)
-                | ty::PredicateKind::AliasRelate(..)
-                | ty::PredicateKind::Clause(ty::ClauseKind::ConstEvaluatable(..))
-                | ty::PredicateKind::ConstEquate(..)
-                | ty::PredicateKind::Ambiguous => None,
-            },
-        )
+    #[instrument(level = "debug", skip(self), ret)]
+    fn predicate_has_self_ty(
+        &self,
+        predicate: ty::Predicate<'tcx>,
+        expected_vid: ty::TyVid,
+    ) -> bool {
+        match predicate.kind().skip_binder() {
+            ty::PredicateKind::Clause(ty::ClauseKind::Trait(data)) => {
+                self.type_matches_expected_vid(expected_vid, data.self_ty())
+            }
+            ty::PredicateKind::Clause(ty::ClauseKind::Projection(data)) => {
+                self.type_matches_expected_vid(expected_vid, data.projection_ty.self_ty())
+            }
+            ty::PredicateKind::Clause(ty::ClauseKind::ConstArgHasType(..))
+            | ty::PredicateKind::Subtype(..)
+            | ty::PredicateKind::Coerce(..)
+            | ty::PredicateKind::Clause(ty::ClauseKind::RegionOutlives(..))
+            | ty::PredicateKind::Clause(ty::ClauseKind::TypeOutlives(..))
+            | ty::PredicateKind::Clause(ty::ClauseKind::WellFormed(..))
+            | ty::PredicateKind::ObjectSafe(..)
+            | ty::PredicateKind::NormalizesTo(..)
+            | ty::PredicateKind::AliasRelate(..)
+            | ty::PredicateKind::Clause(ty::ClauseKind::ConstEvaluatable(..))
+            | ty::PredicateKind::ConstEquate(..)
+            | ty::PredicateKind::Ambiguous => false,
+        }
     }
 
     #[instrument(level = "debug", skip(self), ret)]
-    fn self_type_matches_expected_vid(&self, self_ty: Ty<'tcx>, expected_vid: ty::TyVid) -> bool {
-        let self_ty = self.shallow_resolve(self_ty);
-        debug!(?self_ty);
+    fn type_matches_expected_vid(&self, expected_vid: ty::TyVid, ty: Ty<'tcx>) -> bool {
+        let ty = self.shallow_resolve(ty);
+        debug!(?ty);
 
-        match *self_ty.kind() {
+        match *ty.kind() {
             ty::Infer(ty::TyVar(found_vid)) => {
-                let found_vid = self.root_var(found_vid);
-                debug!("self_type_matches_expected_vid - found_vid={:?}", found_vid);
-                expected_vid == found_vid
+                self.root_var(expected_vid) == self.root_var(found_vid)
             }
             _ => false,
         }
     }
-}
+
+    pub(crate) fn obligations_for_self_ty_next(
+        &self,
+        self_ty: ty::TyVid,
+    ) -> Vec<traits::PredicateObligation<'tcx>> {
+        let obligations = self.fulfillment_cx.borrow().pending_obligations();
+        debug!(?obligations);
+        let mut obligations_for_self_ty = vec![];
+        for obligation in obligations {
+            let mut visitor = NestedObligationsForSelfTy {
+                fcx: self,
+                self_ty,
+                obligations_for_self_ty: &mut obligations_for_self_ty,
+                root_cause: &obligation.cause,
+            };
+
+            let goal = Goal::new(self.tcx, obligation.param_env, obligation.predicate);
+            self.visit_proof_tree(goal, &mut visitor);
+        }
+
+        obligations_for_self_ty.retain_mut(|obligation| {
+            obligation.predicate = self.resolve_vars_if_possible(obligation.predicate);
+            !obligation.predicate.has_placeholders()
+        });
+        obligations_for_self_ty
+    }
+}
+
+struct NestedObligationsForSelfTy<'a, 'tcx> {
+    fcx: &'a FnCtxt<'a, 'tcx>,
+    self_ty: ty::TyVid,
+    root_cause: &'a ObligationCause<'tcx>,
+    obligations_for_self_ty: &'a mut Vec<traits::PredicateObligation<'tcx>>,
+}
+
+impl<'a, 'tcx> ProofTreeVisitor<'tcx> for NestedObligationsForSelfTy<'a, 'tcx> {
+    type Result = ();
+
+    fn span(&self) -> Span {
+        self.root_cause.span
+    }
+
+    fn config(&self) -> InspectConfig {
+        // Using an intentionally low depth to minimize the chance of future
+        // breaking changes in case we adapt the approach later on. This also
+        // avoids any hangs for exponentially growing proof trees.
+        InspectConfig { max_depth: 5 }
+    }
+
+    fn visit_goal(&mut self, inspect_goal: &InspectGoal<'_, 'tcx>) {
+        let tcx = self.fcx.tcx;
+        let goal = inspect_goal.goal();
+        if self.fcx.predicate_has_self_ty(goal.predicate, self.self_ty) {
+            self.obligations_for_self_ty.push(traits::Obligation::new(
+                tcx,
+                self.root_cause.clone(),
+                goal.param_env,
+                goal.predicate,
+            ));
+        }
+
+        // If there's a unique way to prove a given goal, recurse into
+        // that candidate. This means that for `impl<F: FnOnce(u32)> Trait<F> for () {}`
+        // and a `(): Trait<?0>` goal we recurse into the impl and look at
+        // the nested `?0: FnOnce(u32)` goal.
+        if let Some(candidate) = inspect_goal.unique_applicable_candidate() {
+            candidate.visit_nested_no_probe(self)
+        }
+    }
+}

compiler/rustc_middle/src/traits/solve/inspect.rs

@@ -102,6 +102,7 @@ pub struct Probe<'tcx> {
     /// What happened inside of this probe in chronological order.
     pub steps: Vec<ProbeStep<'tcx>>,
     pub kind: ProbeKind<'tcx>,
+    pub final_state: CanonicalState<'tcx, ()>,
 }
 
 impl Debug for Probe<'_> {

compiler/rustc_trait_selection/src/solve/assembly/mod.rs

@@ -1,7 +1,7 @@
 //! Code shared by trait and projection goals for candidate assembly.
 
-use super::{EvalCtxt, SolverMode};
 use crate::solve::GoalSource;
+use crate::solve::{inspect, EvalCtxt, SolverMode};
 use crate::traits::coherence;
 use rustc_hir::def_id::DefId;
 use rustc_infer::traits::query::NoSolution;
@@ -16,6 +16,7 @@ use rustc_middle::ty::{fast_reject, TypeFoldable};
 use rustc_middle::ty::{ToPredicate, TypeVisitableExt};
 use rustc_span::{ErrorGuaranteed, DUMMY_SP};
 use std::fmt::Debug;
+use std::mem;
 
 pub(super) mod structural_traits;
 
@@ -315,20 +316,17 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
     }
 
     fn forced_ambiguity(&mut self, cause: MaybeCause) -> Vec<Candidate<'tcx>> {
-        let source = CandidateSource::BuiltinImpl(BuiltinImplSource::Misc);
-        let certainty = Certainty::Maybe(cause);
         // This may fail if `try_evaluate_added_goals` overflows because it
         // fails to reach a fixpoint but ends up getting an error after
         // running for some additional step.
         //
-        // FIXME: Add a test for this. It seems to be necessary for typenum but
-        // is incredibly hard to minimize as it may rely on being inside of a
-        // trait solver cycle.
-        let result = self.evaluate_added_goals_and_make_canonical_response(certainty);
-        let mut dummy_probe = self.inspect.new_probe();
-        dummy_probe.probe_kind(ProbeKind::TraitCandidate { source, result });
-        self.inspect.finish_probe(dummy_probe);
-        if let Ok(result) = result { vec![Candidate { source, result }] } else { vec![] }
+        // cc trait-system-refactor-initiative#105
+        let source = CandidateSource::BuiltinImpl(BuiltinImplSource::Misc);
+        let certainty = Certainty::Maybe(cause);
+        let result = self
+            .probe_trait_candidate(source)
+            .enter(|this| this.evaluate_added_goals_and_make_canonical_response(certainty));
+        if let Ok(cand) = result { vec![cand] } else { vec![] }
     }
 
     #[instrument(level = "debug", skip_all)]
@@ -813,6 +811,11 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
         goal: Goal<'tcx, G>,
         candidates: &mut Vec<Candidate<'tcx>>,
     ) {
+        // HACK: We temporarily remove the `ProofTreeBuilder` to
+        // avoid adding `Trait` candidates to the candidates used
+        // to prove the current goal.
+        let inspect = mem::replace(&mut self.inspect, inspect::ProofTreeBuilder::new_noop());
+
         let tcx = self.tcx();
         let trait_goal: Goal<'tcx, ty::TraitPredicate<'tcx>> =
             goal.with(tcx, goal.predicate.trait_ref(tcx));
@@ -846,6 +849,7 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
                 }
             }
         }
+        self.inspect = inspect;
     }
 
     /// If there are multiple ways to prove a trait or projection goal, we have