Make lookup of associated item by name O(log n)

Author: ecstatic-morse

src/librustc/arena.rs

@@ -48,6 +48,7 @@ macro_rules! arena_types {
             [] item_local_set: rustc_hir::ItemLocalSet,
             [decode] mir_const_qualif: rustc_index::bit_set::BitSet<rustc::mir::Local>,
             [] trait_impls_of: rustc::ty::trait_def::TraitImpls,
+            [] associated_items: rustc::ty::AssociatedItems,
             [] dropck_outlives:
                 rustc::infer::canonical::Canonical<'tcx,
                     rustc::infer::canonical::QueryResponse<'tcx,

src/librustc/query/mod.rs

@@ -325,7 +325,7 @@ rustc_queries! {
         query associated_item(_: DefId) -> ty::AssocItem {}
 
         /// Collects the associated items defined on a trait or impl.
-        query associated_items(key: DefId) -> &'tcx [ty::AssocItem] {
+        query associated_items(key: DefId) -> &'tcx ty::AssociatedItems {
             desc { |tcx| "collecting associated items of {}", tcx.def_path_str(key) }
         }
 

src/librustc/traits/specialization_graph.rs

@@ -81,8 +81,8 @@ impl<'tcx> Node {
     }
 
     /// Iterate over the items defined directly by the given (impl or trait) node.
-    pub fn items(&self, tcx: TyCtxt<'tcx>) -> &'tcx [ty::AssocItem] {
-        tcx.associated_items(self.def_id())
+    pub fn items(&self, tcx: TyCtxt<'tcx>) -> impl 'tcx + Iterator<Item = &'tcx ty::AssocItem> {
+        tcx.associated_items(self.def_id()).in_definition_order()
     }
 
     /// Finds an associated item defined in this node.
@@ -99,7 +99,7 @@ impl<'tcx> Node {
         use crate::ty::AssocKind::*;
 
         tcx.associated_items(self.def_id())
-            .iter()
+            .filter_by_name_unhygienic(trait_item_name.name)
             .find(move |impl_item| {
                 match (trait_item_kind, impl_item.kind) {
                 | (Const, Const)

src/librustc/ty/adjustment.rs

@@ -122,7 +122,7 @@ impl<'tcx> OverloadedDeref<'tcx> {
         };
         let method_def_id = tcx
             .associated_items(trait_def_id.unwrap())
-            .iter()
+            .in_definition_order()
             .find(|m| m.kind == ty::AssocKind::Method)
             .unwrap()
             .def_id;

src/librustc/ty/instance.rs

@@ -337,7 +337,7 @@ impl<'tcx> Instance<'tcx> {
         let fn_once = tcx.lang_items().fn_once_trait().unwrap();
         let call_once = tcx
             .associated_items(fn_once)
-            .iter()
+            .in_definition_order()
             .find(|it| it.kind == ty::AssocKind::Method)
             .unwrap()
             .def_id;

src/librustc/ty/mod.rs

@@ -30,6 +30,7 @@ use rustc_attr as attr;
 use rustc_data_structures::captures::Captures;
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::fx::FxIndexMap;
+use rustc_data_structures::sorted_map::SortedIndexMultiMap;
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 use rustc_data_structures::sync::{self, par_iter, Lrc, ParallelIterator};
 use rustc_hir as hir;
@@ -264,6 +265,81 @@ impl AssocItem {
     }
 }
 
+/// A list of `ty::AssocItem`s in definition order that allows for efficient lookup by name.
+///
+/// When doing lookup by name, we try to postpone hygienic comparison for as long as possible since
+/// it is relatively expensive. Instead, items are indexed by `Symbol` and hygienic comparison is
+/// done only on items with the same name.
+#[derive(Debug, Clone, PartialEq, HashStable)]
+pub struct AssociatedItems {
+    items: SortedIndexMultiMap<u32, Symbol, ty::AssocItem>,
+}
+
+impl AssociatedItems {
+    /// Constructs an `AssociatedItems` map from a series of `ty::AssocItem`s in definition order.
+    pub fn new(items_in_def_order: Vec<ty::AssocItem>) -> Self {
+        let items = items_in_def_order.into_iter().map(|item| (item.ident.name, item)).collect();
+        AssociatedItems { items }
+    }
+
+    /// Returns a slice of associated items in the order they were defined.
+    ///
+    /// New code should avoid relying on definition order. If you need a particular associated item
+    /// for a known trait, make that trait a lang item instead of indexing this array.
+    pub fn in_definition_order(&self) -> impl '_ + Iterator<Item = &ty::AssocItem> {
+        self.items.iter().map(|(_, v)| v)
+    }
+
+    /// Returns an iterator over all associated items with the given name, ignoring hygiene.
+    pub fn filter_by_name_unhygienic(
+        &self,
+        name: Symbol,
+    ) -> impl '_ + Iterator<Item = &ty::AssocItem> {
+        self.items.get_by_key(&name)
+    }
+
+    /// Returns an iterator over all associated items with the given name.
+    ///
+    /// Multiple items may have the same name if they are in different `Namespace`s. For example,
+    /// an associated type can have the same name as a method. Use one of the `find_by_name_and_*`
+    /// methods below if you know which item you are looking for.
+    pub fn filter_by_name(
+        &'a self,
+        tcx: TyCtxt<'a>,
+        ident: Ident,
+        parent_def_id: DefId,
+    ) -> impl 'a + Iterator<Item = &'a ty::AssocItem> {
+        self.filter_by_name_unhygienic(ident.name)
+            .filter(move |item| tcx.hygienic_eq(ident, item.ident, parent_def_id))
+    }
+
+    /// Returns the associated item with the given name and `AssocKind`, if one exists.
+    pub fn find_by_name_and_kind(
+        &self,
+        tcx: TyCtxt<'_>,
+        ident: Ident,
+        kind: AssocKind,
+        parent_def_id: DefId,
+    ) -> Option<&ty::AssocItem> {
+        self.filter_by_name_unhygienic(ident.name)
+            .filter(|item| item.kind == kind)
+            .find(|item| tcx.hygienic_eq(ident, item.ident, parent_def_id))
+    }
+
+    /// Returns the associated item with the given name in the given `Namespace`, if one exists.
+    pub fn find_by_name_and_namespace(
+        &self,
+        tcx: TyCtxt<'_>,
+        ident: Ident,
+        ns: Namespace,
+        parent_def_id: DefId,
+    ) -> Option<&ty::AssocItem> {
+        self.filter_by_name_unhygienic(ident.name)
+            .filter(|item| item.kind.namespace() == ns)
+            .find(|item| tcx.hygienic_eq(ident, item.ident, parent_def_id))
+    }
+}
+
 #[derive(Clone, Debug, PartialEq, Eq, Copy, RustcEncodable, RustcDecodable, HashStable)]
 pub enum Visibility {
     /// Visible everywhere (including in other crates).
@@ -2738,14 +2814,14 @@ impl<'tcx> TyCtxt<'tcx> {
             .for_each(|&body_id| f(self.hir().body_owner_def_id(body_id)));
     }
 
-    pub fn provided_trait_methods(self, id: DefId) -> impl Iterator<Item = &'tcx AssocItem> {
+    pub fn provided_trait_methods(self, id: DefId) -> impl 'tcx + Iterator<Item = &'tcx AssocItem> {
         self.associated_items(id)
-            .iter()
+            .in_definition_order()
             .filter(|item| item.kind == AssocKind::Method && item.defaultness.has_value())
     }
 
     pub fn trait_relevant_for_never(self, did: DefId) -> bool {
-        self.associated_items(did).iter().any(|item| item.relevant_for_never())
+        self.associated_items(did).in_definition_order().any(|item| item.relevant_for_never())
     }
 
     pub fn opt_item_name(self, def_id: DefId) -> Option<Ident> {

src/librustc/ty/sty.rs

@@ -1066,11 +1066,7 @@ impl<'tcx> ProjectionTy<'tcx> {
     ) -> ProjectionTy<'tcx> {
         let item_def_id = tcx
             .associated_items(trait_ref.def_id)
-            .iter()
-            .find(|item| {
-                item.kind == ty::AssocKind::Type
-                    && tcx.hygienic_eq(item_name, item.ident, trait_ref.def_id)
-            })
+            .find_by_name_and_kind(tcx, item_name, ty::AssocKind::Type, trait_ref.def_id)
             .unwrap()
             .def_id;
 

src/librustc/ty/util.rs

@@ -357,7 +357,7 @@ impl<'tcx> TyCtxt<'tcx> {
         let mut dtor_did = None;
         let ty = self.type_of(adt_did);
         self.for_each_relevant_impl(drop_trait, ty, |impl_did| {
-            if let Some(item) = self.associated_items(impl_did).first() {
+            if let Some(item) = self.associated_items(impl_did).in_definition_order().nth(0) {
                 if validate(self, impl_did).is_ok() {
                     dtor_did = Some(item.def_id);
                 }

src/librustc_infer/traits/mod.rs

@@ -537,7 +537,7 @@ fn vtable_methods<'tcx>(
     tcx.arena.alloc_from_iter(supertraits(tcx, trait_ref).flat_map(move |trait_ref| {
         let trait_methods = tcx
             .associated_items(trait_ref.def_id())
-            .iter()
+            .in_definition_order()
             .filter(|item| item.kind == ty::AssocKind::Method);
 
         // Now list each method's DefId and InternalSubsts (for within its trait).

src/librustc_infer/traits/object_safety.rs

@@ -213,7 +213,7 @@ fn object_safety_violations_for_trait(
     // Check methods for violations.
     let mut violations: Vec<_> = tcx
         .associated_items(trait_def_id)
-        .iter()
+        .in_definition_order()
         .filter(|item| item.kind == ty::AssocKind::Method)
         .filter_map(|item| {
             object_safety_violation_for_method(tcx, trait_def_id, &item)
@@ -289,7 +289,7 @@ fn object_safety_violations_for_trait(
 
     violations.extend(
         tcx.associated_items(trait_def_id)
-            .iter()
+            .in_definition_order()
             .filter(|item| item.kind == ty::AssocKind::Const)
             .map(|item| ObjectSafetyViolation::AssocConst(item.ident.name, item.ident.span)),
     );
@@ -646,7 +646,7 @@ fn object_ty_for_trait<'tcx>(
     let mut associated_types = traits::supertraits(tcx, ty::Binder::dummy(trait_ref))
         .flat_map(|super_trait_ref| {
             tcx.associated_items(super_trait_ref.def_id())
-                .iter()
+                .in_definition_order()
                 .map(move |item| (super_trait_ref, item))
         })
         .filter(|(_, item)| item.kind == ty::AssocKind::Type)

src/librustc_infer/traits/util.rs

@@ -586,7 +586,7 @@ pub fn count_own_vtable_entries(tcx: TyCtxt<'tcx>, trait_ref: ty::PolyTraitRef<'
     let mut entries = 0;
     // Count number of methods and add them to the total offset.
     // Skip over associated types and constants.
-    for trait_item in tcx.associated_items(trait_ref.def_id()) {
+    for trait_item in tcx.associated_items(trait_ref.def_id()).in_definition_order() {
         if trait_item.kind == ty::AssocKind::Method {
             entries += 1;
         }
@@ -606,7 +606,7 @@ pub fn get_vtable_index_of_object_method<N>(
     // add them to the total offset.
     // Skip over associated types and constants.
     let mut entries = object.vtable_base;
-    for trait_item in tcx.associated_items(object.upcast_trait_ref.def_id()) {
+    for trait_item in tcx.associated_items(object.upcast_trait_ref.def_id()).in_definition_order() {
         if trait_item.def_id == method_def_id {
             // The item with the ID we were given really ought to be a method.
             assert_eq!(trait_item.kind, ty::AssocKind::Method);

src/librustc_infer/traits/wf.rs

@@ -318,7 +318,10 @@ impl<'a, 'tcx> WfPredicates<'a, 'tcx> {
             };
 
         if let Elaborate::All = elaborate {
-            let trait_assoc_items = tcx.associated_items(trait_ref.def_id);
+            // FIXME: Make `extend_cause_with_original_assoc_item_obligation` take an iterator
+            // instead of a slice.
+            let trait_assoc_items: Vec<_> =
+                tcx.associated_items(trait_ref.def_id).in_definition_order().copied().collect();
 
             let predicates = obligations.iter().map(|obligation| obligation.predicate).collect();
             let implied_obligations = traits::elaborate_predicates(tcx, predicates);
@@ -327,7 +330,7 @@ impl<'a, 'tcx> WfPredicates<'a, 'tcx> {
                 extend_cause_with_original_assoc_item_obligation(
                     &mut cause,
                     &pred,
-                    trait_assoc_items,
+                    &*trait_assoc_items,
                 );
                 traits::Obligation::new(cause, param_env, pred)
             });

src/librustc_mir/shim.rs

@@ -68,7 +68,7 @@ fn make_shim<'tcx>(tcx: TyCtxt<'tcx>, instance: ty::InstanceDef<'tcx>) -> &'tcx
             let fn_mut = tcx.lang_items().fn_mut_trait().unwrap();
             let call_mut = tcx
                 .associated_items(fn_mut)
-                .iter()
+                .in_definition_order()
                 .find(|it| it.kind == ty::AssocKind::Method)
                 .unwrap()
                 .def_id;

src/librustc_mir/util/elaborate_drops.rs

@@ -539,7 +539,7 @@ where
         debug!("destructor_call_block({:?}, {:?})", self, succ);
         let tcx = self.tcx();
         let drop_trait = tcx.lang_items().drop_trait().unwrap();
-        let drop_fn = tcx.associated_items(drop_trait)[0];
+        let drop_fn = tcx.associated_items(drop_trait).in_definition_order().nth(0).unwrap();
         let ty = self.place_ty(self.place);
         let substs = tcx.mk_substs_trait(ty, &[]);
 

src/librustc_mir_build/hair/cx/mod.rs

@@ -167,17 +167,19 @@ impl<'a, 'tcx> Cx<'a, 'tcx> {
         params: &[GenericArg<'tcx>],
     ) -> &'tcx ty::Const<'tcx> {
         let substs = self.tcx.mk_substs_trait(self_ty, params);
-        for item in self.tcx.associated_items(trait_def_id) {
-            // The unhygienic comparison here is acceptable because this is only
-            // used on known traits.
-            if item.kind == ty::AssocKind::Method && item.ident.name == method_name {
-                let method_ty = self.tcx.type_of(item.def_id);
-                let method_ty = method_ty.subst(self.tcx, substs);
-                return ty::Const::zero_sized(self.tcx, method_ty);
-            }
-        }
 
-        bug!("found no method `{}` in `{:?}`", method_name, trait_def_id);
+        // The unhygienic comparison here is acceptable because this is only
+        // used on known traits.
+        let item = self
+            .tcx
+            .associated_items(trait_def_id)
+            .filter_by_name_unhygienic(method_name)
+            .find(|item| item.kind == ty::AssocKind::Method)
+            .expect("trait method not found");
+
+        let method_ty = self.tcx.type_of(item.def_id);
+        let method_ty = method_ty.subst(self.tcx, substs);
+        ty::Const::zero_sized(self.tcx, method_ty)
     }
 
     crate fn all_fields(&mut self, adt_def: &ty::AdtDef, variant_index: VariantIdx) -> Vec<Field> {

src/librustc_passes/stability.rs

@@ -468,8 +468,8 @@ impl Visitor<'tcx> for Checker<'tcx> {
                         let trait_item_def_id = self
                             .tcx
                             .associated_items(trait_did)
-                            .iter()
-                            .find(|item| item.ident.name == impl_item.ident.name)
+                            .filter_by_name_unhygienic(impl_item.ident.name)
+                            .next()
                             .map(|item| item.def_id);
                         if let Some(def_id) = trait_item_def_id {
                             // Pass `None` to skip deprecation warnings.

src/librustc_save_analysis/lib.rs

@@ -408,7 +408,6 @@ impl<'l, 'tcx> SaveContext<'l, 'tcx> {
                         qualname.push_str(&self.tcx.hir().hir_to_pretty_string(self_ty.hir_id));
 
                         let trait_id = self.tcx.trait_id_of_impl(impl_id);
-                        let mut decl_id = None;
                         let mut docs = String::new();
                         let mut attrs = vec![];
                         let hir_id = self.tcx.hir().node_to_hir_id(id);
@@ -417,15 +416,18 @@ impl<'l, 'tcx> SaveContext<'l, 'tcx> {
                             attrs = item.attrs.to_vec();
                         }
 
+                        let mut decl_id = None;
                         if let Some(def_id) = trait_id {
                             // A method in a trait impl.
                             qualname.push_str(" as ");
                             qualname.push_str(&self.tcx.def_path_str(def_id));
-                            self.tcx
+
+                            decl_id = self
+                                .tcx
                                 .associated_items(def_id)
-                                .iter()
-                                .find(|item| item.ident.name == ident.name)
-                                .map(|item| decl_id = Some(item.def_id));
+                                .filter_by_name_unhygienic(ident.name)
+                                .next()
+                                .map(|item| item.def_id);
                         }
                         qualname.push_str(">");
 
@@ -716,12 +718,11 @@ impl<'l, 'tcx> SaveContext<'l, 'tcx> {
             Res::Def(HirDefKind::Method, decl_id) => {
                 let def_id = if decl_id.is_local() {
                     let ti = self.tcx.associated_item(decl_id);
+
                     self.tcx
                         .associated_items(ti.container.id())
-                        .iter()
-                        .find(|item| {
-                            item.ident.name == ti.ident.name && item.defaultness.has_value()
-                        })
+                        .filter_by_name_unhygienic(ti.ident.name)
+                        .find(|item| item.defaultness.has_value())
                         .map(|item| item.def_id)
                 } else {
                     None

src/librustc_ty/ty.rs

@@ -210,10 +210,10 @@ fn associated_item_def_ids(tcx: TyCtxt<'_>, def_id: DefId) -> &[DefId] {
     }
 }
 
-fn associated_items<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> &'tcx [ty::AssocItem] {
-    tcx.arena.alloc_from_iter(
-        tcx.associated_item_def_ids(def_id).iter().map(|did| tcx.associated_item(*did)),
-    )
+fn associated_items<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> &'tcx ty::AssociatedItems {
+    let items =
+        tcx.associated_item_def_ids(def_id).iter().map(|did| tcx.associated_item(*did)).collect();
+    tcx.arena.alloc(ty::AssociatedItems::new(items))
 }
 
 fn def_span(tcx: TyCtxt<'_>, def_id: DefId) -> Span {