1 files changed, 53 insertions, 30 deletions

diff --git a/crates/ra_mbe/src/syntax_bridge.rs b/crates/ra_mbe/src/syntax_bridge.rs
index 592fcf527..9653f7fef 100644
--- a/crates/ra_mbe/src/syntax_bridge.rs
+++ b/crates/ra_mbe/src/syntax_bridge.rs
@@ -14,12 +14,18 @@ use crate::subtree_source::SubtreeTokenSource;
 use crate::ExpandError;
 
 /// Maps `tt::TokenId` to the relative range of the original token.
-#[derive(Default)]
+#[derive(Debug, PartialEq, Eq, Default)]
 pub struct TokenMap {
     /// Maps `tt::TokenId` to the *relative* source range.
     tokens: Vec<TextRange>,
 }
 
+/// Maps relative range of the expanded syntax node to `tt::TokenId`
+#[derive(Debug, PartialEq, Eq, Default)]
+pub struct RevTokenMap {
+    pub ranges: Vec<(TextRange, tt::TokenId)>,
+}
+
 /// Convert the syntax tree (what user has written) to a `TokenTree` (what macro
 /// will consume).
 pub fn ast_to_token_tree(ast: &ast::TokenTree) -> Option<(tt::Subtree, TokenMap)> {
@@ -52,7 +58,7 @@ pub fn syntax_node_to_token_tree(node: &SyntaxNode) -> Option<(tt::Subtree, Toke
 fn fragment_to_syntax_node(
     tt: &tt::Subtree,
     fragment_kind: FragmentKind,
-) -> Result<Parse<SyntaxNode>, ExpandError> {
+) -> Result<(Parse<SyntaxNode>, RevTokenMap), ExpandError> {
     let tmp;
     let tokens = match tt {
         tt::Subtree { delimiter: tt::Delimiter::None, token_trees } => token_trees.as_slice(),
@@ -69,38 +75,33 @@ fn fragment_to_syntax_node(
         return Err(ExpandError::ConversionError);
     }
     //FIXME: would be cool to report errors
-    let parse = tree_sink.inner.finish();
-    Ok(parse)
+    let (parse, range_map) = tree_sink.finish();
+    Ok((parse, range_map))
 }
 
-/// Parses the token tree (result of macro expansion) to an expression
-pub fn token_tree_to_expr(tt: &tt::Subtree) -> Result<Parse<ast::Expr>, ExpandError> {
-    let parse = fragment_to_syntax_node(tt, Expr)?;
-    parse.cast().ok_or_else(|| crate::ExpandError::ConversionError)
-}
-
-/// Parses the token tree (result of macro expansion) to a Pattern
-pub fn token_tree_to_pat(tt: &tt::Subtree) -> Result<Parse<ast::Pat>, ExpandError> {
-    let parse = fragment_to_syntax_node(tt, Pattern)?;
-    parse.cast().ok_or_else(|| crate::ExpandError::ConversionError)
-}
-
-/// Parses the token tree (result of macro expansion) to a Type
-pub fn token_tree_to_ty(tt: &tt::Subtree) -> Result<Parse<ast::TypeRef>, ExpandError> {
-    let parse = fragment_to_syntax_node(tt, Type)?;
-    parse.cast().ok_or_else(|| crate::ExpandError::ConversionError)
-}
-
-/// Parses the token tree (result of macro expansion) as a sequence of stmts
-pub fn token_tree_to_macro_stmts(tt: &tt::Subtree) -> Result<Parse<ast::MacroStmts>, ExpandError> {
-    let parse = fragment_to_syntax_node(tt, Statements)?;
-    parse.cast().ok_or_else(|| crate::ExpandError::ConversionError)
+macro_rules! impl_token_tree_conversions {
+    ($($(#[$attr:meta])* $name:ident => ($kind:ident, $t:ty) ),*) => {
+        $(
+            $(#[$attr])*
+            pub fn $name(tt: &tt::Subtree) -> Result<(Parse<$t>, RevTokenMap), ExpandError> {
+                let (parse, map) = fragment_to_syntax_node(tt, $kind)?;
+                parse.cast().ok_or_else(|| crate::ExpandError::ConversionError).map(|p| (p, map))
+            }
+        )*
+    }
 }
 
-/// Parses the token tree (result of macro expansion) as a sequence of items
-pub fn token_tree_to_items(tt: &tt::Subtree) -> Result<Parse<ast::MacroItems>, ExpandError> {
-    let parse = fragment_to_syntax_node(tt, Items)?;
-    parse.cast().ok_or_else(|| crate::ExpandError::ConversionError)
+impl_token_tree_conversions! {
+    /// Parses the token tree (result of macro expansion) to an expression
+    token_tree_to_expr => (Expr, ast::Expr),
+    /// Parses the token tree (result of macro expansion) to a Pattern
+    token_tree_to_pat => (Pattern, ast::Pat),
+    /// Parses the token tree (result of macro expansion) to a Type
+    token_tree_to_ty => (Type, ast::TypeRef),
+    /// Parses the token tree (result of macro expansion) as a sequence of stmts
+    token_tree_to_macro_stmts => (Statements, ast::MacroStmts),
+    /// Parses the token tree (result of macro expansion) as a sequence of items
+    token_tree_to_items => (Items, ast::MacroItems)
 }
 
 impl TokenMap {
@@ -116,6 +117,12 @@ impl TokenMap {
     }
 }
 
+impl RevTokenMap {
+    fn add(&mut self, relative_range: TextRange, token_id: tt::TokenId) {
+        self.ranges.push((relative_range, token_id.clone()))
+    }
+}
+
 /// Returns the textual content of a doc comment block as a quoted string
 /// That is, strips leading `///` (or `/**`, etc)
 /// and strips the ending `*/`
@@ -262,6 +269,7 @@ struct TtTreeSink<'a> {
     cursor: Cursor<'a>,
     text_pos: TextUnit,
     inner: SyntaxTreeBuilder,
+    range_map: RevTokenMap,
 
     // Number of roots
     // Use for detect ill-form tree which is not single root
@@ -276,8 +284,13 @@ impl<'a> TtTreeSink<'a> {
             text_pos: 0.into(),
             inner: SyntaxTreeBuilder::default(),
             roots: smallvec::SmallVec::new(),
+            range_map: RevTokenMap::default(),
         }
     }
+
+    fn finish(self) -> (Parse<SyntaxNode>, RevTokenMap) {
+        (self.inner.finish(), self.range_map)
+    }
 }
 
 fn delim_to_str(d: tt::Delimiter, closing: bool) -> SmolStr {
@@ -307,6 +320,15 @@ impl<'a> TreeSink for TtTreeSink<'a> {
 
             match self.cursor.token_tree() {
                 Some(tt::TokenTree::Leaf(leaf)) => {
+                    // Mark the range if needed
+                    if let tt::Leaf::Ident(ident) = leaf {
+                        if kind == IDENT {
+                            let range =
+                                TextRange::offset_len(self.text_pos, TextUnit::of_str(&ident.text));
+                            self.range_map.add(range, ident.id);
+                        }
+                    }
+
                     self.cursor = self.cursor.bump();
                     self.buf += &format!("{}", leaf);
                 }
@@ -337,6 +359,7 @@ impl<'a> TreeSink for TtTreeSink<'a> {
         {
             if curr.spacing == tt::Spacing::Alone {
                 self.inner.token(WHITESPACE, " ".into());
+                self.text_pos += TextUnit::of_char(' ');
             }
         }
     }