Fix original_source find order

1 files changed, 22 insertions, 7 deletions

diff --git a/crates/ra_hir_expand/src/lib.rs b/crates/ra_hir_expand/src/lib.rs
index 94e1e466a..d1a43fe6c 100644
--- a/crates/ra_hir_expand/src/lib.rs
+++ b/crates/ra_hir_expand/src/lib.rs
@@ -214,7 +214,17 @@ pub struct ExpansionInfo {
     exp_map: Arc<mbe::TokenMap>,
 }
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum ExpansionOrigin {
+    Call,
+    Def,
+}
+
 impl ExpansionInfo {
+    pub fn call_node(&self) -> Option<InFile<SyntaxNode>> {
+        Some(self.arg.with_value(self.arg.value.parent()?))
+    }
+
     pub fn map_token_down(&self, token: InFile<&SyntaxToken>) -> Option<InFile<SyntaxToken>> {
         assert_eq!(token.file_id, self.arg.file_id);
         let range = token.value.text_range().checked_sub(self.arg.value.text_range().start())?;
@@ -228,21 +238,26 @@ impl ExpansionInfo {
         Some(self.expanded.with_value(token))
     }
 
-    pub fn map_token_up(&self, token: InFile<&SyntaxToken>) -> Option<InFile<SyntaxToken>> {
+    pub fn map_token_up(
+        &self,
+        token: InFile<&SyntaxToken>,
+    ) -> Option<(InFile<SyntaxToken>, ExpansionOrigin)> {
         let token_id = self.exp_map.token_by_range(token.value.text_range())?;
 
         let (token_id, origin) = self.macro_def.0.map_id_up(token_id);
-        let (token_map, tt) = match origin {
-            mbe::Origin::Call => (&self.macro_arg.1, self.arg.clone()),
-            mbe::Origin::Def => {
-                (&self.macro_def.1, self.def.as_ref().map(|tt| tt.syntax().clone()))
-            }
+        let (token_map, tt, origin) = match origin {
+            mbe::Origin::Call => (&self.macro_arg.1, self.arg.clone(), ExpansionOrigin::Call),
+            mbe::Origin::Def => (
+                &self.macro_def.1,
+                self.def.as_ref().map(|tt| tt.syntax().clone()),
+                ExpansionOrigin::Def,
+            ),
         };
 
         let range = token_map.range_by_token(token_id)?;
         let token = algo::find_covering_element(&tt.value, range + tt.value.text_range().start())
             .into_token()?;
-        Some(tt.with_value(token))
+        Some((tt.with_value(token), origin))
     }
 }