dedup obligations in opt_normalize_projection_type

Author: the8472

compiler/rustc_infer/src/lib.rs

@@ -17,6 +17,7 @@
 #![feature(box_patterns)]
 #![feature(derive_default_enum)]
 #![feature(extend_one)]
+#![feature(hash_raw_entry)]
 #![feature(let_else)]
 #![feature(never_type)]
 #![feature(control_flow_enum)]

compiler/rustc_infer/src/traits/mod.rs

@@ -8,10 +8,14 @@ mod project;
 mod structural_impls;
 pub mod util;
 
+use rustc_data_structures::fx::FxHashMap;
 use rustc_hir as hir;
 use rustc_middle::ty::error::{ExpectedFound, TypeError};
 use rustc_middle::ty::{self, Const, Ty, TyCtxt};
 use rustc_span::Span;
+use std::borrow::{Borrow, Cow};
+use std::collections::hash_map::RawEntryMut;
+use std::hash::Hash;
 
 pub use self::FulfillmentErrorCode::*;
 pub use self::ImplSource::*;
@@ -105,6 +109,65 @@ pub enum FulfillmentErrorCode<'tcx> {
     CodeAmbiguity,
 }
 
+pub struct ObligationsDedup<'a, 'tcx, T> {
+    obligations: &'a mut Vec<Obligation<'tcx, T>>,
+}
+
+impl<'a, 'tcx, T: 'tcx> ObligationsDedup<'a, 'tcx, T>
+where
+    T: Clone + Hash + Eq,
+{
+    pub fn from_vec(vec: &'a mut Vec<Obligation<'tcx, T>>) -> Self {
+        ObligationsDedup { obligations: vec }
+    }
+
+    pub fn extend<'b>(&mut self, iter: impl ExactSizeIterator<Item = Cow<'b, Obligation<'tcx, T>>>)
+    where
+        'tcx: 'b,
+    {
+        // obligation tracing has shown that initial batches added to an empty vec do not
+        // contain any duplicates, so there's no need to attempt deduplication
+        if self.obligations.is_empty() {
+            *self.obligations = iter.into_iter().map(Cow::into_owned).collect();
+            return;
+        }
+
+        let initial_size = self.obligations.len();
+        let iter = iter.into_iter();
+        let expected_new = iter.len();
+        let combined_size = initial_size + expected_new;
+
+        if combined_size <= 16 || combined_size < initial_size.next_power_of_two() {
+            // small case/not crossing a power of two. don't bother with dedup
+            self.obligations.extend(iter.map(Cow::into_owned));
+        } else {
+            // crossing power of two threshold. this would incur a vec growth anyway if we didn't do
+            // anything. piggyback a dedup on that
+            let obligations = std::mem::take(self.obligations);
+
+            let mut seen = FxHashMap::default();
+            seen.reserve(initial_size);
+
+            *self.obligations = obligations
+                .into_iter()
+                .map(Cow::Owned)
+                .chain(iter)
+                .filter_map(|obligation| {
+                    match seen.raw_entry_mut().from_key(obligation.borrow()) {
+                        RawEntryMut::Occupied(..) => {
+                            return None;
+                        }
+                        RawEntryMut::Vacant(vacant) => {
+                            vacant.insert(obligation.clone().into_owned(), ());
+                        }
+                    }
+                    Some(obligation.into_owned())
+                })
+                .collect();
+        }
+    }
+}
+
 impl<'tcx, O> Obligation<'tcx, O> {
     pub fn new(
         cause: ObligationCause<'tcx>,

compiler/rustc_infer/src/traits/project.rs

@@ -173,6 +173,20 @@ impl<'tcx> ProjectionCache<'_, 'tcx> {
         Ok(())
     }
 
+    pub fn try_start_borrowed<'a, T>(
+        &'a mut self,
+        key: ProjectionCacheKey<'tcx>,
+        with: impl FnOnce(&'_ ProjectionCacheEntry<'tcx>) -> T + 'a,
+    ) -> Option<T> {
+        let mut map = self.map();
+        if let Some(entry) = map.get(&key) {
+            return Some(with(entry));
+        }
+
+        map.insert(key, ProjectionCacheEntry::InProgress);
+        None
+    }
+
     /// Indicates that `key` was normalized to `value`.
     pub fn insert_ty(&mut self, key: ProjectionCacheKey<'tcx>, value: NormalizedTy<'tcx>) {
         debug!(

compiler/rustc_trait_selection/src/traits/project.rs

@@ -30,8 +30,10 @@ use rustc_middle::ty::subst::Subst;
 use rustc_middle::ty::{self, ToPredicate, Ty, TyCtxt};
 use rustc_span::symbol::sym;
 
+use std::borrow::Cow;
 use std::collections::BTreeMap;
 
+use crate::traits::ObligationsDedup;
 pub use rustc_middle::traits::Reveal;
 
 pub type PolyProjectionObligation<'tcx> = Obligation<'tcx, ty::PolyProjectionPredicate<'tcx>>;
@@ -839,6 +841,8 @@ fn opt_normalize_projection_type<'a, 'b, 'tcx>(
     // mode, which could lead to using incorrect cache results.
     let use_cache = !selcx.is_intercrate();
 
+    let mut obligations = ObligationsDedup::from_vec(obligations);
+
     let projection_ty = infcx.resolve_vars_if_possible(projection_ty);
     let cache_key = ProjectionCacheKey::new(projection_ty);
 
@@ -850,65 +854,76 @@ fn opt_normalize_projection_type<'a, 'b, 'tcx>(
     // bounds. It might be the case that we want two distinct caches,
     // or else another kind of cache entry.
 
-    let cache_result = if use_cache {
-        infcx.inner.borrow_mut().projection_cache().try_start(cache_key)
-    } else {
-        Ok(())
-    };
-    match cache_result {
-        Ok(()) => debug!("no cache"),
-        Err(ProjectionCacheEntry::Ambiguous) => {
-            // If we found ambiguity the last time, that means we will continue
-            // to do so until some type in the key changes (and we know it
-            // hasn't, because we just fully resolved it).
-            debug!("found cache entry: ambiguous");
-            return Ok(None);
-        }
-        Err(ProjectionCacheEntry::InProgress) => {
-            // Under lazy normalization, this can arise when
-            // bootstrapping.  That is, imagine an environment with a
-            // where-clause like `A::B == u32`. Now, if we are asked
-            // to normalize `A::B`, we will want to check the
-            // where-clauses in scope. So we will try to unify `A::B`
-            // with `A::B`, which can trigger a recursive
-            // normalization.
-
-            debug!("found cache entry: in-progress");
-
-            // Cache that normalizing this projection resulted in a cycle. This
-            // should ensure that, unless this happens within a snapshot that's
-            // rolled back, fulfillment or evaluation will notice the cycle.
+    if use_cache {
+        let result =
+            infcx.inner.borrow_mut().projection_cache().try_start_borrowed(cache_key, |cached| {
+                match cached {
+                    ProjectionCacheEntry::NormalizedTy { ty, complete: _ } => {
+                        // This is the hottest path in this function.
+                        //
+                        // If we find the value in the cache, then return it along
+                        // with the obligations that went along with it. Note
+                        // that, when using a fulfillment context, these
+                        // obligations could in principle be ignored: they have
+                        // already been registered when the cache entry was
+                        // created (and hence the new ones will quickly be
+                        // discarded as duplicated). But when doing trait
+                        // evaluation this is not the case, and dropping the trait
+                        // evaluations can causes ICEs (e.g., #43132).
+                        debug!(?ty, "found normalized ty");
+                        obligations.extend(ty.obligations.iter().map(Cow::Borrowed));
+                        Ok(Some(ty.value))
+                    }
+                    cached @ _ => Err(cached.clone()),
+                }
+            });
 
-            if use_cache {
-                infcx.inner.borrow_mut().projection_cache().recur(cache_key);
+        match result {
+            Some(Ok(ret)) => return Ok(ret),
+            Some(Err(cached)) => {
+                return match cached {
+                    ProjectionCacheEntry::Ambiguous => {
+                        // If we found ambiguity the last time, that means we will continue
+                        // to do so until some type in the key changes (and we know it
+                        // hasn't, because we just fully resolved it).
+                        debug!("found cache entry: ambiguous");
+                        Ok(None)
+                    }
+                    ProjectionCacheEntry::InProgress => {
+                        // Under lazy normalization, this can arise when
+                        // bootstrapping.  That is, imagine an environment with a
+                        // where-clause like `A::B == u32`. Now, if we are asked
+                        // to normalize `A::B`, we will want to check the
+                        // where-clauses in scope. So we will try to unify `A::B`
+                        // with `A::B`, which can trigger a recursive
+                        // normalization.
+
+                        debug!("found cache entry: in-progress");
+
+                        // Cache that normalizing this projection resulted in a cycle. This
+                        // should ensure that, unless this happens within a snapshot that's
+                        // rolled back, fulfillment or evaluation will notice the cycle.
+
+                        if use_cache {
+                            infcx.inner.borrow_mut().projection_cache().recur(cache_key);
+                        }
+                        Err(InProgress)
+                    }
+                    ProjectionCacheEntry::Recur => {
+                        debug!("recur cache");
+                        Err(InProgress)
+                    }
+                    ProjectionCacheEntry::Error => {
+                        debug!("opt_normalize_projection_type: found error");
+                        let result =
+                            normalize_to_error(selcx, param_env, projection_ty, cause, depth);
+                        obligations.extend(result.obligations.into_iter().map(Cow::Owned));
+                        Ok(Some(result.value))
+                    }
+                    _ => unreachable!("unexpected variant"),
+                };
             }
-            return Err(InProgress);
-        }
-        Err(ProjectionCacheEntry::Recur) => {
-            debug!("recur cache");
-            return Err(InProgress);
-        }
-        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
-            // with the obligations that went along with it. Note
-            // that, when using a fulfillment context, these
-            // obligations could in principle be ignored: they have
-            // already been registered when the cache entry was
-            // created (and hence the new ones will quickly be
-            // discarded as duplicated). But when doing trait
-            // evaluation this is not the case, and dropping the trait
-            // evaluations can causes ICEs (e.g., #43132).
-            debug!(?ty, "found normalized ty");
-            obligations.extend(ty.obligations);
-            return Ok(Some(ty.value));
-        }
-        Err(ProjectionCacheEntry::Error) => {
-            debug!("opt_normalize_projection_type: found error");
-            let result = normalize_to_error(selcx, param_env, projection_ty, cause, depth);
-            obligations.extend(result.obligations);
-            return Ok(Some(result.value));
+            _ => {}
         }
     }
 
@@ -955,7 +970,7 @@ fn opt_normalize_projection_type<'a, 'b, 'tcx>(
             if use_cache {
                 infcx.inner.borrow_mut().projection_cache().insert_ty(cache_key, result.clone());
             }
-            obligations.extend(result.obligations);
+            obligations.extend(result.obligations.into_iter().map(Cow::Owned));
             Ok(Some(result.value))
         }
         Ok(ProjectedTy::NoProgress(projected_ty)) => {
@@ -985,7 +1000,7 @@ fn opt_normalize_projection_type<'a, 'b, 'tcx>(
                 infcx.inner.borrow_mut().projection_cache().error(cache_key);
             }
             let result = normalize_to_error(selcx, param_env, projection_ty, cause, depth);
-            obligations.extend(result.obligations);
+            obligations.extend(result.obligations.into_iter().map(Cow::Owned));
             Ok(Some(result.value))
         }
     }

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

@@ -2337,6 +2337,10 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
         // This code is hot enough that it's worth avoiding the allocation
         // required for the FxHashSet when possible. Special-casing lengths 0,
         // 1 and 2 covers roughly 75-80% of the cases.
+        //
+        // Ideally we would perform deduplication incrementally in the predicates
+        // loop above to prevent excessive Vec growth but that would require
+        // a Vec::range_retain or similar method.
         if obligations.len() <= 1 {
             // No possibility of duplicates.
         } else if obligations.len() == 2 {