Reduce allocations in attribute collection

Veykril · Veykril · commit cc04cfc982b0 · 2022-01-30T22:18:32.000+01:00
diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs
@@ -664,7 +664,7 @@ fn emit_def_diagnostic(db: &dyn HirDatabase, acc: &mut Vec<AnyDiagnostic>, diag:
                     let attr = node
                         .doc_comments_and_attrs()
                         .nth((*invoc_attr_index) as usize)
-                        .and_then(Either::right)
+                        .and_then(Either::left)
                         .unwrap_or_else(|| panic!("cannot find attribute #{}", invoc_attr_index));
                     (
                         ast_id.with_value(SyntaxNodePtr::from(AstPtr::new(&attr))),
diff --git a/crates/hir_def/src/attr.rs b/crates/hir_def/src/attr.rs
@@ -525,38 +525,36 @@ impl AttrsWithOwner {
 
 fn inner_attributes(
     syntax: &SyntaxNode,
-) -> Option<(impl Iterator<Item = ast::Attr>, impl Iterator<Item = ast::Comment>)> {
-    let (attrs, docs) = match_ast! {
+) -> Option<impl Iterator<Item = Either<ast::Attr, ast::Comment>>> {
+    let node = match_ast! {
         match syntax {
-            ast::SourceFile(it) => (it.attrs(), ast::DocCommentIter::from_syntax_node(it.syntax())),
-            ast::ExternBlock(it) => {
-                let extern_item_list = it.extern_item_list()?;
-                (extern_item_list.attrs(), ast::DocCommentIter::from_syntax_node(extern_item_list.syntax()))
-            },
-            ast::Fn(it) => {
-                let body = it.body()?;
-                let stmt_list = body.stmt_list()?;
-                (stmt_list.attrs(), ast::DocCommentIter::from_syntax_node(body.syntax()))
-            },
-            ast::Impl(it) => {
-                let assoc_item_list = it.assoc_item_list()?;
-                (assoc_item_list.attrs(), ast::DocCommentIter::from_syntax_node(assoc_item_list.syntax()))
-            },
-            ast::Module(it) => {
-                let item_list = it.item_list()?;
-                (item_list.attrs(), ast::DocCommentIter::from_syntax_node(item_list.syntax()))
+            ast::SourceFile(_) => syntax.clone(),
+            ast::ExternBlock(it) => it.extern_item_list()?.syntax().clone(),
+            ast::Fn(it) => it.body()?.stmt_list()?.syntax().clone(),
+            ast::Impl(it) => it.assoc_item_list()?.syntax().clone(),
+            ast::Module(it) => it.item_list()?.syntax().clone(),
+            ast::BlockExpr(it) => {
+                use syntax::SyntaxKind::{BLOCK_EXPR , EXPR_STMT};
+                // Block expressions accept outer and inner attributes, but only when they are the outer
+                // expression of an expression statement or the final expression of another block expression.
+                let may_carry_attributes = matches!(
+                    it.syntax().parent().map(|it| it.kind()),
+                     Some(BLOCK_EXPR | EXPR_STMT)
+                );
+                if !may_carry_attributes {
+                    return None
+                }
+                syntax.clone()
             },
-            // FIXME: BlockExpr's only accept inner attributes in specific cases
-            // Excerpt from the reference:
-            // Block expressions accept outer and inner attributes, but only when they are the outer
-            // expression of an expression statement or the final expression of another block expression.
-            ast::BlockExpr(_it) => return None,
             _ => return None,
         }
     };
-    let attrs = attrs.filter(|attr| attr.kind().is_inner());
-    let docs = docs.filter(|doc| doc.is_inner());
-    Some((attrs, docs))
+
+    let attrs = ast::AttrDocCommentIter::from_syntax_node(&node).filter(|el| match el {
+        Either::Left(attr) => attr.kind().is_inner(),
+        Either::Right(comment) => comment.is_inner(),
+    });
+    Some(attrs)
 }
 
 #[derive(Debug)]
@@ -833,24 +831,16 @@ fn attrs_from_item_tree<N: ItemTreeNode>(id: ItemTreeId<N>, db: &dyn DefDatabase
 fn collect_attrs(
     owner: &dyn ast::HasAttrs,
 ) -> impl Iterator<Item = (AttrId, Either<ast::Attr, ast::Comment>)> {
-    let (inner_attrs, inner_docs) = inner_attributes(owner.syntax())
-        .map_or((None, None), |(attrs, docs)| (Some(attrs), Some(docs)));
-
-    let outer_attrs = owner.attrs().filter(|attr| attr.kind().is_outer());
-    let attrs = outer_attrs
-        .chain(inner_attrs.into_iter().flatten())
-        .map(|attr| (attr.syntax().text_range().start(), Either::Left(attr)));
-
-    let outer_docs =
-        ast::DocCommentIter::from_syntax_node(owner.syntax()).filter(ast::Comment::is_outer);
-    let docs = outer_docs
-        .chain(inner_docs.into_iter().flatten())
-        .map(|docs_text| (docs_text.syntax().text_range().start(), Either::Right(docs_text)));
-    // sort here by syntax node offset because the source can have doc attributes and doc strings be interleaved
-    docs.chain(attrs)
-        .sorted_by_key(|&(offset, _)| offset)
+    let inner_attrs = inner_attributes(owner.syntax()).into_iter().flatten();
+    let outer_attrs =
+        ast::AttrDocCommentIter::from_syntax_node(owner.syntax()).filter(|el| match el {
+            Either::Left(attr) => attr.kind().is_outer(),
+            Either::Right(comment) => comment.is_outer(),
+        });
+    outer_attrs
+        .chain(inner_attrs)
         .enumerate()
-        .map(|(id, (_, attr))| (AttrId { ast_index: id as u32 }, attr))
+        .map(|(id, attr)| (AttrId { ast_index: id as u32 }, attr))
 }
 
 pub(crate) fn variants_attrs_source_map(
diff --git a/crates/hir_def/src/child_by_source.rs b/crates/hir_def/src/child_by_source.rs
@@ -117,7 +117,7 @@ impl ChildBySource for ItemScope {
             |(ast_id, calls)| {
                 let adt = ast_id.to_node(db.upcast());
                 calls.for_each(|(attr_id, calls)| {
-                    if let Some(Either::Right(attr)) =
+                    if let Some(Either::Left(attr)) =
                         adt.doc_comments_and_attrs().nth(attr_id.ast_index as usize)
                     {
                         res[keys::DERIVE_MACRO_CALL].insert(attr, (attr_id, calls.into()));
diff --git a/crates/hir_expand/src/db.rs b/crates/hir_expand/src/db.rs
@@ -157,7 +157,7 @@ pub fn expand_speculative(
             let attr = item
                 .doc_comments_and_attrs()
                 .nth(invoc_attr_index as usize)
-                .and_then(Either::right)?;
+                .and_then(Either::left)?;
             match attr.token_tree() {
                 Some(token_tree) => {
                     let (mut tree, map) = syntax_node_to_token_tree(attr.token_tree()?.syntax());
@@ -323,7 +323,7 @@ fn censor_for_macro_input(loc: &MacroCallLoc, node: &SyntaxNode) -> FxHashSet<Sy
                 ast::Item::cast(node.clone())?
                     .doc_comments_and_attrs()
                     .nth(invoc_attr_index as usize)
-                    .and_then(Either::right)
+                    .and_then(Either::left)
                     .map(|attr| attr.syntax().clone())
                     .into_iter()
                     .collect()
diff --git a/crates/hir_expand/src/hygiene.rs b/crates/hir_expand/src/hygiene.rs
@@ -191,7 +191,7 @@ fn make_hygiene_info(
                 .to_node(db)
                 .doc_comments_and_attrs()
                 .nth(invoc_attr_index as usize)
-                .and_then(Either::right)?
+                .and_then(Either::left)?
                 .token_tree()?;
             Some(InFile::new(ast_id.file_id, tt))
         }
diff --git a/crates/hir_expand/src/lib.rs b/crates/hir_expand/src/lib.rs
@@ -205,7 +205,7 @@ impl HirFileId {
                             .to_node(db)
                             .doc_comments_and_attrs()
                             .nth(invoc_attr_index as usize)
-                            .and_then(Either::right)?
+                            .and_then(Either::left)?
                             .token_tree()?;
                         Some(InFile::new(ast_id.file_id, tt))
                     }
@@ -382,7 +382,7 @@ impl MacroCallKind {
                     .doc_comments_and_attrs()
                     .nth(derive_attr_index as usize)
                     .expect("missing derive")
-                    .expect_right("derive is a doc comment?")
+                    .expect_left("derive is a doc comment?")
                     .syntax()
                     .text_range()
             }
@@ -391,7 +391,7 @@ impl MacroCallKind {
                 .doc_comments_and_attrs()
                 .nth(invoc_attr_index as usize)
                 .expect("missing attribute")
-                .expect_right("attribute macro is a doc comment?")
+                .expect_left("attribute macro is a doc comment?")
                 .syntax()
                 .text_range(),
         };
@@ -483,7 +483,7 @@ impl ExpansionInfo {
                     let attr = item
                         .doc_comments_and_attrs()
                         .nth(*invoc_attr_index as usize)
-                        .and_then(Either::right)?;
+                        .and_then(Either::left)?;
                     match attr.token_tree() {
                         Some(token_tree)
                             if token_tree.syntax().text_range().contains_range(token_range) =>
diff --git a/crates/syntax/src/ast.rs b/crates/syntax/src/ast.rs
@@ -27,8 +27,8 @@ pub use self::{
     operators::{ArithOp, BinaryOp, CmpOp, LogicOp, Ordering, RangeOp, UnaryOp},
     token_ext::{CommentKind, CommentPlacement, CommentShape, IsString, QuoteOffsets, Radix},
     traits::{
-        DocCommentIter, HasArgList, HasAttrs, HasDocComments, HasGenericParams, HasLoopBody,
-        HasModuleItem, HasName, HasTypeBounds, HasVisibility,
+        AttrDocCommentIter, DocCommentIter, HasArgList, HasAttrs, HasDocComments, HasGenericParams,
+        HasLoopBody, HasModuleItem, HasName, HasTypeBounds, HasVisibility,
     },
 };
 
diff --git a/crates/syntax/src/ast/node_ext.rs b/crates/syntax/src/ast/node_ext.rs
@@ -160,14 +160,9 @@ impl ast::Attr {
     }
 
     pub fn kind(&self) -> AttrKind {
-        let first_token = self.syntax().first_token();
-        let first_token_kind = first_token.as_ref().map(SyntaxToken::kind);
-        let second_token_kind =
-            first_token.and_then(|token| token.next_token()).as_ref().map(SyntaxToken::kind);
-
-        match (first_token_kind, second_token_kind) {
-            (Some(T![#]), Some(T![!])) => AttrKind::Inner,
-            _ => AttrKind::Outer,
+        match self.excl_token() {
+            Some(_) => AttrKind::Inner,
+            None => AttrKind::Outer,
         }
     }
 
diff --git a/crates/syntax/src/ast/traits.rs b/crates/syntax/src/ast/traits.rs
@@ -76,8 +76,8 @@ pub trait HasDocComments: HasAttrs {
     fn doc_comments(&self) -> DocCommentIter {
         DocCommentIter { iter: self.syntax().children_with_tokens() }
     }
-    fn doc_comments_and_attrs(&self) -> AttrCommentIter {
-        AttrCommentIter { iter: self.syntax().children_with_tokens() }
+    fn doc_comments_and_attrs(&self) -> AttrDocCommentIter {
+        AttrDocCommentIter { iter: self.syntax().children_with_tokens() }
     }
 }
 
@@ -113,17 +113,23 @@ impl Iterator for DocCommentIter {
     }
 }
 
-pub struct AttrCommentIter {
+pub struct AttrDocCommentIter {
     iter: SyntaxElementChildren,
 }
 
-impl Iterator for AttrCommentIter {
-    type Item = Either<ast::Comment, ast::Attr>;
+impl AttrDocCommentIter {
+    pub fn from_syntax_node(syntax_node: &ast::SyntaxNode) -> AttrDocCommentIter {
+        AttrDocCommentIter { iter: syntax_node.children_with_tokens() }
+    }
+}
+
+impl Iterator for AttrDocCommentIter {
+    type Item = Either<ast::Attr, ast::Comment>;
     fn next(&mut self) -> Option<Self::Item> {
         self.iter.by_ref().find_map(|el| match el {
-            SyntaxElement::Node(node) => ast::Attr::cast(node).map(Either::Right),
+            SyntaxElement::Node(node) => ast::Attr::cast(node).map(Either::Left),
             SyntaxElement::Token(tok) => {
-                ast::Comment::cast(tok).filter(ast::Comment::is_doc).map(Either::Left)
+                ast::Comment::cast(tok).filter(ast::Comment::is_doc).map(Either::Right)
             }
         })
     }

Original file line number	Diff line number	Diff line change
`@@ -664,7 +664,7 @@ fn emit_def_diagnostic(db: &dyn HirDatabase, acc: &mut Vec<AnyDiagnostic>, diag:`
`664`	`664`	`let attr = node`
`665`	`665`	`.doc_comments_and_attrs()`
`666`	`666`	`.nth((*invoc_attr_index) as usize)`
`667`		`- .and_then(Either::right)`
	`667`	`+ .and_then(Either::left)`
`668`	`668`	`.unwrap_or_else(\|\| panic!("cannot find attribute #{}", invoc_attr_index));`
`669`	`669`	`(`
`670`	`670`	`ast_id.with_value(SyntaxNodePtr::from(AstPtr::new(&attr))),`
Original file line number	Diff line number	Diff line change
`@@ -117,7 +117,7 @@ impl ChildBySource for ItemScope {`
`117`	`117`	`\|(ast_id, calls)\| {`
`118`	`118`	`let adt = ast_id.to_node(db.upcast());`
`119`	`119`	`calls.for_each(\|(attr_id, calls)\| {`
`120`		`- if let Some(Either::Right(attr)) =`
	`120`	`+ if let Some(Either::Left(attr)) =`
`121`	`121`	`adt.doc_comments_and_attrs().nth(attr_id.ast_index as usize)`
`122`	`122`	`{`
`123`	`123`	`res[keys::DERIVE_MACRO_CALL].insert(attr, (attr_id, calls.into()));`
Original file line number	Diff line number	Diff line change
`@@ -191,7 +191,7 @@ fn make_hygiene_info(`
`191`	`191`	`.to_node(db)`
`192`	`192`	`.doc_comments_and_attrs()`
`193`	`193`	`.nth(invoc_attr_index as usize)`
`194`		`- .and_then(Either::right)?`
	`194`	`+ .and_then(Either::left)?`
`195`	`195`	`.token_tree()?;`
`196`	`196`	`Some(InFile::new(ast_id.file_id, tt))`
`197`	`197`	`}`
Original file line number	Diff line number	Diff line change
`@@ -160,14 +160,9 @@ impl ast::Attr {`
`160`	`160`	`}`
`161`	`161`
`162`	`162`	`pub fn kind(&self) -> AttrKind {`
`163`		`- let first_token = self.syntax().first_token();`
`164`		`- let first_token_kind = first_token.as_ref().map(SyntaxToken::kind);`
`165`		`- let second_token_kind =`
`166`		`- first_token.and_then(\|token\| token.next_token()).as_ref().map(SyntaxToken::kind);`
`167`		`-`
`168`		`- match (first_token_kind, second_token_kind) {`
`169`		`- (Some(T![#]), Some(T![!])) => AttrKind::Inner,`
`170`		`- _ => AttrKind::Outer,`
	`163`	`+ match self.excl_token() {`
	`164`	`+ Some(_) => AttrKind::Inner,`
	`165`	`+ None => AttrKind::Outer,`
`171`	`166`	`}`
`172`	`167`	`}`
`173`	`168`
Original file line number	Diff line number	Diff line change
`@@ -76,8 +76,8 @@ pub trait HasDocComments: HasAttrs {`
`76`	`76`	`fn doc_comments(&self) -> DocCommentIter {`
`77`	`77`	`DocCommentIter { iter: self.syntax().children_with_tokens() }`
`78`	`78`	`}`
`79`		`- fn doc_comments_and_attrs(&self) -> AttrCommentIter {`
`80`		`- AttrCommentIter { iter: self.syntax().children_with_tokens() }`
	`79`	`+ fn doc_comments_and_attrs(&self) -> AttrDocCommentIter {`
	`80`	`+ AttrDocCommentIter { iter: self.syntax().children_with_tokens() }`
`81`	`81`	`}`
`82`	`82`	`}`
`83`	`83`
`@@ -113,17 +113,23 @@ impl Iterator for DocCommentIter {`
`113`	`113`	`}`
`114`	`114`	`}`
`115`	`115`
`116`		`-pub struct AttrCommentIter {`
	`116`	`+pub struct AttrDocCommentIter {`
`117`	`117`	`iter: SyntaxElementChildren,`
`118`	`118`	`}`
`119`	`119`
`120`		`-impl Iterator for AttrCommentIter {`
`121`		`- type Item = Either<ast::Comment, ast::Attr>;`
	`120`	`+impl AttrDocCommentIter {`
	`121`	`+ pub fn from_syntax_node(syntax_node: &ast::SyntaxNode) -> AttrDocCommentIter {`
	`122`	`+ AttrDocCommentIter { iter: syntax_node.children_with_tokens() }`
	`123`	`+ }`
	`124`	`+}`
	`125`	`+`
	`126`	`+impl Iterator for AttrDocCommentIter {`
	`127`	`+ type Item = Either<ast::Attr, ast::Comment>;`
`122`	`128`	`fn next(&mut self) -> Option<Self::Item> {`
`123`	`129`	`self.iter.by_ref().find_map(\|el\| match el {`
`124`		`- SyntaxElement::Node(node) => ast::Attr::cast(node).map(Either::Right),`
	`130`	`+ SyntaxElement::Node(node) => ast::Attr::cast(node).map(Either::Left),`
`125`	`131`	`SyntaxElement::Token(tok) => {`
`126`		`- ast::Comment::cast(tok).filter(ast::Comment::is_doc).map(Either::Left)`
	`132`	`+ ast::Comment::cast(tok).filter(ast::Comment::is_doc).map(Either::Right)`
`127`	`133`	`}`
`128`	`134`	`})`
`129`	`135`	`}`