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

Use Lrc and Lock instead of Rc and RefCell.

Author: gavinleroy

compiler/rustc_hir_typeck/src/lib.rs

@@ -139,7 +139,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::<&dyn Fn(&FnCtxt<'_, 'tcx>)>)
 }
 
 /// Used only to get `TypeckResults` for type inference during error recovery.
@@ -149,14 +149,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::<&dyn Fn(&FnCtxt<'_, 'tcx>)>)
 }
 
-#[instrument(level = "debug", skip(tcx, fallback), ret)]
+pub fn inspect_typeck<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    def_id: LocalDefId,
+    inspect: impl FnOnce(&FnCtxt<'_, '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, inspect_ctxt), ret)]
 fn typeck_with_fallback<'tcx>(
     tcx: TyCtxt<'tcx>,
     def_id: LocalDefId,
     fallback: impl Fn() -> Ty<'tcx> + 'tcx,
+    inspect_ctxt: Option<impl FnOnce(&FnCtxt<'_, 'tcx>)>,
 ) -> &'tcx ty::TypeckResults<'tcx> {
     // Closures' typeck results come from their outermost function,
     // as they are part of the same "inference environment".
@@ -277,6 +287,10 @@ fn typeck_with_fallback<'tcx>(
         fcx.report_ambiguity_errors();
     }
 
+    if let Some(inspect) = inspect_ctxt {
+        inspect(&fcx);
+    }
+
     if let None = fcx.infcx.tainted_by_errors() {
         fcx.check_transmutes();
     }

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,
+            fulfilled_obligations: Lrc::clone(&self.fulfilled_obligations),
         }
     }
 }

compiler/rustc_infer/src/infer/mod.rs

@@ -17,7 +17,7 @@ use crate::traits::{self, ObligationCause, PredicateObligations, TraitEngine, Tr
 
 use rustc_data_structures::fx::FxIndexMap;
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
-use rustc_data_structures::sync::Lrc;
+use rustc_data_structures::sync::{Lock, Lrc};
 use rustc_data_structures::undo_log::Rollback;
 use rustc_data_structures::unify as ut;
 use rustc_errors::{DiagCtxt, DiagnosticBuilder, ErrorGuaranteed};
@@ -334,6 +334,8 @@ pub struct InferCtxt<'tcx> {
     pub intercrate: bool,
 
     next_trait_solver: bool,
+
+    pub fulfilled_obligations: Lrc<Lock<Vec<traits::FulfilledObligation<'tcx>>>>,
 }
 
 impl<'tcx> ty::InferCtxtLike for InferCtxt<'tcx> {
@@ -708,6 +710,7 @@ impl<'tcx> InferCtxtBuilder<'tcx> {
             universe: Cell::new(ty::UniverseIndex::ROOT),
             intercrate,
             next_trait_solver,
+            fulfilled_obligations: Default::default(),
         }
     }
 }

compiler/rustc_infer/src/traits/mod.rs

@@ -109,6 +109,11 @@ impl<'tcx> PolyTraitObligation<'tcx> {
     }
 }
 
+pub enum FulfilledObligation<'tcx> {
+    Success(PredicateObligation<'tcx>),
+    Failure(FulfillmentError<'tcx>),
+}
+
 // `PredicateObligation` is used a lot. Make sure it doesn't unintentionally get bigger.
 #[cfg(all(target_arch = "x86_64", target_pointer_width = "64"))]
 static_assert_size!(PredicateObligation<'_>, 48);
@@ -117,6 +122,7 @@ pub type PredicateObligations<'tcx> = Vec<PredicateObligation<'tcx>>;
 
 pub type Selection<'tcx> = ImplSource<'tcx, PredicateObligation<'tcx>>;
 
+#[derive(Clone)]
 pub struct FulfillmentError<'tcx> {
     pub obligation: PredicateObligation<'tcx>,
     pub code: FulfillmentErrorCode<'tcx>,

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 fulfilled 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.

compiler/rustc_trait_selection/src/solve/fulfill.rs

@@ -1,11 +1,12 @@
 use std::mem;
 
+use rustc_data_structures::sync::{Lock, Lrc};
 use rustc_infer::infer::InferCtxt;
 use rustc_infer::traits::solve::MaybeCause;
 use rustc_infer::traits::Obligation;
 use rustc_infer::traits::{
-    query::NoSolution, FulfillmentError, FulfillmentErrorCode, MismatchedProjectionTypes,
-    PredicateObligation, SelectionError, TraitEngine,
+    query::NoSolution, FulfilledObligation, FulfillmentError, FulfillmentErrorCode,
+    MismatchedProjectionTypes, PredicateObligation, SelectionError, TraitEngine,
 };
 use rustc_middle::ty;
 use rustc_middle::ty::error::{ExpectedFound, TypeError};
@@ -32,11 +33,40 @@ 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,
+
+    tracked_obligations: Option<Lrc<Lock<Vec<FulfilledObligation<'tcx>>>>>,
 }
 
 impl<'tcx> FulfillmentCtxt<'tcx> {
     pub fn new(infcx: &InferCtxt<'tcx>) -> FulfillmentCtxt<'tcx> {
-        FulfillmentCtxt { obligations: Vec::new(), usable_in_snapshot: infcx.num_open_snapshots() }
+        let tracked_obligations = if infcx.tcx.sess.opts.unstable_opts.track_trait_obligations {
+            Some(Lrc::clone(&infcx.fulfilled_obligations))
+        } else {
+            None
+        };
+
+        FulfillmentCtxt {
+            obligations: Vec::new(),
+            usable_in_snapshot: infcx.num_open_snapshots(),
+            tracked_obligations,
+        }
+    }
+
+    fn track_fulfillment_errors<'b, 'a: 'b>(
+        &'a self,
+        errors: impl IntoIterator<Item = &'b FulfillmentError<'tcx>>,
+    ) {
+        if let Some(tracked_obligations) = &self.tracked_obligations {
+            tracked_obligations.borrow_mut().extend(
+                errors.into_iter().map(|error| FulfilledObligation::Failure(error.clone())),
+            );
+        }
+    }
+
+    fn track_fulfillment_success(&self, predicate: &PredicateObligation<'tcx>) {
+        if let Some(tracked_obligations) = &self.tracked_obligations {
+            tracked_obligations.borrow_mut().push(FulfilledObligation::Success(predicate.clone()));
+        }
     }
 }
 
@@ -52,7 +82,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 +112,11 @@ impl<'tcx> TraitEngine<'tcx> for FulfillmentCtxt<'tcx> {
                     root_obligation: obligation,
                 }
             })
-            .collect()
+            .collect();
+
+        self.track_fulfillment_errors(&errors);
+
+        errors
     }
 
     fn select_where_possible(&mut self, infcx: &InferCtxt<'tcx>) -> Vec<FulfillmentError<'tcx>> {
@@ -97,7 +132,10 @@ impl<'tcx> TraitEngine<'tcx> for FulfillmentCtxt<'tcx> {
                 let goal = obligation.clone().into();
                 let (changed, certainty, nested_goals) =
                     match infcx.evaluate_root_goal(goal, GenerateProofTree::IfEnabled).0 {
-                        Ok(result) => result,
+                        Ok(result) => {
+                            self.track_fulfillment_success(&obligation);
+                            result
+                        }
                         Err(NoSolution) => {
                             errors.push(FulfillmentError {
                                 obligation: obligation.clone(),
@@ -176,6 +214,8 @@ impl<'tcx> TraitEngine<'tcx> for FulfillmentCtxt<'tcx> {
             }
         }
 
+        self.track_fulfillment_errors(&errors);
+
         errors
     }
 

tests/ui/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}]")
+        }
+    }
+}