Auto merge of #89831 - Aaron1011:project-caching-speedup, r=jackh726

Re-introduce concept of projection cache 'completion'

Instead of clearing out the cache entirely, we store
the intermediate evaluation result into the cache entry.
This accomplishes several things:

* We avoid the performance hit associated with re-evaluating
  the sub-obligations
* We avoid causing issues with incremental compilation, since
  the final evaluation result is always the same
* We avoid affecting other uses of the same `InferCtxt` which
  might care about 'side effects' from processing the sub-obligations
  (e,g. region constraints). Only code that is specifically aware
   of the new 'complete' code is affected

Author: bors

compiler/rustc_infer/src/traits/project.rs

@@ -10,7 +10,7 @@ use rustc_data_structures::{
 };
 use rustc_middle::ty::{self, Ty};
 
-pub use rustc_middle::traits::Reveal;
+pub use rustc_middle::traits::{EvaluationResult, Reveal};
 
 pub(crate) type UndoLog<'tcx> =
     snapshot_map::UndoLog<ProjectionCacheKey<'tcx>, ProjectionCacheEntry<'tcx>>;
@@ -92,7 +92,42 @@ pub enum ProjectionCacheEntry<'tcx> {
     Ambiguous,
     Recur,
     Error,
-    NormalizedTy(NormalizedTy<'tcx>),
+    NormalizedTy {
+        ty: NormalizedTy<'tcx>,
+        /// If we were able to successfully evaluate the
+        /// corresponding cache entry key during predicate
+        /// evaluation, then this field stores the final
+        /// result obtained from evaluating all of the projection
+        /// sub-obligations. During evaluation, we will skip
+        /// evaluating the cached sub-obligations in `ty`
+        /// if this field is set. Evaluation only
+        /// cares about the final result, so we don't
+        /// care about any region constraint side-effects
+        /// produced by evaluating the sub-boligations.
+        ///
+        /// Additionally, we will clear out the sub-obligations
+        /// entirely if we ever evaluate the cache entry (along
+        /// with all its sub obligations) to `EvaluatedToOk`.
+        /// This affects all users of the cache, not just evaluation.
+        /// Since a result of `EvaluatedToOk` means that there were
+        /// no region obligations that need to be tracked, it's
+        /// fine to forget about the sub-obligations - they
+        /// don't provide any additional information. However,
+        /// we do *not* discard any obligations when we see
+        /// `EvaluatedToOkModuloRegions` - we don't know
+        /// which sub-obligations may introduce region constraints,
+        /// so we keep them all to be safe.
+        ///
+        /// When we are not performing evaluation
+        /// (e.g. in `FulfillmentContext`), we ignore this field,
+        /// and always re-process the cached sub-obligations
+        /// (which may have been cleared out - see the above
+        /// paragraph).
+        /// This ensures that we do not lose any regions
+        /// constraints that arise from processing the
+        /// sub-obligations.
+        complete: Option<EvaluationResult>,
+    },
 }
 
 impl<'tcx> ProjectionCacheStorage<'tcx> {
@@ -149,10 +184,41 @@ impl<'tcx> ProjectionCache<'_, 'tcx> {
             debug!("Not overwriting Recur");
             return;
         }
-        let fresh_key = map.insert(key, ProjectionCacheEntry::NormalizedTy(value));
+        let fresh_key =
+            map.insert(key, ProjectionCacheEntry::NormalizedTy { ty: value, complete: None });
         assert!(!fresh_key, "never started projecting `{:?}`", key);
     }
 
+    /// Mark the relevant projection cache key as having its derived obligations
+    /// complete, so they won't have to be re-computed (this is OK to do in a
+    /// snapshot - if the snapshot is rolled back, the obligations will be
+    /// marked as incomplete again).
+    pub fn complete(&mut self, key: ProjectionCacheKey<'tcx>, result: EvaluationResult) {
+        let mut map = self.map();
+        match map.get(&key) {
+            Some(&ProjectionCacheEntry::NormalizedTy { ref ty, complete: _ }) => {
+                info!("ProjectionCacheEntry::complete({:?}) - completing {:?}", key, ty);
+                let mut ty = ty.clone();
+                if result == EvaluationResult::EvaluatedToOk {
+                    ty.obligations = vec![];
+                }
+                map.insert(key, ProjectionCacheEntry::NormalizedTy { ty, complete: Some(result) });
+            }
+            ref value => {
+                // Type inference could "strand behind" old cache entries. Leave
+                // them alone for now.
+                info!("ProjectionCacheEntry::complete({:?}) - ignoring {:?}", key, value);
+            }
+        };
+    }
+
+    pub fn is_complete(&mut self, key: ProjectionCacheKey<'tcx>) -> Option<EvaluationResult> {
+        self.map().get(&key).and_then(|res| match res {
+            ProjectionCacheEntry::NormalizedTy { ty: _, complete } => *complete,
+            _ => None,
+        })
+    }
+
     /// Indicates that trying to normalize `key` resulted in
     /// ambiguity. No point in trying it again then until we gain more
     /// type information (in which case, the "fully resolved" key will

compiler/rustc_trait_selection/src/lib.rs

@@ -16,6 +16,7 @@
 #![feature(drain_filter)]
 #![feature(derive_default_enum)]
 #![feature(hash_drain_filter)]
+#![feature(label_break_value)]
 #![feature(let_else)]
 #![feature(never_type)]
 #![feature(crate_visibility_modifier)]

compiler/rustc_trait_selection/src/traits/fulfill.rs

@@ -4,6 +4,7 @@ use rustc_data_structures::obligation_forest::ProcessResult;
 use rustc_data_structures::obligation_forest::{Error, ForestObligation, Outcome};
 use rustc_data_structures::obligation_forest::{ObligationForest, ObligationProcessor};
 use rustc_errors::ErrorReported;
+use rustc_infer::traits::ProjectionCacheKey;
 use rustc_infer::traits::{SelectionError, TraitEngine, TraitEngineExt as _, TraitObligation};
 use rustc_middle::mir::interpret::ErrorHandled;
 use rustc_middle::thir::abstract_const::NotConstEvaluatable;
@@ -20,12 +21,14 @@ use super::wf;
 use super::CodeAmbiguity;
 use super::CodeProjectionError;
 use super::CodeSelectionError;
+use super::EvaluationResult;
 use super::Unimplemented;
 use super::{FulfillmentError, FulfillmentErrorCode};
 use super::{ObligationCause, PredicateObligation};
 
 use crate::traits::error_reporting::InferCtxtExt as _;
 use crate::traits::project::PolyProjectionObligation;
+use crate::traits::project::ProjectionCacheKeyExt as _;
 use crate::traits::query::evaluate_obligation::InferCtxtExt as _;
 
 impl<'tcx> ForestObligation for PendingPredicateObligation<'tcx> {
@@ -709,6 +712,20 @@ impl<'a, 'b, 'tcx> FulfillProcessor<'a, 'b, 'tcx> {
             // no type variables present, can use evaluation for better caching.
             // FIXME: consider caching errors too.
             if self.selcx.infcx().predicate_must_hold_considering_regions(obligation) {
+                if let Some(key) = ProjectionCacheKey::from_poly_projection_predicate(
+                    &mut self.selcx,
+                    project_obligation.predicate,
+                ) {
+                    // If `predicate_must_hold_considering_regions` succeeds, then we've
+                    // evaluated all sub-obligations. We can therefore mark the 'root'
+                    // obligation as complete, and skip evaluating sub-obligations.
+                    self.selcx
+                        .infcx()
+                        .inner
+                        .borrow_mut()
+                        .projection_cache()
+                        .complete(key, EvaluationResult::EvaluatedToOk);
+                }
                 return ProcessResult::Changed(vec![]);
             } else {
                 tracing::debug!("Does NOT hold: {:?}", obligation);

compiler/rustc_trait_selection/src/traits/project.rs

@@ -889,7 +889,7 @@ fn opt_normalize_projection_type<'a, 'b, 'tcx>(
             debug!("recur cache");
             return Err(InProgress);
         }
-        Err(ProjectionCacheEntry::NormalizedTy(ty)) => {
+        Err(ProjectionCacheEntry::NormalizedTy { ty, complete: _ }) => {
             // This is the hottest path in this function.
             //
             // If we find the value in the cache, then return it along

compiler/rustc_trait_selection/src/traits/select/mod.rs

@@ -25,6 +25,8 @@ use super::{ObligationCause, PredicateObligation, TraitObligation};
 
 use crate::infer::{InferCtxt, InferOk, TypeFreshener};
 use crate::traits::error_reporting::InferCtxtExt;
+use crate::traits::project::ProjectionCacheKeyExt;
+use crate::traits::ProjectionCacheKey;
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_data_structures::stack::ensure_sufficient_stack;
 use rustc_data_structures::sync::Lrc;
@@ -550,8 +552,54 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
                     let project_obligation = obligation.with(data);
                     match project::poly_project_and_unify_type(self, &project_obligation) {
                         Ok(Ok(Some(mut subobligations))) => {
-                            self.add_depth(subobligations.iter_mut(), obligation.recursion_depth);
-                            self.evaluate_predicates_recursively(previous_stack, subobligations)
+                            'compute_res: {
+                                // If we've previously marked this projection as 'complete', thne
+                                // use the final cached result (either `EvaluatedToOk` or
+                                // `EvaluatedToOkModuloRegions`), and skip re-evaluating the
+                                // sub-obligations.
+                                if let Some(key) =
+                                    ProjectionCacheKey::from_poly_projection_predicate(self, data)
+                                {
+                                    if let Some(cached_res) = self
+                                        .infcx
+                                        .inner
+                                        .borrow_mut()
+                                        .projection_cache()
+                                        .is_complete(key)
+                                    {
+                                        break 'compute_res Ok(cached_res);
+                                    }
+                                }
+
+                                self.add_depth(
+                                    subobligations.iter_mut(),
+                                    obligation.recursion_depth,
+                                );
+                                let res = self.evaluate_predicates_recursively(
+                                    previous_stack,
+                                    subobligations,
+                                );
+                                if let Ok(res) = res {
+                                    if res == EvaluatedToOk || res == EvaluatedToOkModuloRegions {
+                                        if let Some(key) =
+                                            ProjectionCacheKey::from_poly_projection_predicate(
+                                                self, data,
+                                            )
+                                        {
+                                            // If the result is something that we can cache, then mark this
+                                            // entry as 'complete'. This will allow us to skip evaluating the
+                                            // suboligations at all the next time we evaluate the projection
+                                            // predicate.
+                                            self.infcx
+                                                .inner
+                                                .borrow_mut()
+                                                .projection_cache()
+                                                .complete(key, res);
+                                        }
+                                    }
+                                }
+                                res
+                            }
                         }
                         Ok(Ok(None)) => Ok(EvaluatedToAmbig),
                         Ok(Err(project::InProgress)) => Ok(EvaluatedToRecur),