make for_all_relevant_impls O(1) again

A change in #41911 had made `for_all_relevant_impls` do a linear scan over
all impls, instead of using an HashMap. Use an HashMap again to avoid
quadratic blowup when there is a large number of structs with impls.

I think this fixes #43141 completely, but I want better measurements in
order to be sure. As a perf patch, please don't roll this up.

Author: arielb1

src/librustc/traits/error_reporting.rs

@@ -278,8 +278,8 @@ impl<'a, 'gcx, 'tcx> InferCtxt<'a, 'gcx, 'tcx> {
         let mut self_match_impls = vec![];
         let mut fuzzy_match_impls = vec![];
 
-        self.tcx.trait_def(trait_ref.def_id)
-            .for_each_relevant_impl(self.tcx, trait_self_ty, |def_id| {
+        self.tcx.for_each_relevant_impl(
+            trait_ref.def_id, trait_self_ty, |def_id| {
                 let impl_substs = self.fresh_substs_for_item(obligation.cause.span, def_id);
                 let impl_trait_ref = tcx
                     .impl_trait_ref(def_id)
@@ -396,10 +396,9 @@ impl<'a, 'gcx, 'tcx> InferCtxt<'a, 'gcx, 'tcx> {
                                               trait_ref.skip_binder().self_ty(),
                                               true);
         let mut impl_candidates = Vec::new();
-        let trait_def = self.tcx.trait_def(trait_ref.def_id());
 
         match simp {
-            Some(simp) => trait_def.for_each_impl(self.tcx, |def_id| {
+            Some(simp) => self.tcx.for_each_impl(trait_ref.def_id(), |def_id| {
                 let imp = self.tcx.impl_trait_ref(def_id).unwrap();
                 let imp_simp = fast_reject::simplify_type(self.tcx,
                                                           imp.self_ty(),
@@ -411,7 +410,7 @@ impl<'a, 'gcx, 'tcx> InferCtxt<'a, 'gcx, 'tcx> {
                 }
                 impl_candidates.push(imp);
             }),
-            None => trait_def.for_each_impl(self.tcx, |def_id| {
+            None => self.tcx.for_each_impl(trait_ref.def_id(), |def_id| {
                 impl_candidates.push(
                     self.tcx.impl_trait_ref(def_id).unwrap());
             })

src/librustc/traits/select.rs

@@ -1430,10 +1430,8 @@ impl<'cx, 'gcx, 'tcx> SelectionContext<'cx, 'gcx, 'tcx> {
     {
         debug!("assemble_candidates_from_impls(obligation={:?})", obligation);
 
-        let def = self.tcx().trait_def(obligation.predicate.def_id());
-
-        def.for_each_relevant_impl(
-            self.tcx(),
+        self.tcx().for_each_relevant_impl(
+            obligation.predicate.def_id(),
             obligation.predicate.0.trait_ref.self_ty(),
             |impl_def_id| {
                 self.probe(|this, snapshot| { /* [1] */

src/librustc/traits/specialize/mod.rs

@@ -300,7 +300,8 @@ pub(super) fn specialization_graph_provider<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx
                                                       -> Rc<specialization_graph::Graph> {
     let mut sg = specialization_graph::Graph::new();
 
-    let mut trait_impls: Vec<DefId> = tcx.trait_impls_of(trait_id).iter().collect();
+    let mut trait_impls = Vec::new();
+    tcx.for_each_impl(trait_id, |impl_did| trait_impls.push(impl_did));
 
     // The coherence checking implementation seems to rely on impls being
     // iterated over (roughly) in definition order, so we are sorting by

src/librustc/ty/maps.rs

@@ -470,12 +470,6 @@ impl<'tcx> QueryDescription for queries::trait_impls_of<'tcx> {
     }
 }
 
-impl<'tcx> QueryDescription for queries::relevant_trait_impls_for<'tcx> {
-    fn describe(tcx: TyCtxt, (def_id, ty): (DefId, SimplifiedType)) -> String {
-        format!("relevant impls for: `({}, {:?})`", tcx.item_path_str(def_id), ty)
-    }
-}
-
 impl<'tcx> QueryDescription for queries::is_object_safe<'tcx> {
     fn describe(tcx: TyCtxt, def_id: DefId) -> String {
         format!("determine object safety of trait `{}`", tcx.item_path_str(def_id))
@@ -966,10 +960,7 @@ define_maps! { <'tcx>
     [] const_is_rvalue_promotable_to_static: ConstIsRvaluePromotableToStatic(DefId) -> bool,
     [] is_mir_available: IsMirAvailable(DefId) -> bool,
 
-    [] trait_impls_of: TraitImpls(DefId) -> ty::trait_def::TraitImpls,
-    // Note that TraitDef::for_each_relevant_impl() will do type simplication for you.
-    [] relevant_trait_impls_for: relevant_trait_impls_for((DefId, SimplifiedType))
-        -> ty::trait_def::TraitImpls,
+    [] trait_impls_of: TraitImpls(DefId) -> Rc<ty::trait_def::TraitImpls>,
     [] specialization_graph_of: SpecializationGraph(DefId) -> Rc<specialization_graph::Graph>,
     [] is_object_safe: ObjectSafety(DefId) -> bool,
 
@@ -1045,10 +1036,6 @@ fn crate_variances<'tcx>(_: CrateNum) -> DepConstructor<'tcx> {
     DepConstructor::CrateVariances
 }
 
-fn relevant_trait_impls_for<'tcx>((def_id, t): (DefId, SimplifiedType)) -> DepConstructor<'tcx> {
-    DepConstructor::RelevantTraitImpls(def_id, t)
-}
-
 fn is_copy_dep_node<'tcx>(key: ty::ParamEnvAnd<'tcx, Ty<'tcx>>) -> DepConstructor<'tcx> {
     let def_id = ty::item_path::characteristic_def_id_of_type(key.value)
         .unwrap_or(DefId::local(CRATE_DEF_INDEX));

src/librustc/ty/mod.rs

@@ -2530,7 +2530,6 @@ pub fn provide(providers: &mut ty::maps::Providers) {
         param_env,
         trait_of_item,
         trait_impls_of: trait_def::trait_impls_of_provider,
-        relevant_trait_impls_for: trait_def::relevant_trait_impls_provider,
         ..*providers
     };
 }
@@ -2540,7 +2539,6 @@ pub fn provide_extern(providers: &mut ty::maps::Providers) {
         adt_sized_constraint,
         adt_dtorck_constraint,
         trait_impls_of: trait_def::trait_impls_of_provider,
-        relevant_trait_impls_for: trait_def::relevant_trait_impls_provider,
         param_env,
         ..*providers
     };

src/librustc/ty/trait_def.rs

@@ -8,14 +8,17 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
+use hir;
 use hir::def_id::DefId;
 use hir::map::DefPathHash;
 use traits::specialization_graph;
 use ty::fast_reject;
 use ty::fold::TypeFoldable;
 use ty::{Ty, TyCtxt};
+
+use rustc_data_structures::fx::FxHashMap;
+
 use std::rc::Rc;
-use hir;
 
 /// A trait's definition with type information.
 pub struct TraitDef {
@@ -36,60 +39,12 @@ pub struct TraitDef {
     pub def_path_hash: DefPathHash,
 }
 
-// We don't store the list of impls in a flat list because each cached list of
-// `relevant_impls_for` we would then duplicate all blanket impls. By keeping
-// blanket and non-blanket impls separate, we can share the list of blanket
-// impls.
-#[derive(Clone)]
 pub struct TraitImpls {
-    blanket_impls: Rc<Vec<DefId>>,
-    non_blanket_impls: Rc<Vec<DefId>>,
+    blanket_impls: Vec<DefId>,
+    /// Impls indexed by their simplified self-type, for fast lookup.
+    non_blanket_impls: FxHashMap<fast_reject::SimplifiedType, Vec<DefId>>,
 }
 
-impl TraitImpls {
-    pub fn iter(&self) -> TraitImplsIter {
-        TraitImplsIter {
-            blanket_impls: self.blanket_impls.clone(),
-            non_blanket_impls: self.non_blanket_impls.clone(),
-            index: 0
-        }
-    }
-}
-
-#[derive(Clone)]
-pub struct TraitImplsIter {
-    blanket_impls: Rc<Vec<DefId>>,
-    non_blanket_impls: Rc<Vec<DefId>>,
-    index: usize,
-}
-
-impl Iterator for TraitImplsIter {
-    type Item = DefId;
-
-    fn next(&mut self) -> Option<DefId> {
-        if self.index < self.blanket_impls.len() {
-            let bi_index = self.index;
-            self.index += 1;
-            Some(self.blanket_impls[bi_index])
-        } else {
-            let nbi_index = self.index - self.blanket_impls.len();
-            if nbi_index < self.non_blanket_impls.len() {
-                self.index += 1;
-                Some(self.non_blanket_impls[nbi_index])
-            } else {
-                None
-            }
-        }
-    }
-
-    fn size_hint(&self) -> (usize, Option<usize>) {
-        let items_left = (self.blanket_impls.len() + self.non_blanket_impls.len()) - self.index;
-        (items_left, Some(items_left))
-    }
-}
-
-impl ExactSizeIterator for TraitImplsIter {}
-
 impl<'a, 'gcx, 'tcx> TraitDef {
     pub fn new(def_id: DefId,
                unsafety: hir::Unsafety,
@@ -111,20 +66,36 @@ impl<'a, 'gcx, 'tcx> TraitDef {
                      -> specialization_graph::Ancestors {
         specialization_graph::ancestors(tcx, self.def_id, of_impl)
     }
+}
 
-    pub fn for_each_impl<F: FnMut(DefId)>(&self, tcx: TyCtxt<'a, 'gcx, 'tcx>, mut f: F) {
-        for impl_def_id in tcx.trait_impls_of(self.def_id).iter() {
+impl<'a, 'gcx, 'tcx> TyCtxt<'a, 'gcx, 'tcx> {
+    pub fn for_each_impl<F: FnMut(DefId)>(self, def_id: DefId, mut f: F) {
+        let impls = self.trait_impls_of(def_id);
+
+        for &impl_def_id in impls.blanket_impls.iter() {
             f(impl_def_id);
         }
+
+        for v in impls.non_blanket_impls.values() {
+            for &impl_def_id in v {
+                f(impl_def_id);
+            }
+        }
     }
 
     /// Iterate over every impl that could possibly match the
     /// self-type `self_ty`.
-    pub fn for_each_relevant_impl<F: FnMut(DefId)>(&self,
-                                                   tcx: TyCtxt<'a, 'gcx, 'tcx>,
+    pub fn for_each_relevant_impl<F: FnMut(DefId)>(self,
+                                                   def_id: DefId,
                                                    self_ty: Ty<'tcx>,
                                                    mut f: F)
     {
+        let impls = self.trait_impls_of(def_id);
+
+        for &impl_def_id in impls.blanket_impls.iter() {
+            f(impl_def_id);
+        }
+
         // simplify_type(.., false) basically replaces type parameters and
         // projections with infer-variables. This is, of course, done on
         // the impl trait-ref when it is instantiated, but not on the
@@ -137,23 +108,39 @@ impl<'a, 'gcx, 'tcx> TraitDef {
         // replace `S` with anything - this impl of course can't be
         // selected, and as there are hundreds of similar impls,
         // considering them would significantly harm performance.
-        let relevant_impls = if let Some(simplified_self_ty) =
-                fast_reject::simplify_type(tcx, self_ty, true) {
-            tcx.relevant_trait_impls_for((self.def_id, simplified_self_ty))
-        } else {
-            tcx.trait_impls_of(self.def_id)
-        };
 
-        for impl_def_id in relevant_impls.iter() {
-            f(impl_def_id);
+        // This depends on the set of all impls for the trait. That is
+        // unfortunate. When we get red-green recompilation, we would like
+        // to have a way of knowing whether the set of relevant impls
+        // changed. The most naive
+        // way would be to compute the Vec of relevant impls and see whether
+        // it differs between compilations. That shouldn't be too slow by
+        // itself - we do quite a bit of work for each relevant impl anyway.
+        //
+        // If we want to be faster, we could have separate queries for
+        // blanket and non-blanket impls, and compare them separately.
+        //
+        // I think we'll cross that bridge when we get to it.
+        if let Some(simp) = fast_reject::simplify_type(self, self_ty, true) {
+            if let Some(impls) = impls.non_blanket_impls.get(&simp) {
+                for &impl_def_id in impls {
+                    f(impl_def_id);
+                }
+            }
+        } else {
+            for v in impls.non_blanket_impls.values() {
+                for &impl_def_id in v {
+                    f(impl_def_id);
+                }
+            }
         }
     }
 }
 
 // Query provider for `trait_impls_of`.
 pub(super) fn trait_impls_of_provider<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
                                                 trait_id: DefId)
-                                                -> TraitImpls {
+                                                -> Rc<TraitImpls> {
     let remote_impls = if trait_id.is_local() {
         // Traits defined in the current crate can't have impls in upstream
         // crates, so we don't bother querying the cstore.
@@ -163,7 +150,7 @@ pub(super) fn trait_impls_of_provider<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
     };
 
     let mut blanket_impls = Vec::new();
-    let mut non_blanket_impls = Vec::new();
+    let mut non_blanket_impls = FxHashMap();
 
     let local_impls = tcx.hir
                          .trait_impls(trait_id)
@@ -176,47 +163,20 @@ pub(super) fn trait_impls_of_provider<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
             continue
         }
 
-        if fast_reject::simplify_type(tcx, impl_self_ty, false).is_some() {
-            non_blanket_impls.push(impl_def_id);
+        if let Some(simplified_self_ty) =
+            fast_reject::simplify_type(tcx, impl_self_ty, false)
+        {
+            non_blanket_impls
+                .entry(simplified_self_ty)
+                .or_insert(vec![])
+                .push(impl_def_id);
         } else {
             blanket_impls.push(impl_def_id);
         }
     }
 
-    TraitImpls {
-        blanket_impls: Rc::new(blanket_impls),
-        non_blanket_impls: Rc::new(non_blanket_impls),
-    }
-}
-
-// Query provider for `relevant_trait_impls_for`.
-pub(super) fn relevant_trait_impls_provider<'a, 'tcx>(
-    tcx: TyCtxt<'a, 'tcx, 'tcx>,
-    (trait_id, self_ty): (DefId, fast_reject::SimplifiedType))
-    -> TraitImpls
-{
-    let all_trait_impls = tcx.trait_impls_of(trait_id);
-
-    let relevant: Vec<DefId> = all_trait_impls
-        .non_blanket_impls
-        .iter()
-        .cloned()
-        .filter(|&impl_def_id| {
-            let impl_self_ty = tcx.type_of(impl_def_id);
-            let impl_simple_self_ty = fast_reject::simplify_type(tcx,
-                                                                 impl_self_ty,
-                                                                 false).unwrap();
-            impl_simple_self_ty == self_ty
-        })
-        .collect();
-
-    if all_trait_impls.non_blanket_impls.len() == relevant.len() {
-        // If we didn't filter anything out, re-use the existing vec.
-        all_trait_impls
-    } else {
-        TraitImpls {
-            blanket_impls: all_trait_impls.blanket_impls.clone(),
-            non_blanket_impls: Rc::new(relevant),
-        }
-    }
+    Rc::new(TraitImpls {
+        blanket_impls: blanket_impls,
+        non_blanket_impls: non_blanket_impls,
+    })
 }

src/librustc/ty/util.rs

@@ -422,7 +422,7 @@ impl<'a, 'gcx, 'tcx> TyCtxt<'a, 'gcx, 'tcx> {
 
         let mut dtor_did = None;
         let ty = self.type_of(adt_did);
-        self.trait_def(drop_trait).for_each_relevant_impl(self, ty, |impl_did| {
+        self.for_each_relevant_impl(drop_trait, ty, |impl_did| {
             if let Some(item) = self.associated_items(impl_did).next() {
                 if let Ok(()) = validate(self, impl_did) {
                     dtor_did = Some(item.def_id);

src/librustc_lint/builtin.rs

@@ -542,9 +542,8 @@ impl<'a, 'tcx> LateLintPass<'a, 'tcx> for MissingDebugImplementations {
         };
 
         if self.impling_types.is_none() {
-            let debug_def = cx.tcx.trait_def(debug);
             let mut impls = NodeSet();
-            debug_def.for_each_impl(cx.tcx, |d| {
+            cx.tcx.for_each_impl(debug, |d| {
                 if let Some(ty_def) = cx.tcx.type_of(d).ty_to_def_id() {
                     if let Some(node_id) = cx.tcx.hir.as_local_node_id(ty_def) {
                         impls.insert(node_id);

src/librustc_mir/transform/qualify_consts.rs

@@ -244,8 +244,7 @@ impl<'a, 'tcx> Qualifier<'a, 'tcx, 'tcx> {
                 let mut span = None;
 
                 self.tcx
-                    .trait_def(drop_trait_id)
-                    .for_each_relevant_impl(self.tcx, self.mir.return_ty, |impl_did| {
+                    .for_each_relevant_impl(drop_trait_id, self.mir.return_ty, |impl_did| {
                         self.tcx.hir
                             .as_local_node_id(impl_did)
                             .and_then(|impl_node_id| self.tcx.hir.find(impl_node_id))

src/librustc_typeck/check/method/probe.rs

@@ -718,10 +718,8 @@ impl<'a, 'gcx, 'tcx> ProbeContext<'a, 'gcx, 'tcx> {
                                                      import_id: Option<ast::NodeId>,
                                                      trait_def_id: DefId,
                                                      item: ty::AssociatedItem) {
-        let trait_def = self.tcx.trait_def(trait_def_id);
-
         // FIXME(arielb1): can we use for_each_relevant_impl here?
-        trait_def.for_each_impl(self.tcx, |impl_def_id| {
+        self.tcx.for_each_impl(trait_def_id, |impl_def_id| {
             debug!("assemble_extension_candidates_for_trait_impl: trait_def_id={:?} \
                                                                   impl_def_id={:?}",
                    trait_def_id,