3 files changed, 89 insertions, 17 deletions

diff --git a/crates/ra_hir_expand/Cargo.toml b/crates/ra_hir_expand/Cargo.toml
index 9bf5b7918..8f29bf7d9 100644
--- a/crates/ra_hir_expand/Cargo.toml
+++ b/crates/ra_hir_expand/Cargo.toml
@@ -10,6 +10,7 @@ log = "0.4.5"
 ra_arena = { path = "../ra_arena" }
 ra_db = { path = "../ra_db" }
 ra_syntax = { path = "../ra_syntax" }
+ra_parser = { path = "../ra_parser" }
 ra_prof = { path = "../ra_prof" }
 tt = { path = "../ra_tt", package = "ra_tt" }
 mbe = { path = "../ra_mbe", package = "ra_mbe" }
diff --git a/crates/ra_hir_expand/src/db.rs b/crates/ra_hir_expand/src/db.rs
index a4ee9a529..b4dafe1d8 100644
--- a/crates/ra_hir_expand/src/db.rs
+++ b/crates/ra_hir_expand/src/db.rs
@@ -4,6 +4,7 @@ use std::sync::Arc;
 
 use mbe::MacroRules;
 use ra_db::{salsa, SourceDatabase};
+use ra_parser::FragmentKind;
 use ra_prof::profile;
 use ra_syntax::{AstNode, Parse, SyntaxNode};
 
@@ -22,9 +23,12 @@ pub trait AstDatabase: SourceDatabase {
 
     #[salsa::interned]
     fn intern_macro(&self, macro_call: MacroCallLoc) -> MacroCallId;
-    fn macro_arg(&self, id: MacroCallId) -> Option<Arc<tt::Subtree>>;
-    fn macro_def(&self, id: MacroDefId) -> Option<Arc<mbe::MacroRules>>;
-    fn parse_macro(&self, macro_file: MacroFile) -> Option<Parse<SyntaxNode>>;
+    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,
+    ) -> Option<(Parse<SyntaxNode>, Arc<mbe::RevTokenMap>)>;
     fn macro_expand(&self, macro_call: MacroCallId) -> Result<Arc<tt::Subtree>, String>;
 }
 
@@ -34,10 +38,13 @@ pub(crate) fn ast_id_map(db: &dyn AstDatabase, file_id: HirFileId) -> Arc<AstIdM
     Arc::new(map)
 }
 
-pub(crate) fn macro_def(db: &dyn AstDatabase, id: MacroDefId) -> Option<Arc<MacroRules>> {
+pub(crate) fn macro_def(
+    db: &dyn AstDatabase,
+    id: MacroDefId,
+) -> Option<Arc<(mbe::MacroRules, mbe::TokenMap)>> {
     let macro_call = id.ast_id.to_node(db);
     let arg = macro_call.token_tree()?;
-    let (tt, _) = mbe::ast_to_token_tree(&arg).or_else(|| {
+    let (tt, tmap) = mbe::ast_to_token_tree(&arg).or_else(|| {
         log::warn!("fail on macro_def to token tree: {:#?}", arg);
         None
     })?;
@@ -45,15 +52,18 @@ pub(crate) fn macro_def(db: &dyn AstDatabase, id: MacroDefId) -> Option<Arc<Macr
         log::warn!("fail on macro_def parse: {:#?}", tt);
         None
     })?;
-    Some(Arc::new(rules))
+    Some(Arc::new((rules, tmap)))
 }
 
-pub(crate) fn macro_arg(db: &dyn AstDatabase, id: MacroCallId) -> Option<Arc<tt::Subtree>> {
+pub(crate) fn macro_arg(
+    db: &dyn AstDatabase,
+    id: MacroCallId,
+) -> 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, _) = mbe::ast_to_token_tree(&arg)?;
-    Some(Arc::new(tt))
+    let (tt, tmap) = mbe::ast_to_token_tree(&arg)?;
+    Some(Arc::new((tt, tmap)))
 }
 
 pub(crate) fn macro_expand(
@@ -64,7 +74,7 @@ pub(crate) fn macro_expand(
     let macro_arg = db.macro_arg(id).ok_or("Fail to args in to tt::TokenTree")?;
 
     let macro_rules = db.macro_def(loc.def).ok_or("Fail to find macro definition")?;
-    let tt = macro_rules.expand(&macro_arg).map_err(|err| format!("{:?}", err))?;
+    let tt = macro_rules.0.expand(&macro_arg.0).map_err(|err| format!("{:?}", err))?;
     // Set a hard limit for the expanded tt
     let count = tt.count();
     if count > 65536 {
@@ -77,7 +87,7 @@ pub(crate) fn parse_or_expand(db: &dyn AstDatabase, file_id: HirFileId) -> Optio
     match file_id.0 {
         HirFileIdRepr::FileId(file_id) => Some(db.parse(file_id).tree().syntax().clone()),
         HirFileIdRepr::MacroFile(macro_file) => {
-            db.parse_macro(macro_file).map(|it| it.syntax_node())
+            db.parse_macro(macro_file).map(|(it, _)| it.syntax_node())
         }
     }
 }
@@ -85,8 +95,9 @@ pub(crate) fn parse_or_expand(db: &dyn AstDatabase, file_id: HirFileId) -> Optio
 pub(crate) fn parse_macro(
     db: &dyn AstDatabase,
     macro_file: MacroFile,
-) -> Option<Parse<SyntaxNode>> {
+) -> Option<(Parse<SyntaxNode>, Arc<mbe::RevTokenMap>)> {
     let _p = profile("parse_macro_query");
+
     let macro_call_id = macro_file.macro_call_id;
     let tt = db
         .macro_expand(macro_call_id)
@@ -97,8 +108,11 @@ pub(crate) fn parse_macro(
             log::warn!("fail on macro_parse: (reason: {})", err,);
         })
         .ok()?;
-    match macro_file.macro_file_kind {
-        MacroFileKind::Items => mbe::token_tree_to_items(&tt).ok().map(Parse::to_syntax),
-        MacroFileKind::Expr => mbe::token_tree_to_expr(&tt).ok().map(Parse::to_syntax),
-    }
+
+    let fragment_kind = match macro_file.macro_file_kind {
+        MacroFileKind::Items => FragmentKind::Items,
+        MacroFileKind::Expr => FragmentKind::Expr,
+    };
+    let (parse, rev_token_map) = mbe::token_tree_to_syntax_node(&tt, fragment_kind).ok()?;
+    Some((parse, Arc::new(rev_token_map)))
 }
diff --git a/crates/ra_hir_expand/src/lib.rs b/crates/ra_hir_expand/src/lib.rs
index dd07a16b4..151d1d785 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,
+    SyntaxNode, TextRange, TextUnit,
 };
 
 use crate::ast_id_map::FileAstId;
@@ -66,6 +67,30 @@ impl HirFileId {
             }
         }
     }
+
+    /// Return expansion information if it is a macro-expansion file
+    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_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 exp_map = db.parse_macro(macro_file)?.1;
+                let macro_arg = db.macro_arg(macro_file.macro_call_id)?;
+
+                let arg_start = (loc.ast_id.file_id, arg_start);
+                let def_start = (loc.def.ast_id.file_id, def_start);
+
+                Some(ExpansionInfo { arg_start, def_start, macro_arg, macro_def, exp_map, shift })
+            }
+        }
+    }
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
@@ -112,6 +137,38 @@ 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_start: (HirFileId, TextUnit),
+    pub(crate) def_start: (HirFileId, TextUnit),
+    pub(crate) shift: u32,
+
+    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)> {
+        let token_id = look_in_rev_map(&self.exp_map, from)?;
+
+        let (token_map, (file_id, start_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)
+        };
+
+        let range = token_map.relative_range_of(token_id)?;
+
+        return Some((file_id, range + start_offset));
+
+        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)
+        }
+    }
+}
+
 /// `AstId` points to an AST node in any file.
 ///
 /// It is stable across reparses, and can be used as salsa key/value.