Auto merge of #69072 - ecstatic-morse:associated-items, r=petrochenkov

O(log n) lookup of associated items by name

Resolves #68957, in which compile time is quadratic in the number of associated items. This PR makes name lookup use binary search instead of a linear scan to improve its asymptotic performance. As a result, the pathological case from that issue now runs in 8 seconds on my local machine, as opposed to many minutes on the current stable.

Currently, method resolution must do a linear scan through all associated items of a type to find one with a certain name. This PR changes the result of the `associated_items` query to a data structure that preserves the definition order of associated items (which is used, e.g., for the layout of trait object vtables) while adding an index of those items sorted by (unhygienic) name. When doing name lookup, we first find all items with the same `Symbol` using binary search, then run hygienic comparison to find the one we are looking for. Ideally, this would be implemented using an insertion-order preserving, hash-based multi-map, but one is not readily available.

Someone who is more familiar with identifier hygiene could probably make this better by auditing the uses of the `AssociatedItems` interface. My goal was to preserve the current behavior exactly, even if it seemed strange (I left at least one FIXME to this effect). For example, some places use comparison with `ident.modern()` and some places use `tcx.hygienic_eq` which requires the `DefId` of the containing `impl`. I don't know whether those approaches are equivalent or which one should be preferred.

Author: bors

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

@@ -333,7 +333,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,10 +30,11 @@ 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;
-use rustc_hir::def::{CtorKind, CtorOf, DefKind, Res};
+use rustc_hir::def::{CtorKind, CtorOf, DefKind, Namespace, Res};
 use rustc_hir::def_id::{CrateNum, DefId, DefIdMap, LocalDefId, CRATE_DEF_INDEX, LOCAL_CRATE};
 use rustc_hir::{Constness, GlobMap, Node, TraitMap};
 use rustc_index::vec::{Idx, IndexVec};
@@ -216,6 +217,13 @@ impl AssocKind {
             ty::AssocKind::Const => "associated constant",
         }
     }
+
+    pub fn namespace(&self) -> Namespace {
+        match *self {
+            ty::AssocKind::OpaqueTy | ty::AssocKind::Type => Namespace::TypeNS,
+            ty::AssocKind::Const | ty::AssocKind::Method => Namespace::ValueNS,
+        }
+    }
 }
 
 impl AssocItem {
@@ -257,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: impl IntoIterator<Item = 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).
@@ -2731,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_data_structures/sorted_map.rs

@@ -4,6 +4,10 @@ use std::iter::FromIterator;
 use std::mem;
 use std::ops::{Bound, Index, IndexMut, RangeBounds};
 
+mod index_map;
+
+pub use index_map::SortedIndexMultiMap;
+
 /// `SortedMap` is a data structure with similar characteristics as BTreeMap but
 /// slightly different trade-offs: lookup, insertion, and removal are O(log(N))
 /// and elements can be iterated in order cheaply.