hir typeck: look into nested goals

lcnr · lcnr · commit fe86a846c83b · 2024-04-24T15:24:15.000Z
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.
diff --git a/compiler/rustc_hir_typeck/src/closure.rs b/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)
                 })?
             }
diff --git a/compiler/rustc_hir_typeck/src/fn_ctxt/_impl.rs b/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
diff --git a/compiler/rustc_hir_typeck/src/fn_ctxt/inspect_obligations.rs b/compiler/rustc_hir_typeck/src/fn_ctxt/inspect_obligations.rs
@@ -1,50 +1,54 @@
 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_trait_selection::solve::inspect::ProofTreeInferCtxtExt;
+use rustc_trait_selection::solve::inspect::{InspectConfig, InspectGoal, ProofTreeVisitor};
 
 impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
     #[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.self_type_matches_expected_vid(data.self_ty(), expected_vid)
+            }
+            ty::PredicateKind::Clause(ty::ClauseKind::Projection(data)) => {
+                self.self_type_matches_expected_vid(data.projection_ty.self_ty(), expected_vid)
+            }
+            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)]
@@ -53,12 +57,63 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         debug!(?self_ty);
 
         match *self_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
-            }
+            ty::Infer(ty::TyVar(found_vid)) => 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();
+        let mut obligations_for_self_ty = vec![];
+        for obligation in obligations {
+            let mut visitor = NestedObligationsForSelfTy {
+                fcx: self,
+                ty_var_root: self.root_var(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
+    }
+}
+
+struct NestedObligationsForSelfTy<'a, 'tcx> {
+    fcx: &'a FnCtxt<'a, 'tcx>,
+    ty_var_root: 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 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: 3 }
+    }
+
+    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.ty_var_root) {
+            self.obligations_for_self_ty.push(traits::Obligation::new(
+                tcx,
+                self.root_cause.clone(),
+                goal.param_env,
+                goal.predicate,
+            ));
+        }
+
+        if let Some(candidate) = inspect_goal.unique_applicable_candidate() {
+            candidate.visit_nested_no_probe(self)
+        }
+    }
+}
diff --git a/compiler/rustc_trait_selection/src/solve/inspect/analyse.rs b/compiler/rustc_trait_selection/src/solve/inspect/analyse.rs
@@ -20,6 +20,10 @@ use rustc_middle::ty;
 use crate::solve::inspect::ProofTreeBuilder;
 use crate::solve::{GenerateProofTree, InferCtxtEvalExt};
 
+pub struct InspectConfig {
+    pub max_depth: usize,
+}
+
 pub struct InspectGoal<'a, 'tcx> {
     infcx: &'a InferCtxt<'tcx>,
     depth: usize,
@@ -48,15 +52,42 @@ impl<'a, 'tcx> InspectCandidate<'a, 'tcx> {
         self.result.map(|c| c.value.certainty)
     }
 
+    pub fn visit_nested_no_probe<V: ProofTreeVisitor<'tcx>>(&self, visitor: &mut V) -> V::Result {
+        if self.goal.depth < visitor.config().max_depth {
+            let infcx = self.goal.infcx;
+            let mut instantiated_goals = vec![];
+            for goal in &self.nested_goals {
+                let goal = ProofTreeBuilder::instantiate_canonical_state(
+                    infcx,
+                    self.goal.goal.param_env,
+                    self.goal.orig_values,
+                    *goal,
+                );
+                instantiated_goals.push(goal);
+            }
+
+            for &goal in &instantiated_goals {
+                let (_, proof_tree) = infcx.evaluate_root_goal(goal, GenerateProofTree::Yes);
+                let proof_tree = proof_tree.unwrap();
+                try_visit!(visitor.visit_goal(&InspectGoal::new(
+                    infcx,
+                    self.goal.depth + 1,
+                    &proof_tree,
+                )));
+            }
+        }
+
+        V::Result::output()
+    }
+
     /// Visit the nested goals of this candidate.
     ///
     /// FIXME(@lcnr): we have to slightly adapt this API
     /// to also use it to compute the most relevant goal
     /// for fulfillment errors. Will do that once we actually
     /// need it.
     pub fn visit_nested<V: ProofTreeVisitor<'tcx>>(&self, visitor: &mut V) -> V::Result {
-        // HACK: An arbitrary cutoff to avoid dealing with overflow and cycles.
-        if self.goal.depth <= 10 {
+        if self.goal.depth < visitor.config().max_depth {
             let infcx = self.goal.infcx;
             infcx.probe(|_| {
                 let mut instantiated_goals = vec![];
@@ -167,6 +198,16 @@ impl<'a, 'tcx> InspectGoal<'a, 'tcx> {
         }
     }
 
+    /// Returns the single candidate applicable for the current goal, if it exists.
+    ///
+    /// Returns `None` if there are no multiple applicable candidates.
+    pub fn unique_applicable_candidate(&'a self) -> Option<InspectCandidate<'a, 'tcx>> {
+        // FIXME(-Znext-solver): This does not handle impl candidates
+        // hidden by env candidates.
+        let mut candidates = self.candidates();
+        candidates.pop().filter(|_| candidates.is_empty())
+    }
+
     pub fn candidates(&'a self) -> Vec<InspectCandidate<'a, 'tcx>> {
         let mut candidates = vec![];
         let last_eval_step = match self.evaluation.evaluation.kind {
@@ -213,6 +254,10 @@ impl<'a, 'tcx> InspectGoal<'a, 'tcx> {
 pub trait ProofTreeVisitor<'tcx> {
     type Result: VisitorResult = ();
 
+    fn config(&self) -> InspectConfig {
+        InspectConfig { max_depth: 10 }
+    }
+
     fn visit_goal(&mut self, goal: &InspectGoal<'_, 'tcx>) -> Self::Result;
 }
 
@@ -223,10 +268,8 @@ impl<'tcx> InferCtxt<'tcx> {
         goal: Goal<'tcx, ty::Predicate<'tcx>>,
         visitor: &mut V,
     ) -> V::Result {
-        self.probe(|_| {
-            let (_, proof_tree) = self.evaluate_root_goal(goal, GenerateProofTree::Yes);
-            let proof_tree = proof_tree.unwrap();
-            visitor.visit_goal(&InspectGoal::new(self, 0, &proof_tree))
-        })
+        let (_, proof_tree) = self.evaluate_root_goal(goal, GenerateProofTree::Yes);
+        let proof_tree = proof_tree.unwrap();
+        visitor.visit_goal(&InspectGoal::new(self, 0, &proof_tree))
     }
 }
diff --git a/compiler/rustc_trait_selection/src/traits/coherence.rs b/compiler/rustc_trait_selection/src/traits/coherence.rs
@@ -1157,5 +1157,5 @@ fn search_ambiguity_causes<'tcx>(
     goal: Goal<'tcx, ty::Predicate<'tcx>>,
     causes: &mut FxIndexSet<IntercrateAmbiguityCause<'tcx>>,
 ) {
-    infcx.visit_proof_tree(goal, &mut AmbiguityCausesVisitor { causes });
+    infcx.probe(|_| infcx.visit_proof_tree(goal, &mut AmbiguityCausesVisitor { causes }));
 }
diff --git a/tests/ui/closures/infer-signature-from-impl.next.stderr b/tests/ui/closures/infer-signature-from-impl.next.stderr
diff --git a/tests/ui/closures/infer-signature-from-impl.rs b/tests/ui/closures/infer-signature-from-impl.rs
@@ -1,8 +1,7 @@
 //@ revisions: current next
 //@ ignore-compare-mode-next-solver (explicit revisions)
 //@[next] compile-flags: -Znext-solver
-//@[next] known-bug: trait-system-refactor-initiative#71
-//@[current] check-pass
+//@ check-pass
 
 trait Foo {}
 fn needs_foo<T>(_: T)

Original file line number	Diff line number	Diff line change
`@@ -1157,5 +1157,5 @@ fn search_ambiguity_causes<'tcx>(`
`1157`	`1157`	`goal: Goal<'tcx, ty::Predicate<'tcx>>,`
`1158`	`1158`	`causes: &mut FxIndexSet<IntercrateAmbiguityCause<'tcx>>,`
`1159`	`1159`	`) {`
`1160`		`- infcx.visit_proof_tree(goal, &mut AmbiguityCausesVisitor { causes });`
	`1160`	`+ infcx.probe(\|_\| infcx.visit_proof_tree(goal, &mut AmbiguityCausesVisitor { causes }));`
`1161`	`1161`	`}`