Auto merge of #68837 - jonas-schievink:assoc-item-lookup-2, r=<try>

Make associated item collection a query

Before this change, every time associated items were iterated over (which rustc does *a lot* – this can probably be further optimized), there would be N+1 queries to fetch all assoc. items. Now there's just one after they've been computed once.

Author: bors

src/librustc/query/mod.rs

@@ -310,6 +310,11 @@ rustc_queries! {
         /// Maps from a trait item to the trait item "descriptor".
         query associated_item(_: DefId) -> ty::AssocItem {}
 
+        /// Collects the associated items defined on a trait or impl.
+        query associated_items(key: DefId) -> ty::AssocItemsIterator<'tcx> {
+            desc { |tcx| "collecting associated items of {}", tcx.def_path_str(key) }
+        }
+
         query impl_trait_ref(_: DefId) -> Option<ty::TraitRef<'tcx>> {}
         query impl_polarity(_: DefId) -> ty::ImplPolarity {}
 

src/librustc/ty/mod.rs

@@ -2746,19 +2746,6 @@ impl<'tcx> TyCtxt<'tcx> {
         variant.fields.iter().position(|field| self.hygienic_eq(ident, field.ident, variant.def_id))
     }
 
-    pub fn associated_items(self, def_id: DefId) -> AssocItemsIterator<'tcx> {
-        // Ideally, we would use `-> impl Iterator` here, but it falls
-        // afoul of the conservative "capture [restrictions]" we put
-        // in place, so we use a hand-written iterator.
-        //
-        // [restrictions]: https://github.com/rust-lang/rust/issues/34511#issuecomment-373423999
-        AssocItemsIterator {
-            tcx: self,
-            def_ids: self.associated_item_def_ids(def_id),
-            next_index: 0,
-        }
-    }
-
     /// Returns `true` if the impls are the same polarity and the trait either
     /// has no items or is annotated #[marker] and prevents item overrides.
     pub fn impls_are_allowed_to_overlap(
@@ -2998,20 +2985,22 @@ impl<'tcx> TyCtxt<'tcx> {
     }
 }
 
-#[derive(Clone)]
+#[derive(Copy, Clone, HashStable)]
 pub struct AssocItemsIterator<'tcx> {
-    tcx: TyCtxt<'tcx>,
-    def_ids: &'tcx [DefId],
-    next_index: usize,
+    pub items: &'tcx [AssocItem],
 }
 
-impl Iterator for AssocItemsIterator<'_> {
+impl<'tcx> Iterator for AssocItemsIterator<'tcx> {
     type Item = AssocItem;
 
+    #[inline]
     fn next(&mut self) -> Option<AssocItem> {
-        let def_id = self.def_ids.get(self.next_index)?;
-        self.next_index += 1;
-        Some(self.tcx.associated_item(*def_id))
+        if let Some((first, rest)) = self.items.split_first() {
+            self.items = rest;
+            Some(*first)
+        } else {
+            None
+        }
     }
 }
 

src/librustc_ty/ty.rs

@@ -206,6 +206,14 @@ fn associated_item_def_ids(tcx: TyCtxt<'_>, def_id: DefId) -> &[DefId] {
     }
 }
 
+fn associated_items<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> ty::AssocItemsIterator<'tcx> {
+    ty::AssocItemsIterator {
+        items: tcx.arena.alloc_from_iter(
+            tcx.associated_item_def_ids(def_id).iter().map(|did| tcx.associated_item(*did)),
+        ),
+    }
+}
+
 fn def_span(tcx: TyCtxt<'_>, def_id: DefId) -> Span {
     tcx.hir().span_if_local(def_id).unwrap()
 }
@@ -356,6 +364,7 @@ pub fn provide(providers: &mut ty::query::Providers<'_>) {
         asyncness,
         associated_item,
         associated_item_def_ids,
+        associated_items,
         adt_sized_constraint,
         def_span,
         param_env,