4 files changed, 35 insertions, 66 deletions

diff --git a/crates/ra_hir_expand/src/db.rs b/crates/ra_hir_expand/src/db.rs
index e6d2e8d9d..b789c6e7b 100644
--- a/crates/ra_hir_expand/src/db.rs
+++ b/crates/ra_hir_expand/src/db.rs
@@ -22,8 +22,8 @@ pub trait AstDatabase: SourceDatabase {
 
     #[salsa::interned]
     fn intern_macro(&self, macro_call: MacroCallLoc) -> MacroCallId;
-    fn macro_arg(&self, id: MacroCallId) -> Option<(Arc<tt::Subtree>, Arc<mbe::TokenMap>)>;
-    fn macro_def(&self, id: MacroDefId) -> Option<(Arc<mbe::MacroRules>, Arc<mbe::TokenMap>)>;
+    fn macro_arg(&self, id: MacroCallId) -> Option<Arc<(tt::Subtree, mbe::TokenMap)>>;
+    fn macro_def(&self, id: MacroDefId) -> Option<Arc<(mbe::MacroRules, mbe::TokenMap)>>;
     fn parse_macro(
         &self,
         macro_file: MacroFile,
@@ -40,7 +40,7 @@ pub(crate) fn ast_id_map(db: &dyn AstDatabase, file_id: HirFileId) -> Arc<AstIdM
 pub(crate) fn macro_def(
     db: &dyn AstDatabase,
     id: MacroDefId,
-) -> Option<(Arc<mbe::MacroRules>, Arc<mbe::TokenMap>)> {
+) -> Option<Arc<(mbe::MacroRules, mbe::TokenMap)>> {
     let macro_call = id.ast_id.to_node(db);
     let arg = macro_call.token_tree()?;
     let (tt, tmap) = mbe::ast_to_token_tree(&arg).or_else(|| {
@@ -51,18 +51,18 @@ pub(crate) fn macro_def(
         log::warn!("fail on macro_def parse: {:#?}", tt);
         None
     })?;
-    Some((Arc::new(rules), Arc::new(tmap)))
+    Some(Arc::new((rules, tmap)))
 }
 
 pub(crate) fn macro_arg(
     db: &dyn AstDatabase,
     id: MacroCallId,
-) -> Option<(Arc<tt::Subtree>, Arc<mbe::TokenMap>)> {
+) -> Option<Arc<(tt::Subtree, mbe::TokenMap)>> {
     let loc = db.lookup_intern_macro(id);
     let macro_call = loc.ast_id.to_node(db);
     let arg = macro_call.token_tree()?;
     let (tt, tmap) = mbe::ast_to_token_tree(&arg)?;
-    Some((Arc::new(tt), Arc::new(tmap)))
+    Some(Arc::new((tt, tmap)))
 }
 
 pub(crate) fn macro_expand(
diff --git a/crates/ra_hir_expand/src/lib.rs b/crates/ra_hir_expand/src/lib.rs
index 1a1d6bdd4..b219b8fbf 100644
--- a/crates/ra_hir_expand/src/lib.rs
+++ b/crates/ra_hir_expand/src/lib.rs
@@ -12,11 +12,12 @@ pub mod hygiene;
 pub mod diagnostics;
 
 use std::hash::{Hash, Hasher};
+use std::sync::Arc;
 
 use ra_db::{salsa, CrateId, FileId};
 use ra_syntax::{
     ast::{self, AstNode},
-    SyntaxNode, TextRange,
+    SyntaxNode, TextRange, TextUnit,
 };
 
 use crate::ast_id_map::FileAstId;
@@ -68,29 +69,25 @@ impl HirFileId {
     }
 
     /// Return expansion information if it is a macro-expansion file
-    pub fn parent_expansion(self, db: &dyn db::AstDatabase) -> Option<ExpansionInfo> {
+    pub fn expansion_info(self, db: &dyn db::AstDatabase) -> Option<ExpansionInfo> {
         match self.0 {
             HirFileIdRepr::FileId(_) => None,
             HirFileIdRepr::MacroFile(macro_file) => {
                 let loc: MacroCallLoc = db.lookup_intern_macro(macro_file.macro_call_id);
 
-                let arg_range = loc.ast_id.to_node(db).token_tree()?.syntax().text_range();
-                let def_range = loc.def.ast_id.to_node(db).token_tree()?.syntax().text_range();
+                let arg_start = loc.ast_id.to_node(db).token_tree()?.syntax().text_range().start();
+                let def_start =
+                    loc.def.ast_id.to_node(db).token_tree()?.syntax().text_range().start();
 
                 let macro_def = db.macro_def(loc.def)?;
                 let shift = macro_def.0.shift();
-                let rev_map = db.parse_macro(macro_file)?.1;
+                let exp_map = db.parse_macro(macro_file)?.1;
+                let macro_arg = db.macro_arg(macro_file.macro_call_id)?;
 
-                let arg_token_map = db.macro_arg(macro_file.macro_call_id)?.1;
-                let def_token_map = macro_def.1;
+                let arg_start = (loc.ast_id.file_id, arg_start);
+                let def_start = (loc.def.ast_id.file_id, def_start);
 
-                let arg_map = rev_map.map_ranges(&arg_token_map, arg_range, shift);
-                let def_map = rev_map.map_ranges(&def_token_map, def_range, 0);
-
-                let arg_file = loc.ast_id.file_id;
-                let def_file = loc.def.ast_id.file_id;
-
-                Some(ExpansionInfo { arg_file, def_file, arg_map, def_map })
+                Some(ExpansionInfo { arg_start, def_start, macro_arg, macro_def, exp_map, shift })
             }
         }
     }
@@ -143,28 +140,30 @@ impl MacroCallId {
 #[derive(Debug, Clone, PartialEq, Eq)]
 /// ExpansionInfo mainly describes how to map text range between src and expanded macro
 pub struct ExpansionInfo {
-    pub(crate) arg_file: HirFileId,
-    pub(crate) def_file: HirFileId,
+    pub(crate) arg_start: (HirFileId, TextUnit),
+    pub(crate) def_start: (HirFileId, TextUnit),
+    pub(crate) shift: u32,
 
-    pub(crate) arg_map: Vec<(TextRange, TextRange)>,
-    pub(crate) def_map: Vec<(TextRange, TextRange)>,
+    pub(crate) macro_def: Arc<(mbe::MacroRules, mbe::TokenMap)>,
+    pub(crate) macro_arg: Arc<(tt::Subtree, mbe::TokenMap)>,
+    pub(crate) exp_map: Arc<mbe::RevTokenMap>,
 }
 
 impl ExpansionInfo {
     pub fn find_range(&self, from: TextRange) -> Option<(HirFileId, TextRange)> {
-        for (src, dest) in &self.arg_map {
-            if src.is_subrange(&from) {
-                return Some((self.arg_file, *dest));
-            }
+        fn look_in_rev_map(exp_map: &mbe::RevTokenMap, from: TextRange) -> Option<tt::TokenId> {
+            exp_map.ranges.iter().find(|&it| it.0.is_subrange(&from)).map(|it| it.1)
         }
 
-        for (src, dest) in &self.def_map {
-            if src.is_subrange(&from) {
-                return Some((self.def_file, *dest));
-            }
-        }
+        let token_id = look_in_rev_map(&self.exp_map, from)?;
+        let (token_map, file_offset, token_id) = if token_id.0 >= self.shift {
+            (&self.macro_arg.1, self.arg_start, tt::TokenId(token_id.0 - self.shift).into())
+        } else {
+            (&self.macro_def.1, self.def_start, token_id)
+        };
 
-        None
+        let range = token_map.relative_range_of(token_id)?;
+        Some((file_offset.0, TextRange::offset_len(range.start() + file_offset.1, range.len())))
     }
 }
 
diff --git a/crates/ra_ide_api/src/display/navigation_target.rs b/crates/ra_ide_api/src/display/navigation_target.rs
index b77e19231..1bf81e7d5 100644
--- a/crates/ra_ide_api/src/display/navigation_target.rs
+++ b/crates/ra_ide_api/src/display/navigation_target.rs
@@ -36,7 +36,7 @@ fn find_range_from_node(
 ) -> (FileId, TextRange) {
     let text_range = node.text_range();
     let (file_id, text_range) = src
-        .parent_expansion(db)
+        .expansion_info(db)
         .and_then(|expansion_info| expansion_info.find_range(text_range))
         .unwrap_or((src, text_range));
 
diff --git a/crates/ra_mbe/src/syntax_bridge.rs b/crates/ra_mbe/src/syntax_bridge.rs
index a49e63ace..9653f7fef 100644
--- a/crates/ra_mbe/src/syntax_bridge.rs
+++ b/crates/ra_mbe/src/syntax_bridge.rs
@@ -23,7 +23,7 @@ pub struct TokenMap {
 /// Maps relative range of the expanded syntax node to `tt::TokenId`
 #[derive(Debug, PartialEq, Eq, Default)]
 pub struct RevTokenMap {
-    ranges: Vec<(TextRange, tt::TokenId)>,
+    pub ranges: Vec<(TextRange, tt::TokenId)>,
 }
 
 /// Convert the syntax tree (what user has written) to a `TokenTree` (what macro
@@ -121,36 +121,6 @@ impl RevTokenMap {
     fn add(&mut self, relative_range: TextRange, token_id: tt::TokenId) {
         self.ranges.push((relative_range, token_id.clone()))
     }
-
-    /// Map a given token map to (Expanded syntax node, Input tokens) text-ranges pair
-    ///
-    /// This function do the following things:
-    ///
-    /// 1. Undo the increment of token-id `shift`:
-    ///     When we output a token from from macro argument, we increased its id
-    ///     by `shift` (so it's guaranteed to not to collide with anything from the definition)
-    ///     We undo the increment here to rollback to its original token id.
-    /// 2. Offset the input tokens (`to`) by `parent` text-range:
-    ///     We transforms the input tokens text-ranges from relative to original first token
-    ///     to parent text-range
-    /// 3. Maps expanded tokens text-ranges to parent text-ranges
-    ///
-    pub fn map_ranges(
-        &self,
-        to: &TokenMap,
-        parent: TextRange,
-        shift: u32,
-    ) -> Vec<(TextRange, TextRange)> {
-        self.ranges
-            .iter()
-            .filter_map(|(r, tid)| {
-                let adjusted_id = tt::TokenId(tid.0.checked_sub(shift)?);
-                let to_range = to.relative_range_of(adjusted_id)?;
-
-                Some((*r, TextRange::offset_len(to_range.start() + parent.start(), to_range.len())))
-            })
-            .collect()
-    }
 }
 
 /// Returns the textual content of a doc comment block as a quoted string