Add trait obligation tracking to FulfillCtxt and expose FnCtxt in rustc_infer using callback.

gavinleroy · gavinleroy · commit f9046e617cb7 · 2024-01-18T00:40:14.000+01:00
Pass each obligation to an fn callback with its respective inference context. This avoids needing to keep around copies of obligations or inference contexts.
diff --git a/compiler/rustc_hir_typeck/src/lib.rs b/compiler/rustc_hir_typeck/src/lib.rs
@@ -60,6 +60,7 @@ use rustc_hir::{HirIdMap, Node};
 use rustc_hir_analysis::astconv::AstConv;
 use rustc_hir_analysis::check::check_abi;
 use rustc_infer::infer::type_variable::{TypeVariableOrigin, TypeVariableOriginKind};
+use rustc_infer::traits::ObligationInspector;
 use rustc_middle::query::Providers;
 use rustc_middle::traits;
 use rustc_middle::ty::{self, Ty, TyCtxt};
@@ -139,7 +140,7 @@ fn used_trait_imports(tcx: TyCtxt<'_>, def_id: LocalDefId) -> &UnordSet<LocalDef
 
 fn typeck<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> &ty::TypeckResults<'tcx> {
     let fallback = move || tcx.type_of(def_id.to_def_id()).instantiate_identity();
-    typeck_with_fallback(tcx, def_id, fallback)
+    typeck_with_fallback(tcx, def_id, fallback, None)
 }
 
 /// Used only to get `TypeckResults` for type inference during error recovery.
@@ -149,14 +150,24 @@ fn diagnostic_only_typeck<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> &ty::T
         let span = tcx.hir().span(tcx.local_def_id_to_hir_id(def_id));
         Ty::new_error_with_message(tcx, span, "diagnostic only typeck table used")
     };
-    typeck_with_fallback(tcx, def_id, fallback)
+    typeck_with_fallback(tcx, def_id, fallback, None)
 }
 
-#[instrument(level = "debug", skip(tcx, fallback), ret)]
+pub fn inspect_typeck<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    def_id: LocalDefId,
+    inspect: ObligationInspector<'tcx>,
+) -> &'tcx ty::TypeckResults<'tcx> {
+    let fallback = move || tcx.type_of(def_id.to_def_id()).instantiate_identity();
+    typeck_with_fallback(tcx, def_id, fallback, Some(inspect))
+}
+
+#[instrument(level = "debug", skip(tcx, fallback, inspector), ret)]
 fn typeck_with_fallback<'tcx>(
     tcx: TyCtxt<'tcx>,
     def_id: LocalDefId,
     fallback: impl Fn() -> Ty<'tcx> + 'tcx,
+    inspector: Option<ObligationInspector<'tcx>>,
 ) -> &'tcx ty::TypeckResults<'tcx> {
     // Closures' typeck results come from their outermost function,
     // as they are part of the same "inference environment".
@@ -178,6 +189,9 @@ fn typeck_with_fallback<'tcx>(
     let param_env = tcx.param_env(def_id);
 
     let inh = Inherited::new(tcx, def_id);
+    if let Some(inspector) = inspector {
+        inh.infcx.attach_obligation_inspector(inspector);
+    }
     let mut fcx = FnCtxt::new(&inh, param_env, def_id);
 
     if let Some(hir::FnSig { header, decl, .. }) = fn_sig {
diff --git a/compiler/rustc_infer/src/infer/at.rs b/compiler/rustc_infer/src/infer/at.rs
@@ -90,6 +90,7 @@ impl<'tcx> InferCtxt<'tcx> {
             universe: self.universe.clone(),
             intercrate,
             next_trait_solver: self.next_trait_solver,
+            obligation_inspector: self.obligation_inspector.clone(),
         }
     }
 }
diff --git a/compiler/rustc_infer/src/infer/mod.rs b/compiler/rustc_infer/src/infer/mod.rs
@@ -13,7 +13,9 @@ use rustc_middle::infer::unify_key::{ConstVidKey, EffectVidKey};
 use self::opaque_types::OpaqueTypeStorage;
 pub(crate) use self::undo_log::{InferCtxtUndoLogs, Snapshot, UndoLog};
 
-use crate::traits::{self, ObligationCause, PredicateObligations, TraitEngine, TraitEngineExt};
+use crate::traits::{
+    self, ObligationCause, ObligationInspector, PredicateObligations, TraitEngine, TraitEngineExt,
+};
 
 use rustc_data_structures::fx::FxIndexMap;
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
@@ -334,6 +336,8 @@ pub struct InferCtxt<'tcx> {
     pub intercrate: bool,
 
     next_trait_solver: bool,
+
+    pub obligation_inspector: Cell<Option<ObligationInspector<'tcx>>>,
 }
 
 impl<'tcx> ty::InferCtxtLike for InferCtxt<'tcx> {
@@ -708,6 +712,7 @@ impl<'tcx> InferCtxtBuilder<'tcx> {
             universe: Cell::new(ty::UniverseIndex::ROOT),
             intercrate,
             next_trait_solver,
+            obligation_inspector: Cell::new(None),
         }
     }
 }
@@ -1726,6 +1731,14 @@ impl<'tcx> InferCtxt<'tcx> {
             }
         }
     }
+
+    pub fn attach_obligation_inspector(&self, inspector: ObligationInspector<'tcx>) {
+        debug_assert!(
+            self.obligation_inspector.get().is_none(),
+            "shouldn't override a set obligation inspector"
+        );
+        self.obligation_inspector.set(Some(inspector));
+    }
 }
 
 impl<'tcx> TypeErrCtxt<'_, 'tcx> {
diff --git a/compiler/rustc_infer/src/traits/mod.rs b/compiler/rustc_infer/src/traits/mod.rs
@@ -13,13 +13,16 @@ use std::hash::{Hash, Hasher};
 
 use hir::def_id::LocalDefId;
 use rustc_hir as hir;
+use rustc_middle::traits::query::NoSolution;
+use rustc_middle::traits::solve::Certainty;
 use rustc_middle::ty::error::{ExpectedFound, TypeError};
 use rustc_middle::ty::{self, Const, ToPredicate, Ty, TyCtxt};
 use rustc_span::Span;
 
 pub use self::FulfillmentErrorCode::*;
 pub use self::ImplSource::*;
 pub use self::SelectionError::*;
+use crate::infer::InferCtxt;
 
 pub use self::engine::{TraitEngine, TraitEngineExt};
 pub use self::project::MismatchedProjectionTypes;
@@ -117,6 +120,12 @@ pub type PredicateObligations<'tcx> = Vec<PredicateObligation<'tcx>>;
 
 pub type Selection<'tcx> = ImplSource<'tcx, PredicateObligation<'tcx>>;
 
+/// A callback that can be provided to `inspect_typeck`. Invoked on evaluation
+/// of root obligations.
+pub type ObligationInspector<'tcx> =
+    fn(&InferCtxt<'tcx>, &PredicateObligation<'tcx>, Result<Certainty, NoSolution>);
+
+#[derive(Clone)]
 pub struct FulfillmentError<'tcx> {
     pub obligation: PredicateObligation<'tcx>,
     pub code: FulfillmentErrorCode<'tcx>,
diff --git a/compiler/rustc_session/src/options.rs b/compiler/rustc_session/src/options.rs
@@ -1943,6 +1943,9 @@ written to standard error output)"),
         "for every macro invocation, print its name and arguments (default: no)"),
     track_diagnostics: bool = (false, parse_bool, [UNTRACKED],
         "tracks where in rustc a diagnostic was emitted"),
+    track_trait_obligations: bool = (false, parse_bool, [TRACKED],
+        "tracks evaluated obligations while trait solving, option is only \
+         valid when -Z next-solver=globally (default: no)"),
     // Diagnostics are considered side-effects of a query (see `QuerySideEffects`) and are saved
     // alongside query results and changes to translation options can affect diagnostics - so
     // translation options should be tracked.
diff --git a/compiler/rustc_trait_selection/src/solve/fulfill.rs b/compiler/rustc_trait_selection/src/solve/fulfill.rs
@@ -11,7 +11,7 @@ use rustc_middle::ty;
 use rustc_middle::ty::error::{ExpectedFound, TypeError};
 
 use super::eval_ctxt::GenerateProofTree;
-use super::{Certainty, InferCtxtEvalExt};
+use super::{Certainty, Goal, InferCtxtEvalExt};
 
 /// A trait engine using the new trait solver.
 ///
@@ -32,11 +32,34 @@ pub struct FulfillmentCtxt<'tcx> {
     /// gets rolled back. Because of this we explicitly check that we only
     /// use the context in exactly this snapshot.
     usable_in_snapshot: usize,
+
+    track_obligations: bool,
 }
 
 impl<'tcx> FulfillmentCtxt<'tcx> {
     pub fn new(infcx: &InferCtxt<'tcx>) -> FulfillmentCtxt<'tcx> {
-        FulfillmentCtxt { obligations: Vec::new(), usable_in_snapshot: infcx.num_open_snapshots() }
+        FulfillmentCtxt {
+            obligations: Vec::new(),
+            usable_in_snapshot: infcx.num_open_snapshots(),
+            track_obligations: infcx.tcx.sess.opts.unstable_opts.track_trait_obligations,
+        }
+    }
+
+    fn track_evaluated_obligation(
+        &self,
+        infcx: &InferCtxt<'tcx>,
+        obligation: &PredicateObligation<'tcx>,
+        result: &Result<(bool, Certainty, Vec<Goal<'tcx, ty::Predicate<'tcx>>>), NoSolution>,
+    ) {
+        if self.track_obligations {
+            if let Some(inspector) = infcx.obligation_inspector.get() {
+                let result = match result {
+                    Ok((_, c, _)) => Ok(*c),
+                    Err(NoSolution) => Err(NoSolution),
+                };
+                (inspector)(infcx, &obligation, result);
+            }
+        }
     }
 }
 
@@ -52,7 +75,8 @@ impl<'tcx> TraitEngine<'tcx> for FulfillmentCtxt<'tcx> {
     }
 
     fn collect_remaining_errors(&mut self, infcx: &InferCtxt<'tcx>) -> Vec<FulfillmentError<'tcx>> {
-        self.obligations
+        let errors = self
+            .obligations
             .drain(..)
             .map(|obligation| {
                 let code = infcx.probe(|_| {
@@ -81,7 +105,9 @@ impl<'tcx> TraitEngine<'tcx> for FulfillmentCtxt<'tcx> {
                     root_obligation: obligation,
                 }
             })
-            .collect()
+            .collect();
+
+        errors
     }
 
     fn select_where_possible(&mut self, infcx: &InferCtxt<'tcx>) -> Vec<FulfillmentError<'tcx>> {
@@ -95,65 +121,66 @@ impl<'tcx> TraitEngine<'tcx> for FulfillmentCtxt<'tcx> {
             let mut has_changed = false;
             for obligation in mem::take(&mut self.obligations) {
                 let goal = obligation.clone().into();
-                let (changed, certainty, nested_goals) =
-                    match infcx.evaluate_root_goal(goal, GenerateProofTree::IfEnabled).0 {
-                        Ok(result) => result,
-                        Err(NoSolution) => {
-                            errors.push(FulfillmentError {
-                                obligation: obligation.clone(),
-                                code: match goal.predicate.kind().skip_binder() {
-                                    ty::PredicateKind::Clause(ty::ClauseKind::Projection(_)) => {
-                                        FulfillmentErrorCode::CodeProjectionError(
-                                            // FIXME: This could be a `Sorts` if the term is a type
-                                            MismatchedProjectionTypes { err: TypeError::Mismatch },
-                                        )
-                                    }
-                                    ty::PredicateKind::NormalizesTo(..) => {
-                                        FulfillmentErrorCode::CodeProjectionError(
-                                            MismatchedProjectionTypes { err: TypeError::Mismatch },
-                                        )
-                                    }
-                                    ty::PredicateKind::AliasRelate(_, _, _) => {
-                                        FulfillmentErrorCode::CodeProjectionError(
-                                            MismatchedProjectionTypes { err: TypeError::Mismatch },
-                                        )
-                                    }
-                                    ty::PredicateKind::Subtype(pred) => {
-                                        let (a, b) = infcx.instantiate_binder_with_placeholders(
-                                            goal.predicate.kind().rebind((pred.a, pred.b)),
-                                        );
-                                        let expected_found = ExpectedFound::new(true, a, b);
-                                        FulfillmentErrorCode::CodeSubtypeError(
-                                            expected_found,
-                                            TypeError::Sorts(expected_found),
-                                        )
-                                    }
-                                    ty::PredicateKind::Coerce(pred) => {
-                                        let (a, b) = infcx.instantiate_binder_with_placeholders(
-                                            goal.predicate.kind().rebind((pred.a, pred.b)),
-                                        );
-                                        let expected_found = ExpectedFound::new(false, a, b);
-                                        FulfillmentErrorCode::CodeSubtypeError(
-                                            expected_found,
-                                            TypeError::Sorts(expected_found),
-                                        )
-                                    }
-                                    ty::PredicateKind::Clause(_)
-                                    | ty::PredicateKind::ObjectSafe(_)
-                                    | ty::PredicateKind::Ambiguous => {
-                                        FulfillmentErrorCode::CodeSelectionError(
-                                            SelectionError::Unimplemented,
-                                        )
-                                    }
-                                    ty::PredicateKind::ConstEquate(..) => {
-                                        bug!("unexpected goal: {goal:?}")
-                                    }
-                                },
-                                root_obligation: obligation,
-                            });
-                            continue;
-                        }
-                    };
+                let result = infcx.evaluate_root_goal(goal, GenerateProofTree::IfEnabled).0;
+                self.track_evaluated_obligation(infcx, &obligation, &result);
+                let (changed, certainty, nested_goals) = match result {
+                    Ok(result) => result,
+                    Err(NoSolution) => {
+                        errors.push(FulfillmentError {
+                            obligation: obligation.clone(),
+                            code: match goal.predicate.kind().skip_binder() {
+                                ty::PredicateKind::Clause(ty::ClauseKind::Projection(_)) => {
+                                    FulfillmentErrorCode::CodeProjectionError(
+                                        // FIXME: This could be a `Sorts` if the term is a type
+                                        MismatchedProjectionTypes { err: TypeError::Mismatch },
+                                    )
+                                }
+                                ty::PredicateKind::NormalizesTo(..) => {
+                                    FulfillmentErrorCode::CodeProjectionError(
+                                        MismatchedProjectionTypes { err: TypeError::Mismatch },
+                                    )
+                                }
+                                ty::PredicateKind::AliasRelate(_, _, _) => {
+                                    FulfillmentErrorCode::CodeProjectionError(
+                                        MismatchedProjectionTypes { err: TypeError::Mismatch },
+                                    )
+                                }
+                                ty::PredicateKind::Subtype(pred) => {
+                                    let (a, b) = infcx.instantiate_binder_with_placeholders(
+                                        goal.predicate.kind().rebind((pred.a, pred.b)),
+                                    );
+                                    let expected_found = ExpectedFound::new(true, a, b);
+                                    FulfillmentErrorCode::CodeSubtypeError(
+                                        expected_found,
+                                        TypeError::Sorts(expected_found),
+                                    )
+                                }
+                                ty::PredicateKind::Coerce(pred) => {
+                                    let (a, b) = infcx.instantiate_binder_with_placeholders(
+                                        goal.predicate.kind().rebind((pred.a, pred.b)),
+                                    );
+                                    let expected_found = ExpectedFound::new(false, a, b);
+                                    FulfillmentErrorCode::CodeSubtypeError(
+                                        expected_found,
+                                        TypeError::Sorts(expected_found),
+                                    )
+                                }
+                                ty::PredicateKind::Clause(_)
+                                | ty::PredicateKind::ObjectSafe(_)
+                                | ty::PredicateKind::Ambiguous => {
+                                    FulfillmentErrorCode::CodeSelectionError(
+                                        SelectionError::Unimplemented,
+                                    )
+                                }
+                                ty::PredicateKind::ConstEquate(..) => {
+                                    bug!("unexpected goal: {goal:?}")
+                                }
+                            },
+                            root_obligation: obligation,
+                        });
+                        continue;
+                    }
+                };
                 // Push any nested goals that we get from unifying our canonical response
                 // with our obligation onto the fulfillment context.
                 self.obligations.extend(nested_goals.into_iter().map(|goal| {
diff --git a/tests/ui/traits/track_trait_obligations.rs b/tests/ui/traits/track_trait_obligations.rs
@@ -0,0 +1,32 @@
+// compile-flags: -Ztrack-trait-obligations
+// run-pass
+
+// Just making sure this flag is accepted and doesn't crash the compiler
+use traits::IntoString;
+
+fn does_impl_into_string<T: IntoString>(_: T) {}
+
+fn main() {
+    let v = vec![(0, 1), (2, 3)];
+
+    does_impl_into_string(v);
+}
+
+mod traits {
+    pub trait IntoString {
+        fn to_string(&self) -> String;
+    }
+
+    impl IntoString for (i32, i32) {
+        fn to_string(&self) -> String {
+            format!("({}, {})", self.0, self.1)
+        }
+    }
+
+    impl<T: IntoString> IntoString for Vec<T> {
+        fn to_string(&self) -> String {
+            let s = self.iter().map(|v| v.to_string()).collect::<Vec<_>>().join(", ");
+            format!("[{s}]")
+        }
+    }
+}

Original file line number	Diff line number	Diff line change
`@@ -90,6 +90,7 @@ impl<'tcx> InferCtxt<'tcx> {`
`90`	`90`	`universe: self.universe.clone(),`
`91`	`91`	`intercrate,`
`92`	`92`	`next_trait_solver: self.next_trait_solver,`
	`93`	`+ obligation_inspector: self.obligation_inspector.clone(),`
`93`	`94`	`}`
`94`	`95`	`}`
`95`	`96`	`}`
Original file line number	Diff line number	Diff line change
`@@ -13,7 +13,9 @@ use rustc_middle::infer::unify_key::{ConstVidKey, EffectVidKey};`
`13`	`13`	`use self::opaque_types::OpaqueTypeStorage;`
`14`	`14`	`pub(crate) use self::undo_log::{InferCtxtUndoLogs, Snapshot, UndoLog};`
`15`	`15`
`16`		`-use crate::traits::{self, ObligationCause, PredicateObligations, TraitEngine, TraitEngineExt};`
	`16`	`+use crate::traits::{`
	`17`	`+ self, ObligationCause, ObligationInspector, PredicateObligations, TraitEngine, TraitEngineExt,`
	`18`	`+};`
`17`	`19`
`18`	`20`	`use rustc_data_structures::fx::FxIndexMap;`
`19`	`21`	`use rustc_data_structures::fx::{FxHashMap, FxHashSet};`
`@@ -334,6 +336,8 @@ pub struct InferCtxt<'tcx> {`
`334`	`336`	`pub intercrate: bool,`
`335`	`337`
`336`	`338`	`next_trait_solver: bool,`
	`339`	`+`
	`340`	`+ pub obligation_inspector: Cell<Option<ObligationInspector<'tcx>>>,`
`337`	`341`	`}`
`338`	`342`
`339`	`343`	`impl<'tcx> ty::InferCtxtLike for InferCtxt<'tcx> {`
`@@ -708,6 +712,7 @@ impl<'tcx> InferCtxtBuilder<'tcx> {`
`708`	`712`	`universe: Cell::new(ty::UniverseIndex::ROOT),`
`709`	`713`	`intercrate,`
`710`	`714`	`next_trait_solver,`
	`715`	`+ obligation_inspector: Cell::new(None),`
`711`	`716`	`}`
`712`	`717`	`}`
`713`	`718`	`}`
`@@ -1726,6 +1731,14 @@ impl<'tcx> InferCtxt<'tcx> {`
`1726`	`1731`	`}`
`1727`	`1732`	`}`
`1728`	`1733`	`}`
	`1734`	`+`
	`1735`	`+ pub fn attach_obligation_inspector(&self, inspector: ObligationInspector<'tcx>) {`
	`1736`	`+ debug_assert!(`
	`1737`	`+ self.obligation_inspector.get().is_none(),`
	`1738`	`+ "shouldn't override a set obligation inspector"`
	`1739`	`+ );`
	`1740`	`+ self.obligation_inspector.set(Some(inspector));`
	`1741`	`+ }`
`1729`	`1742`	`}`
`1730`	`1743`
`1731`	`1744`	`impl<'tcx> TypeErrCtxt<'_, 'tcx> {`