Avoid O(n²) when constructing AttrSourceMap

Author: jonas-schievink

crates/hir_def/src/attr.rs

@@ -16,7 +16,7 @@ use mbe::ast_to_token_tree;
 use smallvec::{smallvec, SmallVec};
 use syntax::{
     ast::{self, AstNode, AttrsOwner},
-    match_ast, AstToken, SmolStr, SyntaxNode, TextRange, TextSize,
+    match_ast, AstPtr, AstToken, SmolStr, SyntaxNode, TextRange, TextSize,
 };
 use tt::Subtree;
 
@@ -403,21 +403,27 @@ impl AttrsWithOwner {
                 return AttrSourceMap { attrs };
             }
             AttrDefId::FieldId(id) => {
-                id.parent.child_source(db).map(|source| match &source[id.local_id] {
-                    Either::Left(field) => ast::AttrsOwnerNode::new(field.clone()),
-                    Either::Right(field) => ast::AttrsOwnerNode::new(field.clone()),
-                })
+                let map = db.fields_attrs_source_map(id.parent);
+                let file_id = id.parent.file_id(db);
+                let root = db.parse_or_expand(file_id).unwrap();
+                let owner = match &map[id.local_id] {
+                    Either::Left(it) => ast::AttrsOwnerNode::new(it.to_node(&root)),
+                    Either::Right(it) => ast::AttrsOwnerNode::new(it.to_node(&root)),
+                };
+                InFile::new(file_id, owner)
             }
             AttrDefId::AdtId(adt) => match adt {
                 AdtId::StructId(id) => id.lookup(db).source(db).map(ast::AttrsOwnerNode::new),
                 AdtId::UnionId(id) => id.lookup(db).source(db).map(ast::AttrsOwnerNode::new),
                 AdtId::EnumId(id) => id.lookup(db).source(db).map(ast::AttrsOwnerNode::new),
             },
             AttrDefId::FunctionId(id) => id.lookup(db).source(db).map(ast::AttrsOwnerNode::new),
-            AttrDefId::EnumVariantId(id) => id
-                .parent
-                .child_source(db)
-                .map(|source| ast::AttrsOwnerNode::new(source[id.local_id].clone())),
+            AttrDefId::EnumVariantId(id) => {
+                let map = db.variants_attrs_source_map(id.parent);
+                let file_id = id.parent.lookup(db).id.file_id();
+                let root = db.parse_or_expand(file_id).unwrap();
+                InFile::new(file_id, ast::AttrsOwnerNode::new(map[id.local_id].to_node(&root)))
+            }
             AttrDefId::StaticId(id) => id.lookup(db).source(db).map(ast::AttrsOwnerNode::new),
             AttrDefId::ConstId(id) => id.lookup(db).source(db).map(ast::AttrsOwnerNode::new),
             AttrDefId::TraitId(id) => id.lookup(db).source(db).map(ast::AttrsOwnerNode::new),
@@ -746,3 +752,36 @@ fn collect_attrs(
 
     attrs.into_iter().map(|(_, attr)| attr)
 }
+
+pub(crate) fn variants_attrs_source_map(
+    db: &dyn DefDatabase,
+    def: EnumId,
+) -> Arc<ArenaMap<LocalEnumVariantId, AstPtr<ast::Variant>>> {
+    let mut res = ArenaMap::default();
+    let child_source = def.child_source(db);
+
+    for (idx, variant) in child_source.value.iter() {
+        res.insert(idx, AstPtr::new(variant));
+    }
+
+    Arc::new(res)
+}
+
+pub(crate) fn fields_attrs_source_map(
+    db: &dyn DefDatabase,
+    def: VariantId,
+) -> Arc<ArenaMap<LocalFieldId, Either<AstPtr<ast::TupleField>, AstPtr<ast::RecordField>>>> {
+    let mut res = ArenaMap::default();
+    let child_source = def.child_source(db);
+
+    for (idx, variant) in child_source.value.iter() {
+        res.insert(
+            idx,
+            variant
+                .as_ref()
+                .either(|l| Either::Left(AstPtr::new(l)), |r| Either::Right(AstPtr::new(r))),
+        );
+    }
+
+    Arc::new(res)
+}

crates/hir_def/src/db.rs

@@ -2,9 +2,10 @@
 use std::sync::Arc;
 
 use base_db::{salsa, CrateId, SourceDatabase, Upcast};
+use either::Either;
 use hir_expand::{db::AstDatabase, HirFileId};
 use la_arena::ArenaMap;
-use syntax::SmolStr;
+use syntax::{ast, AstPtr, SmolStr};
 
 use crate::{
     adt::{EnumData, StructData},
@@ -122,6 +123,18 @@ pub trait DefDatabase: InternDatabase + AstDatabase + Upcast<dyn AstDatabase> {
     #[salsa::invoke(Attrs::fields_attrs_query)]
     fn fields_attrs(&self, def: VariantId) -> Arc<ArenaMap<LocalFieldId, Attrs>>;
 
+    #[salsa::invoke(crate::attr::variants_attrs_source_map)]
+    fn variants_attrs_source_map(
+        &self,
+        def: EnumId,
+    ) -> Arc<ArenaMap<LocalEnumVariantId, AstPtr<ast::Variant>>>;
+
+    #[salsa::invoke(crate::attr::fields_attrs_source_map)]
+    fn fields_attrs_source_map(
+        &self,
+        def: VariantId,
+    ) -> Arc<ArenaMap<LocalFieldId, Either<AstPtr<ast::TupleField>, AstPtr<ast::RecordField>>>>;
+
     #[salsa::invoke(AttrsWithOwner::attrs_query)]
     fn attrs(&self, def: AttrDefId) -> AttrsWithOwner;
 

crates/hir_def/src/lib.rs

@@ -453,6 +453,14 @@ impl VariantId {
             }
         }
     }
+
+    pub fn file_id(self, db: &dyn db::DefDatabase) -> HirFileId {
+        match self {
+            VariantId::EnumVariantId(it) => it.parent.lookup(db).id.file_id(),
+            VariantId::StructId(it) => it.lookup(db).id.file_id(),
+            VariantId::UnionId(it) => it.lookup(db).id.file_id(),
+        }
+    }
 }
 
 trait Intern {

crates/ide_db/src/apply_change.rs

@@ -152,6 +152,10 @@ impl RootDatabase {
             hir::db::FileItemTreeQuery
             hir::db::BlockDefMapQuery
             hir::db::CrateDefMapQueryQuery
+            hir::db::FieldsAttrsQuery
+            hir::db::VariantsAttrsQuery
+            hir::db::FieldsAttrsSourceMapQuery
+            hir::db::VariantsAttrsSourceMapQuery
             hir::db::StructDataQuery
             hir::db::UnionDataQuery
             hir::db::EnumDataQuery