5 files changed, 103 insertions, 40 deletions

diff --git a/src/ast/generated.rs b/src/ast/generated.rs
new file mode 100644
index 000000000..612b04f86
--- /dev/null
+++ b/src/ast/generated.rs
@@ -0,0 +1,54 @@
+use std::sync::Arc;
+use {
+    SyntaxNode, SyntaxRoot, TreeRoot, AstNode,
+    SyntaxKind::*,
+};
+
+
+#[derive(Debug)]
+pub struct File<R: TreeRoot = Arc<SyntaxRoot>> {
+    syntax: SyntaxNode<R>,
+}
+
+impl<R: TreeRoot> AstNode<R> for File<R> {
+    fn cast(syntax: SyntaxNode<R>) -> Option<Self> {
+        match syntax.kind() {
+            FILE => Some(File { syntax }),
+            _ => None,
+        }
+    }
+    fn syntax(&self) -> &SyntaxNode<R> { &self.syntax }
+}
+
+
+#[derive(Debug)]
+pub struct FnItem<R: TreeRoot = Arc<SyntaxRoot>> {
+    syntax: SyntaxNode<R>,
+}
+
+impl<R: TreeRoot> AstNode<R> for FnItem<R> {
+    fn cast(syntax: SyntaxNode<R>) -> Option<Self> {
+        match syntax.kind() {
+            FN_ITEM => Some(FnItem { syntax }),
+            _ => None,
+        }
+    }
+    fn syntax(&self) -> &SyntaxNode<R> { &self.syntax }
+}
+
+
+#[derive(Debug)]
+pub struct Name<R: TreeRoot = Arc<SyntaxRoot>> {
+    syntax: SyntaxNode<R>,
+}
+
+impl<R: TreeRoot> AstNode<R> for Name<R> {
+    fn cast(syntax: SyntaxNode<R>) -> Option<Self> {
+        match syntax.kind() {
+            NAME => Some(Name { syntax }),
+            _ => None,
+        }
+    }
+    fn syntax(&self) -> &SyntaxNode<R> { &self.syntax }
+}
+
diff --git a/src/ast/generated.rs.tera b/src/ast/generated.rs.tera
new file mode 100644
index 000000000..f2559383a
--- /dev/null
+++ b/src/ast/generated.rs.tera
@@ -0,0 +1,22 @@
+use std::sync::Arc;
+use {
+    SyntaxNode, SyntaxRoot, TreeRoot, AstNode,
+    SyntaxKind::*,
+};
+{% for node in ast %}
+{% set Name = node.kind | camel  %}
+#[derive(Debug)]
+pub struct {{ Name }}<R: TreeRoot = Arc<SyntaxRoot>> {
+    syntax: SyntaxNode<R>,
+}
+
+impl<R: TreeRoot> AstNode<R> for {{ Name }}<R> {
+    fn cast(syntax: SyntaxNode<R>) -> Option<Self> {
+        match syntax.kind() {
+            {{ node.kind }} => Some({{ Name }} { syntax }),
+            _ => None,
+        }
+    }
+    fn syntax(&self) -> &SyntaxNode<R> { &self.syntax }
+}
+{% endfor %}
diff --git a/src/ast.rs b/src/ast/mod.rs
index a595b9324..dc7e006c9 100644
--- a/src/ast.rs
+++ b/src/ast/mod.rs
@@ -1,57 +1,41 @@
+mod generated;
+
 use std::sync::Arc;
 use {
     SyntaxNode, SyntaxRoot, TreeRoot,
     SyntaxKind::*,
 };
+pub use self::generated::*;
 
-#[derive(Debug)]
-pub struct File<R: TreeRoot = Arc<SyntaxRoot>> {
-    syntax: SyntaxNode<R>,
-}
-
-#[derive(Debug)]
-pub struct Function<R: TreeRoot = Arc<SyntaxRoot>> {
-    syntax: SyntaxNode<R>,
+pub trait AstNode<R: TreeRoot>: Sized {
+    fn cast(syntax: SyntaxNode<R>) -> Option<Self>;
+    fn syntax(&self) -> &SyntaxNode<R>;
 }
 
-#[derive(Debug)]
-pub struct Name<R: TreeRoot = Arc<SyntaxRoot>> {
-    syntax: SyntaxNode<R>,
-}
-
-
 impl File<Arc<SyntaxRoot>> {
     pub fn parse(text: &str) -> Self {
-        File {
-            syntax: ::parse(text),
-        }
+        File::cast(::parse(text)).unwrap()
     }
 }
 
 impl<R: TreeRoot> File<R> {
-    pub fn functions<'a>(&'a self) -> impl Iterator<Item = Function<R>> + 'a {
-        self.syntax
+    pub fn functions<'a>(&'a self) -> impl Iterator<Item = FnItem<R>> + 'a {
+        self.syntax()
             .children()
-            .filter(|node| node.kind() == FN_ITEM)
-            .map(|node| Function { syntax: node })
+            .filter_map(FnItem::cast)
     }
 }
 
-impl<R: TreeRoot> Function<R> {
-    pub fn syntax(&self) -> SyntaxNode<R> {
-        self.syntax.clone()
-    }
-
+impl<R: TreeRoot> FnItem<R> {
     pub fn name(&self) -> Option<Name<R>> {
-        self.syntax
+        self.syntax()
             .children()
-            .filter(|node| node.kind() == NAME)
-            .map(|node| Name { syntax: node })
+            .filter_map(Name::cast)
             .next()
     }
 
     pub fn has_atom_attr(&self, atom: &str) -> bool {
-        self.syntax
+        self.syntax()
             .children()
             .filter(|node| node.kind() == ATTR)
             .any(|attr| {
@@ -81,14 +65,6 @@ impl<R: TreeRoot> Function<R> {
 
 impl<R: TreeRoot> Name<R> {
     pub fn text(&self) -> String {
-        self.syntax.text()
-    }
-}
-
-
-
-impl<R: TreeRoot> File<R> {
-    pub fn syntax(&self) -> SyntaxNode<R> {
-        self.syntax.clone()
+        self.syntax().text()
     }
 }
diff --git a/src/grammar.ron b/src/grammar.ron
index 0443dd798..b6a870d84 100644
--- a/src/grammar.ron
+++ b/src/grammar.ron
@@ -212,5 +212,16 @@ Grammar(
         "PARAM",
         "SELF_PARAM",
         "ARG_LIST",
+    ],
+    ast: [
+        (
+            kind: "FILE"
+        ),
+        (
+            kind: "FN_ITEM"
+        ),
+        (
+            kind: "NAME"
+        ),
     ]
 )
diff --git a/src/lib.rs b/src/lib.rs
index 91b0c9c55..d1e690bb2 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -41,7 +41,7 @@ mod yellow;
 pub mod utils;
 
 pub use {
-    ast::File,
+    ast::{AstNode, File},
     lexer::{tokenize, Token},
     syntax_kinds::SyntaxKind,
     text_unit::{TextRange, TextUnit},