Merge #7162

7162: Introduce queries to avoid problems when performing completion for enums with many variants r=matklad a=danielframpton

This change introduces two new queries to compute:
  1) attributes for all variants of an enum, and
  2) attributes for all fields of a variant.

The purpose of this change is to avoid the current n^2 behavior when rendering completion for variants (which prevents completion for enums with large numbers of variants).

Co-authored-by: Daniel Frampton <[email protected]>

2 files changed, 55 insertions, 12 deletions

diff --git a/crates/hir_def/src/attr.rs b/crates/hir_def/src/attr.rs
index 6b79e7bad..9e6426b31 100644
--- a/crates/hir_def/src/attr.rs
+++ b/crates/hir_def/src/attr.rs
@@ -2,6 +2,7 @@
 
 use std::{ops, sync::Arc};
 
+use arena::map::ArenaMap;
 use base_db::CrateId;
 use cfg::{CfgExpr, CfgOptions};
 use either::Either;
@@ -21,7 +22,8 @@ use crate::{
     nameres::ModuleSource,
     path::{ModPath, PathKind},
     src::HasChildSource,
-    AdtId, AttrDefId, GenericParamId, Lookup,
+    AdtId, AttrDefId, EnumId, GenericParamId, HasModule, LocalEnumVariantId, LocalFieldId, Lookup,
+    VariantId,
 };
 
 /// Holds documentation
@@ -210,16 +212,10 @@ impl Attrs {
                 }
             }
             AttrDefId::FieldId(it) => {
-                let src = it.parent.child_source(db);
-                match &src.value[it.local_id] {
-                    Either::Left(_tuple) => RawAttrs::default(),
-                    Either::Right(record) => RawAttrs::from_attrs_owner(db, src.with_value(record)),
-                }
+                return db.fields_attrs(it.parent)[it.local_id].clone();
             }
-            AttrDefId::EnumVariantId(var_id) => {
-                let src = var_id.parent.child_source(db);
-                let src = src.as_ref().map(|it| &it[var_id.local_id]);
-                RawAttrs::from_attrs_owner(db, src.map(|it| it as &dyn AttrsOwner))
+            AttrDefId::EnumVariantId(it) => {
+                return db.variants_attrs(it.parent)[it.local_id].clone();
             }
             AttrDefId::AdtId(it) => match it {
                 AdtId::StructId(it) => attrs_from_item_tree(it.lookup(db).id, db),
@@ -259,6 +255,46 @@ impl Attrs {
         raw_attrs.filter(db, def.krate(db))
     }
 
+    pub(crate) fn variants_attrs_query(
+        db: &dyn DefDatabase,
+        e: EnumId,
+    ) -> Arc<ArenaMap<LocalEnumVariantId, Attrs>> {
+        let krate = e.lookup(db).container.module(db).krate;
+        let src = e.child_source(db);
+        let mut res = ArenaMap::default();
+
+        for (id, var) in src.value.iter() {
+            let attrs = RawAttrs::from_attrs_owner(db, src.with_value(var as &dyn AttrsOwner))
+                .filter(db, krate);
+
+            res.insert(id, attrs)
+        }
+
+        Arc::new(res)
+    }
+
+    pub(crate) fn fields_attrs_query(
+        db: &dyn DefDatabase,
+        v: VariantId,
+    ) -> Arc<ArenaMap<LocalFieldId, Attrs>> {
+        let krate = v.module(db).krate;
+        let src = v.child_source(db);
+        let mut res = ArenaMap::default();
+
+        for (id, fld) in src.value.iter() {
+            let attrs = match fld {
+                Either::Left(_tuple) => Attrs::default(),
+                Either::Right(record) => {
+                    RawAttrs::from_attrs_owner(db, src.with_value(record)).filter(db, krate)
+                }
+            };
+
+            res.insert(id, attrs);
+        }
+
+        Arc::new(res)
+    }
+
     pub fn by_key(&self, key: &'static str) -> AttrQuery<'_> {
         AttrQuery { attrs: self, key }
     }
diff --git a/crates/hir_def/src/db.rs b/crates/hir_def/src/db.rs
index d1a459066..d3bf5b34c 100644
--- a/crates/hir_def/src/db.rs
+++ b/crates/hir_def/src/db.rs
@@ -1,6 +1,7 @@
 //! Defines database & queries for name resolution.
 use std::sync::Arc;
 
+use arena::map::ArenaMap;
 use base_db::{salsa, CrateId, SourceDatabase, Upcast};
 use hir_expand::{db::AstDatabase, HirFileId};
 use syntax::SmolStr;
@@ -16,8 +17,8 @@ use crate::{
     lang_item::{LangItemTarget, LangItems},
     nameres::CrateDefMap,
     AttrDefId, ConstId, ConstLoc, DefWithBodyId, EnumId, EnumLoc, FunctionId, FunctionLoc,
-    GenericDefId, ImplId, ImplLoc, StaticId, StaticLoc, StructId, StructLoc, TraitId, TraitLoc,
-    TypeAliasId, TypeAliasLoc, UnionId, UnionLoc,
+    GenericDefId, ImplId, ImplLoc, LocalEnumVariantId, LocalFieldId, StaticId, StaticLoc, StructId,
+    StructLoc, TraitId, TraitLoc, TypeAliasId, TypeAliasLoc, UnionId, UnionLoc, VariantId,
 };
 
 #[salsa::query_group(InternDatabaseStorage)]
@@ -92,6 +93,12 @@ pub trait DefDatabase: InternDatabase + AstDatabase + Upcast<dyn AstDatabase> {
     #[salsa::invoke(GenericParams::generic_params_query)]
     fn generic_params(&self, def: GenericDefId) -> Arc<GenericParams>;
 
+    #[salsa::invoke(Attrs::variants_attrs_query)]
+    fn variants_attrs(&self, def: EnumId) -> Arc<ArenaMap<LocalEnumVariantId, Attrs>>;
+
+    #[salsa::invoke(Attrs::fields_attrs_query)]
+    fn fields_attrs(&self, def: VariantId) -> Arc<ArenaMap<LocalFieldId, Attrs>>;
+
     #[salsa::invoke(Attrs::attrs_query)]
     fn attrs(&self, def: AttrDefId) -> Attrs;