Handle `cfg_attr` gating multiple attributes

1 files changed, 19 insertions, 29 deletions

diff --git a/crates/hir_def/src/attr.rs b/crates/hir_def/src/attr.rs
index 97cdbbb9e..7b41b148c 100644
--- a/crates/hir_def/src/attr.rs
+++ b/crates/hir_def/src/attr.rs
@@ -9,6 +9,7 @@ use hir_expand::{hygiene::Hygiene, name::AsName, AstId, InFile};
 use itertools::Itertools;
 use la_arena::ArenaMap;
 use mbe::ast_to_token_tree;
+use smallvec::{smallvec, SmallVec};
 use syntax::{
     ast::{self, AstNode, AttrsOwner},
     match_ast, AstToken, SmolStr, SyntaxNode,
@@ -134,53 +135,42 @@ impl RawAttrs {
         let crate_graph = db.crate_graph();
         let new_attrs = self
             .iter()
-            .filter_map(|attr| {
+            .flat_map(|attr| -> SmallVec<[_; 1]> {
                 let attr = attr.clone();
                 let is_cfg_attr =
                     attr.path.as_ident().map_or(false, |name| *name == hir_expand::name![cfg_attr]);
                 if !is_cfg_attr {
-                    return Some(attr);
+                    return smallvec![attr];
                 }
 
                 let subtree = match &attr.input {
                     Some(AttrInput::TokenTree(it)) => it,
-                    _ => return Some(attr),
+                    _ => return smallvec![attr],
                 };
 
-                // Input subtree is: `(cfg, attr)`
-                // Split it up into a `cfg` and an `attr` subtree.
+                // Input subtree is: `(cfg, $(attr),+)`
+                // Split it up into a `cfg` subtree and the `attr` subtrees.
                 // FIXME: There should be a common API for this.
-                let mut saw_comma = false;
-                let (mut cfg, attr): (Vec<_>, Vec<_>) =
-                    subtree.clone().token_trees.into_iter().partition(|tree| {
-                        if saw_comma {
-                            return false;
-                        }
-
-                        match tree {
-                            tt::TokenTree::Leaf(tt::Leaf::Punct(p)) if p.char == ',' => {
-                                saw_comma = true;
-                            }
-                            _ => {}
-                        }
-
-                        true
-                    });
-                cfg.pop(); // `,` ends up in here
-
-                let attr = Subtree { delimiter: None, token_trees: attr };
-                let cfg = Subtree { delimiter: subtree.delimiter, token_trees: cfg };
+                let mut parts = subtree.token_trees.split(
+                    |tt| matches!(tt, tt::TokenTree::Leaf(tt::Leaf::Punct(p)) if p.char == ','),
+                );
+                let cfg = parts.next().unwrap();
+                let cfg = Subtree { delimiter: subtree.delimiter, token_trees: cfg.to_vec() };
                 let cfg = CfgExpr::parse(&cfg);
+                let attrs = parts.filter(|a| !a.is_empty()).filter_map(|attr| {
+                    let tree = Subtree { delimiter: None, token_trees: attr.to_vec() };
+                    let attr = ast::Attr::parse(&format!("#[{}]", tree)).ok()?;
+                    let hygiene = Hygiene::new_unhygienic(); // FIXME
+                    Attr::from_src(attr, &hygiene)
+                });
 
                 let cfg_options = &crate_graph[krate].cfg_options;
                 if cfg_options.check(&cfg) == Some(false) {
-                    None
+                    smallvec![]
                 } else {
                     cov_mark::hit!(cfg_attr_active);
 
-                    let attr = ast::Attr::parse(&format!("#[{}]", attr)).ok()?;
-                    let hygiene = Hygiene::new_unhygienic(); // FIXME
-                    Attr::from_src(attr, &hygiene)
+                    attrs.collect()
                 }
             })
             .collect();