2 files changed, 79 insertions, 26 deletions

diff --git a/crates/libsyntax2/src/ast/generated.rs b/crates/libsyntax2/src/ast/generated.rs
index 1ec05c950..c575e15df 100644
--- a/crates/libsyntax2/src/ast/generated.rs
+++ b/crates/libsyntax2/src/ast/generated.rs
@@ -23,6 +23,31 @@ impl<R: TreeRoot> AstNode<R> for ArrayType<R> {
 
 impl<R: TreeRoot> ArrayType<R> {}
 
+// Attr
+#[derive(Debug, Clone, Copy)]
+pub struct Attr<R: TreeRoot = Arc<SyntaxRoot>> {
+    syntax: SyntaxNode<R>,
+}
+
+impl<R: TreeRoot> AstNode<R> for Attr<R> {
+    fn cast(syntax: SyntaxNode<R>) -> Option<Self> {
+        match syntax.kind() {
+            ATTR => Some(Attr { syntax }),
+            _ => None,
+        }
+    }
+    fn syntax(&self) -> &SyntaxNode<R> { &self.syntax }
+}
+
+impl<R: TreeRoot> Attr<R> {
+    pub fn value(&self) -> Option<TokenTree<R>> {
+        self.syntax()
+            .children()
+            .filter_map(TokenTree::cast)
+            .next()
+    }
+}
+
 // ConstDef
 #[derive(Debug, Clone, Copy)]
 pub struct ConstDef<R: TreeRoot = Arc<SyntaxRoot>> {
@@ -40,6 +65,7 @@ impl<R: TreeRoot> AstNode<R> for ConstDef<R> {
 }
 
 impl<R: TreeRoot> ast::NameOwner<R> for ConstDef<R> {}
+impl<R: TreeRoot> ast::AttrsOwner<R> for ConstDef<R> {}
 impl<R: TreeRoot> ConstDef<R> {}
 
 // DynTraitType
@@ -77,6 +103,7 @@ impl<R: TreeRoot> AstNode<R> for EnumDef<R> {
 }
 
 impl<R: TreeRoot> ast::NameOwner<R> for EnumDef<R> {}
+impl<R: TreeRoot> ast::AttrsOwner<R> for EnumDef<R> {}
 impl<R: TreeRoot> EnumDef<R> {}
 
 // File
@@ -120,6 +147,7 @@ impl<R: TreeRoot> AstNode<R> for FnDef<R> {
 }
 
 impl<R: TreeRoot> ast::NameOwner<R> for FnDef<R> {}
+impl<R: TreeRoot> ast::AttrsOwner<R> for FnDef<R> {}
 impl<R: TreeRoot> FnDef<R> {}
 
 // FnPointerType
@@ -211,6 +239,7 @@ impl<R: TreeRoot> AstNode<R> for Module<R> {
 }
 
 impl<R: TreeRoot> ast::NameOwner<R> for Module<R> {}
+impl<R: TreeRoot> ast::AttrsOwner<R> for Module<R> {}
 impl<R: TreeRoot> Module<R> {}
 
 // Name
@@ -266,6 +295,7 @@ impl<R: TreeRoot> AstNode<R> for NamedField<R> {
 }
 
 impl<R: TreeRoot> ast::NameOwner<R> for NamedField<R> {}
+impl<R: TreeRoot> ast::AttrsOwner<R> for NamedField<R> {}
 impl<R: TreeRoot> NamedField<R> {}
 
 // NeverType
@@ -436,6 +466,7 @@ impl<R: TreeRoot> AstNode<R> for StaticDef<R> {
 }
 
 impl<R: TreeRoot> ast::NameOwner<R> for StaticDef<R> {}
+impl<R: TreeRoot> ast::AttrsOwner<R> for StaticDef<R> {}
 impl<R: TreeRoot> StaticDef<R> {}
 
 // StructDef
@@ -455,6 +486,7 @@ impl<R: TreeRoot> AstNode<R> for StructDef<R> {
 }
 
 impl<R: TreeRoot> ast::NameOwner<R> for StructDef<R> {}
+impl<R: TreeRoot> ast::AttrsOwner<R> for StructDef<R> {}
 impl<R: TreeRoot> StructDef<R> {
     pub fn fields<'a>(&'a self) -> impl Iterator<Item = NamedField<R>> + 'a {
         self.syntax()
@@ -463,6 +495,24 @@ impl<R: TreeRoot> StructDef<R> {
     }
 }
 
+// TokenTree
+#[derive(Debug, Clone, Copy)]
+pub struct TokenTree<R: TreeRoot = Arc<SyntaxRoot>> {
+    syntax: SyntaxNode<R>,
+}
+
+impl<R: TreeRoot> AstNode<R> for TokenTree<R> {
+    fn cast(syntax: SyntaxNode<R>) -> Option<Self> {
+        match syntax.kind() {
+            TOKEN_TREE => Some(TokenTree { syntax }),
+            _ => None,
+        }
+    }
+    fn syntax(&self) -> &SyntaxNode<R> { &self.syntax }
+}
+
+impl<R: TreeRoot> TokenTree<R> {}
+
 // TraitDef
 #[derive(Debug, Clone, Copy)]
 pub struct TraitDef<R: TreeRoot = Arc<SyntaxRoot>> {
@@ -480,6 +530,7 @@ impl<R: TreeRoot> AstNode<R> for TraitDef<R> {
 }
 
 impl<R: TreeRoot> ast::NameOwner<R> for TraitDef<R> {}
+impl<R: TreeRoot> ast::AttrsOwner<R> for TraitDef<R> {}
 impl<R: TreeRoot> TraitDef<R> {}
 
 // TupleType
@@ -517,6 +568,7 @@ impl<R: TreeRoot> AstNode<R> for TypeDef<R> {
 }
 
 impl<R: TreeRoot> ast::NameOwner<R> for TypeDef<R> {}
+impl<R: TreeRoot> ast::AttrsOwner<R> for TypeDef<R> {}
 impl<R: TreeRoot> TypeDef<R> {}
 
 // TypeRef
diff --git a/crates/libsyntax2/src/ast/mod.rs b/crates/libsyntax2/src/ast/mod.rs
index 679e292a2..fe8f91d15 100644
--- a/crates/libsyntax2/src/ast/mod.rs
+++ b/crates/libsyntax2/src/ast/mod.rs
@@ -2,10 +2,11 @@ mod generated;
 
 use std::sync::Arc;
 
+use itertools::Itertools;
 use smol_str::SmolStr;
 
 use {
-    SyntaxNode, SyntaxRoot, TreeRoot, SyntaxError,
+    SyntaxNode, SyntaxNodeRef, SyntaxRoot, TreeRoot, SyntaxError,
     SyntaxKind::*,
 };
 pub use self::generated::*;
@@ -14,6 +15,9 @@ pub trait AstNode<R: TreeRoot> {
     fn cast(syntax: SyntaxNode<R>) -> Option<Self>
         where Self: Sized;
     fn syntax(&self) -> &SyntaxNode<R>;
+    fn syntax_ref<'a>(&'a self) -> SyntaxNodeRef<'a> where R: 'a {
+        self.syntax().as_ref()
+    }
 }
 
 pub trait NameOwner<R: TreeRoot>: AstNode<R> {
@@ -25,6 +29,14 @@ pub trait NameOwner<R: TreeRoot>: AstNode<R> {
     }
 }
 
+pub trait AttrsOwner<R: TreeRoot>: AstNode<R> {
+    fn attrs<'a>(&'a self) -> Box<Iterator<Item=Attr<R>> + 'a> where R: 'a {
+        let it = self.syntax().children()
+            .filter_map(Attr::cast);
+        Box::new(it)
+    }
+}
+
 impl File<Arc<SyntaxRoot>> {
     pub fn parse(text: &str) -> Self {
         File::cast(::parse(text)).unwrap()
@@ -39,31 +51,20 @@ impl<R: TreeRoot> File<R> {
 
 impl<R: TreeRoot> FnDef<R> {
     pub fn has_atom_attr(&self, atom: &str) -> bool {
-        self.syntax()
-            .children()
-            .filter(|node| node.kind() == ATTR)
-            .any(|attr| {
-                let mut metas = attr.children().filter(|node| node.kind() == META_ITEM);
-                let meta = match metas.next() {
-                    None => return false,
-                    Some(meta) => {
-                        if metas.next().is_some() {
-                            return false;
-                        }
-                        meta
-                    }
-                };
-                let mut children = meta.children();
-                match children.next() {
-                    None => false,
-                    Some(child) => {
-                        if children.next().is_some() {
-                            return false;
-                        }
-                        child.kind() == IDENT && child.text() == atom
-                    }
-                }
-            })
+        self.attrs()
+            .filter_map(|x| x.value())
+            .filter_map(|x| as_atom(x))
+            .any(|x| x == atom)
+    }
+}
+
+fn as_atom<R: TreeRoot>(tt: TokenTree<R>) -> Option<SmolStr> {
+    let syntax = tt.syntax_ref();
+    let (_bra, attr, _ket) = syntax.children().collect_tuple()?;
+    if attr.kind() == IDENT {
+        Some(attr.leaf_text().unwrap())
+    } else {
+        None
     }
 }