1 files changed, 36 insertions, 17 deletions

diff --git a/crates/ra_ide/src/expand.rs b/crates/ra_ide/src/expand.rs
index 216d5cfec..661628ae4 100644
--- a/crates/ra_ide/src/expand.rs
+++ b/crates/ra_ide/src/expand.rs
@@ -3,7 +3,7 @@ use std::iter::successors;
 
 use hir::InFile;
 use ra_db::FileId;
-use ra_syntax::{ast, AstNode, SyntaxNode, SyntaxToken};
+use ra_syntax::{ast, AstNode, SyntaxNode, SyntaxToken, TextRange};
 
 use crate::{db::RootDatabase, FileRange};
 
@@ -17,26 +17,45 @@ pub(crate) fn original_range(db: &RootDatabase, node: InFile<&SyntaxNode>) -> Fi
         }
         Some(it) => it,
     };
-    // FIXME: the following completely wrong.
-    //
-    // *First*, we should try to map first and last tokens of node, and, if that
-    // fails, return the range of the overall macro expansions.
-    //
-    // *Second*, we should handle recurside macro expansions
-
-    let token = node
-        .value
-        .descendants_with_tokens()
-        .filter_map(|it| it.into_token())
-        .find_map(|it| expansion.map_token_up(node.with_value(&it)));
-
-    match token {
-        Some(it) => {
-            FileRange { file_id: it.file_id.original_file(db), range: it.value.text_range() }
+    // FIXME: We should handle recurside macro expansions
+
+    let range = node.value.descendants_with_tokens().find_map(|it| {
+        match it.as_token() {
+            // FIXME: Remove this branch after all `tt::TokenTree`s have a proper `TokenId`,
+            // and return the range of the overall macro expansions if mapping first and last tokens fails.
+            Some(token) => {
+                let token = expansion.map_token_up(node.with_value(&token))?;
+                Some(token.with_value(token.value.text_range()))
+            }
+            None => {
+                // Try to map first and last tokens of node, and, if success, return the union range of mapped tokens
+                let n = it.into_node()?;
+                let first = expansion.map_token_up(node.with_value(&n.first_token()?))?;
+                let last = expansion.map_token_up(node.with_value(&n.last_token()?))?;
+
+                // FIXME: Is is possible ?
+                if first.file_id != last.file_id {
+                    return None;
+                }
+
+                // FIXME: Add union method in TextRange
+                let range = union_range(first.value.text_range(), last.value.text_range());
+                Some(first.with_value(range))
+            }
         }
+    });
+
+    return match range {
+        Some(it) => FileRange { file_id: it.file_id.original_file(db), range: it.value },
         None => {
             FileRange { file_id: node.file_id.original_file(db), range: node.value.text_range() }
         }
+    };
+
+    fn union_range(a: TextRange, b: TextRange) -> TextRange {
+        let start = a.start().min(b.start());
+        let end = a.end().max(b.end());
+        TextRange::from_to(start, end)
     }
 }