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)