2 files changed, 111 insertions, 65 deletions

diff --git a/crates/ra_syntax/src/ast.rs b/crates/ra_syntax/src/ast.rs
index 26fafb469..1ac0201b8 100644
--- a/crates/ra_syntax/src/ast.rs
+++ b/crates/ra_syntax/src/ast.rs
@@ -11,7 +11,10 @@ pub mod make;
 use std::marker::PhantomData;
 
 use crate::{
-    syntax_node::{SyntaxNode, SyntaxNodeChildren, SyntaxToken},
+    syntax_node::{
+        NodeOrToken, SyntaxElement, SyntaxElementChildren, SyntaxNode, SyntaxNodeChildren,
+        SyntaxToken,
+    },
     SmolStr, SyntaxKind,
 };
 
@@ -30,16 +33,24 @@ pub use self::{
 /// conversion itself has zero runtime cost: ast and syntax nodes have exactly
 /// the same representation: a pointer to the tree root and a pointer to the
 /// node itself.
-pub trait AstNode: std::fmt::Display {
+pub trait AstNode: AstElement {
     fn can_cast(kind: SyntaxKind) -> bool
     where
         Self: Sized;
 
-    fn cast(syntax: SyntaxNode) -> Option<Self>
+    fn cast_or_return(syntax: SyntaxNode) -> Result<Self, SyntaxNode>
     where
         Self: Sized;
 
+    fn cast(syntax: SyntaxNode) -> Option<Self>
+    where
+        Self: Sized,
+    {
+        <Self as AstNode>::cast_or_return(syntax).ok()
+    }
+
     fn syntax(&self) -> &SyntaxNode;
+    fn into_syntax(self) -> SyntaxNode;
 }
 
 #[test]
@@ -48,16 +59,51 @@ fn assert_ast_is_object_safe() {
 }
 
 /// Like `AstNode`, but wraps tokens rather than interior nodes.
-pub trait AstToken {
-    fn cast(token: SyntaxToken) -> Option<Self>
+pub trait AstToken: AstElement {
+    fn can_cast(token: SyntaxKind) -> bool
     where
         Self: Sized;
+
+    fn cast_or_return(syntax: SyntaxToken) -> Result<Self, SyntaxToken>
+    where
+        Self: Sized;
+
+    fn cast(syntax: SyntaxToken) -> Option<Self>
+    where
+        Self: Sized,
+    {
+        <Self as AstToken>::cast_or_return(syntax).ok()
+    }
+
     fn syntax(&self) -> &SyntaxToken;
+    fn into_syntax(self) -> SyntaxToken;
+
     fn text(&self) -> &SmolStr {
         self.syntax().text()
     }
 }
 
+/// Like `AstNode`, but wraps either nodes or tokens rather than interior nodes.
+pub trait AstElement: std::fmt::Display {
+    fn can_cast_element(kind: SyntaxKind) -> bool
+    where
+        Self: Sized;
+
+    fn cast_or_return_element(syntax: SyntaxElement) -> Result<Self, SyntaxElement>
+    where
+        Self: Sized;
+
+    fn cast_element(syntax: SyntaxElement) -> Option<Self>
+    where
+        Self: Sized,
+    {
+        <Self as AstElement>::cast_or_return_element(syntax).ok()
+    }
+
+    fn syntax_element(&self) -> NodeOrToken<&SyntaxNode, &SyntaxToken>;
+    fn into_syntax_element(self) -> SyntaxElement;
+}
+
 /// An iterator over `SyntaxNode` children of a particular AST type.
 #[derive(Debug, Clone)]
 pub struct AstChildren<N> {
@@ -86,6 +132,64 @@ fn children<P: AstNode + ?Sized, C: AstNode>(parent: &P) -> AstChildren<C> {
     AstChildren::new(parent.syntax())
 }
 
+/// An iterator over `SyntaxToken` children of a particular AST type.
+#[derive(Debug, Clone)]
+pub struct AstChildTokens<N> {
+    inner: SyntaxElementChildren,
+    ph: PhantomData<N>,
+}
+
+impl<N> AstChildTokens<N> {
+    fn new(parent: &SyntaxNode) -> Self {
+        AstChildTokens { inner: parent.children_with_tokens(), ph: PhantomData }
+    }
+}
+
+impl<N: AstToken> Iterator for AstChildTokens<N> {
+    type Item = N;
+    fn next(&mut self) -> Option<N> {
+        self.inner.by_ref().filter_map(|x| x.into_token()).find_map(N::cast)
+    }
+}
+
+fn child_token_opt<P: AstNode + ?Sized, C: AstToken>(parent: &P) -> Option<C> {
+    child_tokens(parent).next()
+}
+
+fn child_tokens<P: AstNode + ?Sized, C: AstToken>(parent: &P) -> AstChildTokens<C> {
+    AstChildTokens::new(parent.syntax())
+}
+
+/// An iterator over `SyntaxNode` children of a particular AST type.
+#[derive(Debug, Clone)]
+pub struct AstChildElements<N> {
+    inner: SyntaxElementChildren,
+    ph: PhantomData<N>,
+}
+
+impl<N> AstChildElements<N> {
+    fn new(parent: &SyntaxNode) -> Self {
+        AstChildElements { inner: parent.children_with_tokens(), ph: PhantomData }
+    }
+}
+
+impl<N: AstElement> Iterator for AstChildElements<N> {
+    type Item = N;
+    fn next(&mut self) -> Option<N> {
+        self.inner.by_ref().find_map(N::cast_element)
+    }
+}
+
+#[allow(dead_code)]
+fn child_element_opt<P: AstNode + ?Sized, C: AstElement>(parent: &P) -> Option<C> {
+    child_elements(parent).next()
+}
+
+#[allow(dead_code)]
+fn child_elements<P: AstNode + ?Sized, C: AstElement>(parent: &P) -> AstChildElements<C> {
+    AstChildElements::new(parent.syntax())
+}
+
 #[test]
 fn test_doc_comment_none() {
     let file = SourceFile::parse(
diff --git a/crates/ra_syntax/src/ast/tokens.rs b/crates/ra_syntax/src/ast/tokens.rs
index 1a51b8d3b..e8320b57e 100644
--- a/crates/ra_syntax/src/ast/tokens.rs
+++ b/crates/ra_syntax/src/ast/tokens.rs
@@ -1,26 +1,10 @@
 //! There are many AstNodes, but only a few tokens, so we hand-write them here.
 
 use crate::{
-    ast::AstToken,
-    SyntaxKind::{COMMENT, RAW_STRING, STRING, WHITESPACE},
-    SyntaxToken, TextRange, TextUnit,
+    ast::{AstToken, Comment, RawString, String, Whitespace},
+    TextRange, TextUnit,
 };
 
-#[derive(Debug, Clone, PartialEq, Eq, Hash)]
-pub struct Comment(SyntaxToken);
-
-impl AstToken for Comment {
-    fn cast(token: SyntaxToken) -> Option<Self> {
-        match token.kind() {
-            COMMENT => Some(Comment(token)),
-            _ => None,
-        }
-    }
-    fn syntax(&self) -> &SyntaxToken {
-        &self.0
-    }
-}
-
 impl Comment {
     pub fn kind(&self) -> CommentKind {
         kind_by_prefix(self.text())
@@ -89,20 +73,6 @@ fn prefix_by_kind(kind: CommentKind) -> &'static str {
     unreachable!()
 }
 
-pub struct Whitespace(SyntaxToken);
-
-impl AstToken for Whitespace {
-    fn cast(token: SyntaxToken) -> Option<Self> {
-        match token.kind() {
-            WHITESPACE => Some(Whitespace(token)),
-            _ => None,
-        }
-    }
-    fn syntax(&self) -> &SyntaxToken {
-        &self.0
-    }
-}
-
 impl Whitespace {
     pub fn spans_multiple_lines(&self) -> bool {
         let text = self.text();
@@ -168,20 +138,6 @@ pub trait HasStringValue: HasQuotes {
     fn value(&self) -> Option<std::string::String>;
 }
 
-pub struct String(SyntaxToken);
-
-impl AstToken for String {
-    fn cast(token: SyntaxToken) -> Option<Self> {
-        match token.kind() {
-            STRING => Some(String(token)),
-            _ => None,
-        }
-    }
-    fn syntax(&self) -> &SyntaxToken {
-        &self.0
-    }
-}
-
 impl HasStringValue for String {
     fn value(&self) -> Option<std::string::String> {
         let text = self.text().as_str();
@@ -201,20 +157,6 @@ impl HasStringValue for String {
     }
 }
 
-pub struct RawString(SyntaxToken);
-
-impl AstToken for RawString {
-    fn cast(token: SyntaxToken) -> Option<Self> {
-        match token.kind() {
-            RAW_STRING => Some(RawString(token)),
-            _ => None,
-        }
-    }
-    fn syntax(&self) -> &SyntaxToken {
-        &self.0
-    }
-}
-
 impl HasStringValue for RawString {
     fn value(&self) -> Option<std::string::String> {
         let text = self.text().as_str();