Fix original_source find order

3 files changed, 85 insertions, 47 deletions

diff --git a/crates/ra_ide/src/display/navigation_target.rs b/crates/ra_ide/src/display/navigation_target.rs
index 6a6b49afd..6a2bf7273 100644
--- a/crates/ra_ide/src/display/navigation_target.rs
+++ b/crates/ra_ide/src/display/navigation_target.rs
@@ -7,10 +7,14 @@ use ra_syntax::{
     ast::{self, DocCommentsOwner, NameOwner},
     match_ast, AstNode, SmolStr,
     SyntaxKind::{self, BIND_PAT, TYPE_PARAM},
-    TextRange,
+    SyntaxNode, TextRange,
 };
 
-use crate::{db::RootDatabase, expand::original_range, FileSymbol};
+use crate::{
+    db::RootDatabase,
+    expand::{original_range_by_kind, OriginalRangeKind},
+    FileRange, FileSymbol,
+};
 
 use super::short_label::ShortLabel;
 
@@ -416,3 +420,14 @@ pub(crate) fn description_from_symbol(db: &RootDatabase, symbol: &FileSymbol) ->
         }
     }
 }
+
+fn original_range(db: &RootDatabase, node: InFile<&SyntaxNode>) -> FileRange {
+    if let Some(range) = original_range_by_kind(db, node, OriginalRangeKind::CallToken) {
+        return range;
+    }
+    if let Some(range) = original_range_by_kind(db, node, OriginalRangeKind::WholeCall) {
+        return range;
+    }
+
+    FileRange { file_id: node.file_id.original_file(db), range: node.value.text_range() }
+}
diff --git a/crates/ra_ide/src/expand.rs b/crates/ra_ide/src/expand.rs
index 661628ae4..327393dbb 100644
--- a/crates/ra_ide/src/expand.rs
+++ b/crates/ra_ide/src/expand.rs
@@ -1,57 +1,62 @@
 //! Utilities to work with files, produced by macros.
 use std::iter::successors;
 
-use hir::InFile;
+use hir::{ExpansionOrigin, InFile};
 use ra_db::FileId;
 use ra_syntax::{ast, AstNode, SyntaxNode, SyntaxToken, TextRange};
 
 use crate::{db::RootDatabase, FileRange};
 
-pub(crate) fn original_range(db: &RootDatabase, node: InFile<&SyntaxNode>) -> FileRange {
-    let expansion = match node.file_id.expansion_info(db) {
-        None => {
-            return FileRange {
-                file_id: node.file_id.original_file(db),
-                range: node.value.text_range(),
-            }
-        }
-        Some(it) => it,
-    };
+#[derive(Debug, PartialEq, Eq)]
+pub(crate) enum OriginalRangeKind {
+    /// Return range if any token is matched
+    #[allow(dead_code)]
+    Any,
+    /// Return range if token is inside macro_call
+    CallToken,
+    /// Return whole macro call range if matched
+    WholeCall,
+}
+
+pub(crate) fn original_range_by_kind(
+    db: &RootDatabase,
+    node: InFile<&SyntaxNode>,
+    kind: OriginalRangeKind,
+) -> Option<FileRange> {
+    let expansion = node.file_id.expansion_info(db)?;
+
+    // the input node has only one token ?
+    let single = node.value.first_token()? == node.value.last_token()?;
+
     // FIXME: We should handle recurside macro expansions
+    let range = match kind {
+        OriginalRangeKind::WholeCall => expansion.call_node()?.map(|node| node.text_range()),
+        _ => node.value.descendants().find_map(|it| {
+            let first = it.first_token()?;
+            let last = it.last_token()?;
 
-    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()))
+            if !single && first == last {
+                return None;
             }
-            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;
-                }
+            // Try to map first and last tokens of node, and, if success, return the union range of mapped tokens
+            let (first, first_origin) = expansion.map_token_up(node.with_value(&first))?;
+            let (last, last_origin) = expansion.map_token_up(node.with_value(&last))?;
 
-                // FIXME: Add union method in TextRange
-                let range = union_range(first.value.text_range(), last.value.text_range());
-                Some(first.with_value(range))
+            if first.file_id != last.file_id
+                || first_origin != last_origin
+                || (kind == OriginalRangeKind::CallToken && first_origin != ExpansionOrigin::Call)
+            {
+                return None;
             }
-        }
-    });
 
-    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() }
-        }
+            // FIXME: Add union method in TextRange
+            Some(first.with_value(union_range(first.value.text_range(), last.value.text_range())))
+        })?,
     };
 
+    return Some(FileRange { file_id: range.file_id.original_file(db), range: range.value });
+
     fn union_range(a: TextRange, b: TextRange) -> TextRange {
         let start = a.start().min(b.start());
         let end = a.end().max(b.end());
diff --git a/crates/ra_ide/src/goto_definition.rs b/crates/ra_ide/src/goto_definition.rs
index cfe62037f..2c634990d 100644
--- a/crates/ra_ide/src/goto_definition.rs
+++ b/crates/ra_ide/src/goto_definition.rs
@@ -209,7 +209,7 @@ fn named_target(db: &RootDatabase, node: InFile<&SyntaxNode>) -> Option<Navigati
 
 #[cfg(test)]
 mod tests {
-    use test_utils::covers;
+    use test_utils::{assert_eq_text, covers};
 
     use crate::mock_analysis::analysis_and_position;
 
@@ -222,6 +222,24 @@ mod tests {
         nav.assert_match(expected);
     }
 
+    fn check_goto_with_range_content(fixture: &str, expected: &str, expected_range: &str) {
+        let (analysis, pos) = analysis_and_position(fixture);
+
+        let mut navs = analysis.goto_definition(pos).unwrap().unwrap().info;
+        assert_eq!(navs.len(), 1);
+        let nav = navs.pop().unwrap();
+        let file_text = analysis.file_text(pos.file_id).unwrap();
+
+        let actual_full_range = &file_text[nav.full_range()];
+        let actual_range = &file_text[nav.range()];
+
+        test_utils::assert_eq_text!(
+            &format!("{}|{}", actual_full_range, actual_range),
+            expected_range
+        );
+        nav.assert_match(expected);
+    }
+
     #[test]
     fn goto_definition_works_in_items() {
         check_goto(
@@ -339,28 +357,27 @@ mod tests {
 
     #[test]
     fn goto_definition_works_for_macro_defined_fn_with_arg() {
-        check_goto(
+        check_goto_with_range_content(
             "
             //- /lib.rs
             macro_rules! define_fn {
                 ($name:ident) => (fn $name() {})
             }
 
-            define_fn!(
-                foo
-            )
+            define_fn!(foo);
 
             fn bar() {
                <|>foo();
             }
             ",
-            "foo FN_DEF FileId(1) [80; 83) [80; 83)",
+            "foo FN_DEF FileId(1) [64; 80) [75; 78)",
+            "define_fn!(foo);|foo",
         );
     }
 
     #[test]
     fn goto_definition_works_for_macro_defined_fn_no_arg() {
-        check_goto(
+        check_goto_with_range_content(
             "
             //- /lib.rs
             macro_rules! define_fn {
@@ -373,7 +390,8 @@ mod tests {
                <|>foo();
             }
             ",
-            "foo FN_DEF FileId(1) [39; 42) [39; 42)",
+            "foo FN_DEF FileId(1) [51; 64) [51; 64)",
+            "define_fn!();|define_fn!();",
         );
     }