Add macro_expansion_info in hir_expand

5 files changed, 212 insertions, 47 deletions

diff --git a/crates/ra_hir_expand/src/db.rs b/crates/ra_hir_expand/src/db.rs
index a4ee9a529..8abfbb4ff 100644
--- a/crates/ra_hir_expand/src/db.rs
+++ b/crates/ra_hir_expand/src/db.rs
@@ -8,10 +8,16 @@ use ra_prof::profile;
 use ra_syntax::{AstNode, Parse, SyntaxNode};
 
 use crate::{
-    ast_id_map::AstIdMap, HirFileId, HirFileIdRepr, MacroCallId, MacroCallLoc, MacroDefId,
-    MacroFile, MacroFileKind,
+    ast_id_map::AstIdMap, ExpansionInfo, HirFileId, HirFileIdRepr, MacroCallId, MacroCallLoc,
+    MacroDefId, MacroFile, MacroFileKind,
 };
 
+#[derive(Debug, PartialEq, Eq, Clone)]
+pub struct ParseMacroWithInfo {
+    pub parsed: Parse<SyntaxNode>,
+    pub expansion_info: Arc<ExpansionInfo>,
+}
+
 // FIXME: rename to ExpandDatabase
 #[salsa::query_group(AstDatabaseStorage)]
 pub trait AstDatabase: SourceDatabase {
@@ -22,10 +28,16 @@ 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 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 parse_macro(&self, macro_file: MacroFile) -> Option<Parse<SyntaxNode>>;
-    fn macro_expand(&self, macro_call: MacroCallId) -> Result<Arc<tt::Subtree>, String>;
+    fn parse_macro_with_info(&self, macro_file: MacroFile) -> Option<ParseMacroWithInfo>;
+    fn macro_expand(
+        &self,
+        macro_call: MacroCallId,
+    ) -> Result<(Arc<tt::Subtree>, (Arc<mbe::TokenMap>, Arc<mbe::TokenMap>)), String>;
+
+    fn macro_expansion_info(&self, macro_file: MacroFile) -> Option<Arc<ExpansionInfo>>;
 }
 
 pub(crate) fn ast_id_map(db: &dyn AstDatabase, file_id: HirFileId) -> Arc<AstIdMap> {
@@ -34,10 +46,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>, Arc<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,32 +60,36 @@ 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), Arc::new(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>, Arc<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), Arc::new(tmap)))
 }
 
 pub(crate) fn macro_expand(
     db: &dyn AstDatabase,
     id: MacroCallId,
-) -> Result<Arc<tt::Subtree>, String> {
+) -> Result<(Arc<tt::Subtree>, (Arc<mbe::TokenMap>, Arc<mbe::TokenMap>)), String> {
     let loc = db.lookup_intern_macro(id);
     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 {
         return Err(format!("Total tokens count exceed limit : count = {}", count));
     }
-    Ok(Arc::new(tt))
+
+    Ok((Arc::new(tt), (macro_arg.1.clone(), macro_rules.1.clone())))
 }
 
 pub(crate) fn parse_or_expand(db: &dyn AstDatabase, file_id: HirFileId) -> Option<SyntaxNode> {
@@ -87,6 +106,13 @@ pub(crate) fn parse_macro(
     macro_file: MacroFile,
 ) -> Option<Parse<SyntaxNode>> {
     let _p = profile("parse_macro_query");
+    db.parse_macro_with_info(macro_file).map(|r| r.parsed)
+}
+
+pub(crate) fn parse_macro_with_info(
+    db: &dyn AstDatabase,
+    macro_file: MacroFile,
+) -> Option<ParseMacroWithInfo> {
     let macro_call_id = macro_file.macro_call_id;
     let tt = db
         .macro_expand(macro_call_id)
@@ -97,8 +123,39 @@ 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 res = match macro_file.macro_file_kind {
+        MacroFileKind::Items => {
+            mbe::token_tree_to_items(&tt.0).ok().map(|(p, map)| (Parse::to_syntax(p), map))
+        }
+        MacroFileKind::Expr => {
+            mbe::token_tree_to_expr(&tt.0).ok().map(|(p, map)| (Parse::to_syntax(p), map))
+        }
+    };
+
+    res.map(|(parsed, exp_map)| {
+        let (arg_map, def_map) = tt.1;
+        let loc: MacroCallLoc = db.lookup_intern_macro(macro_call_id);
+
+        let def_start =
+            loc.def.ast_id.to_node(db).token_tree().map(|t| t.syntax().text_range().start());
+        let arg_start =
+            loc.ast_id.to_node(db).token_tree().map(|t| t.syntax().text_range().start());
+
+        let arg_map =
+            arg_start.map(|start| exp_map.ranges(&arg_map, start)).unwrap_or_else(|| Vec::new());
+
+        let def_map =
+            def_start.map(|start| exp_map.ranges(&def_map, start)).unwrap_or_else(|| Vec::new());
+
+        let info = ExpansionInfo { arg_map, def_map };
+
+        ParseMacroWithInfo { parsed, expansion_info: Arc::new(info) }
+    })
+}
+
+pub(crate) fn macro_expansion_info(
+    db: &dyn AstDatabase,
+    macro_file: MacroFile,
+) -> Option<Arc<ExpansionInfo>> {
+    db.parse_macro_with_info(macro_file).map(|res| res.expansion_info.clone())
 }
diff --git a/crates/ra_hir_expand/src/lib.rs b/crates/ra_hir_expand/src/lib.rs
index dd07a16b4..194020b45 100644
--- a/crates/ra_hir_expand/src/lib.rs
+++ b/crates/ra_hir_expand/src/lib.rs
@@ -16,7 +16,7 @@ use std::hash::{Hash, Hasher};
 use ra_db::{salsa, CrateId, FileId};
 use ra_syntax::{
     ast::{self, AstNode},
-    SyntaxNode,
+    SyntaxNode, TextRange,
 };
 
 use crate::ast_id_map::FileAstId;
@@ -112,6 +112,39 @@ impl MacroCallId {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+/// ExpansionInfo mainly describle how to map text range between src and expaned macro
+pub struct ExpansionInfo {
+    pub arg_map: Vec<(TextRange, TextRange)>,
+    pub def_map: Vec<(TextRange, TextRange)>,
+}
+
+impl ExpansionInfo {
+    pub fn find_range(
+        &self,
+        from: TextRange,
+        (arg_file_id, def_file_id): (HirFileId, HirFileId),
+    ) -> Option<(HirFileId, TextRange)> {
+        for (src, dest) in &self.arg_map {
+            dbg!((src, *dest, "arg_map"));
+            if src.is_subrange(&from) {
+                dbg!((arg_file_id, *dest));
+                return Some((arg_file_id, *dest));
+            }
+        }
+
+        for (src, dest) in &self.def_map {
+            dbg!((src, *dest, "def_map"));
+            if src.is_subrange(&from) {
+                dbg!((arg_file_id, *dest));
+                return Some((def_file_id, *dest));
+            }
+        }
+
+        None
+    }
+}
+
 /// `AstId` points to an AST node in any file.
 ///
 /// It is stable across reparses, and can be used as salsa key/value.
diff --git a/crates/ra_mbe/src/lib.rs b/crates/ra_mbe/src/lib.rs
index 15f000175..2926b29fd 100644
--- a/crates/ra_mbe/src/lib.rs
+++ b/crates/ra_mbe/src/lib.rs
@@ -32,7 +32,7 @@ pub enum ExpandError {
 
 pub use crate::syntax_bridge::{
     ast_to_token_tree, syntax_node_to_token_tree, token_tree_to_expr, token_tree_to_items,
-    token_tree_to_macro_stmts, token_tree_to_pat, token_tree_to_ty,
+    token_tree_to_macro_stmts, token_tree_to_pat, token_tree_to_ty, TokenMap,
 };
 
 /// This struct contains AST for a single `macro_rules` definition. What might
diff --git a/crates/ra_mbe/src/syntax_bridge.rs b/crates/ra_mbe/src/syntax_bridge.rs
index 8d9217518..5db6647e3 100644
--- a/crates/ra_mbe/src/syntax_bridge.rs
+++ b/crates/ra_mbe/src/syntax_bridge.rs
@@ -15,6 +15,7 @@ use crate::ExpandError;
 use std::sync::atomic::{AtomicU32, Ordering};
 
 /// Maps `tt::TokenId` to the relative range of the original token.
+#[derive(Debug, PartialEq, Eq)]
 pub struct TokenMap {
     /// Maps `tt::TokenId` to the *relative* source range.
     tokens: Vec<TextRange>,
@@ -34,6 +35,13 @@ impl std::default::Default for TokenMap {
     }
 }
 
+/// Maps Relative range of the expanded syntax node to `tt::TokenId`
+#[derive(Debug, PartialEq, Eq, Default)]
+pub struct ExpandedRangeMap {
+    /// Maps `tt::TokenId` to the *relative* source range.
+    ranges: Vec<(TextRange, tt::TokenId)>,
+}
+
 /// Convert the syntax tree (what user has written) to a `TokenTree` (what macro
 /// will consume).
 pub fn ast_to_token_tree(ast: &ast::TokenTree) -> Option<(tt::Subtree, TokenMap)> {
@@ -66,7 +74,7 @@ pub fn syntax_node_to_token_tree(node: &SyntaxNode) -> Option<(tt::Subtree, Toke
 fn fragment_to_syntax_node(
     tt: &tt::Subtree,
     fragment_kind: FragmentKind,
-) -> Result<Parse<SyntaxNode>, ExpandError> {
+) -> Result<(Parse<SyntaxNode>, ExpandedRangeMap), ExpandError> {
     let tmp;
     let tokens = match tt {
         tt::Subtree { delimiter: tt::Delimiter::None, token_trees } => token_trees.as_slice(),
@@ -77,44 +85,55 @@ fn fragment_to_syntax_node(
     };
     let buffer = TokenBuffer::new(&tokens);
     let mut token_source = SubtreeTokenSource::new(&buffer);
-    let mut tree_sink = TtTreeSink::new(buffer.begin());
+    let mut range_map = ExpandedRangeMap::default();
+    let mut tree_sink = TtTreeSink::new(buffer.begin(), &mut range_map);
     ra_parser::parse_fragment(&mut token_source, &mut tree_sink, fragment_kind);
     if tree_sink.roots.len() != 1 {
         return Err(ExpandError::ConversionError);
     }
     //FIXME: would be cool to report errors
     let parse = tree_sink.inner.finish();
-    Ok(parse)
+    Ok((parse, range_map))
 }
 
 /// Parses the token tree (result of macro expansion) to an expression
-pub fn token_tree_to_expr(tt: &tt::Subtree) -> Result<Parse<ast::Expr>, ExpandError> {
-    let parse = fragment_to_syntax_node(tt, Expr)?;
-    parse.cast().ok_or_else(|| crate::ExpandError::ConversionError)
+pub fn token_tree_to_expr(
+    tt: &tt::Subtree,
+) -> Result<(Parse<ast::Expr>, ExpandedRangeMap), ExpandError> {
+    let (parse, map) = fragment_to_syntax_node(tt, Expr)?;
+    parse.cast().ok_or_else(|| crate::ExpandError::ConversionError).map(|p| (p, map))
 }
 
 /// Parses the token tree (result of macro expansion) to a Pattern
-pub fn token_tree_to_pat(tt: &tt::Subtree) -> Result<Parse<ast::Pat>, ExpandError> {
-    let parse = fragment_to_syntax_node(tt, Pattern)?;
-    parse.cast().ok_or_else(|| crate::ExpandError::ConversionError)
+pub fn token_tree_to_pat(
+    tt: &tt::Subtree,
+) -> Result<(Parse<ast::Pat>, ExpandedRangeMap), ExpandError> {
+    let (parse, map) = fragment_to_syntax_node(tt, Pattern)?;
+    parse.cast().ok_or_else(|| crate::ExpandError::ConversionError).map(|p| (p, map))
 }
 
 /// Parses the token tree (result of macro expansion) to a Type
-pub fn token_tree_to_ty(tt: &tt::Subtree) -> Result<Parse<ast::TypeRef>, ExpandError> {
-    let parse = fragment_to_syntax_node(tt, Type)?;
-    parse.cast().ok_or_else(|| crate::ExpandError::ConversionError)
+pub fn token_tree_to_ty(
+    tt: &tt::Subtree,
+) -> Result<(Parse<ast::TypeRef>, ExpandedRangeMap), ExpandError> {
+    let (parse, map) = fragment_to_syntax_node(tt, Type)?;
+    parse.cast().ok_or_else(|| crate::ExpandError::ConversionError).map(|p| (p, map))
 }
 
 /// Parses the token tree (result of macro expansion) as a sequence of stmts
-pub fn token_tree_to_macro_stmts(tt: &tt::Subtree) -> Result<Parse<ast::MacroStmts>, ExpandError> {
-    let parse = fragment_to_syntax_node(tt, Statements)?;
-    parse.cast().ok_or_else(|| crate::ExpandError::ConversionError)
+pub fn token_tree_to_macro_stmts(
+    tt: &tt::Subtree,
+) -> Result<(Parse<ast::MacroStmts>, ExpandedRangeMap), ExpandError> {
+    let (parse, map) = fragment_to_syntax_node(tt, Statements)?;
+    parse.cast().ok_or_else(|| crate::ExpandError::ConversionError).map(|p| (p, map))
 }
 
 /// Parses the token tree (result of macro expansion) as a sequence of items
-pub fn token_tree_to_items(tt: &tt::Subtree) -> Result<Parse<ast::MacroItems>, ExpandError> {
-    let parse = fragment_to_syntax_node(tt, Items)?;
-    parse.cast().ok_or_else(|| crate::ExpandError::ConversionError)
+pub fn token_tree_to_items(
+    tt: &tt::Subtree,
+) -> Result<(Parse<ast::MacroItems>, ExpandedRangeMap), ExpandError> {
+    let (parse, map) = fragment_to_syntax_node(tt, Items)?;
+    parse.cast().ok_or_else(|| crate::ExpandError::ConversionError).map(|p| (p, map))
 }
 
 impl TokenMap {
@@ -133,6 +152,28 @@ impl TokenMap {
     }
 }
 
+impl ExpandedRangeMap {
+    fn set(&mut self, relative_range: TextRange, token_id: &tt::TokenId) {
+        self.ranges.push((relative_range, token_id.clone()))
+    }
+
+    pub fn ranges(&self, to: &TokenMap) -> Vec<(TextRange, TextRange)> {
+        self.ranges
+            .iter()
+            .filter_map(|(r, tid)| {
+                if to.map_id == tid.map_id() {
+                    return None;
+                }
+                if let Some(to_range) = to.relative_range_of(*tid) {
+                    Some((*r, to_range))
+                } else {
+                    None
+                }
+            })
+            .collect()
+    }
+}
+
 /// Returns the textual content of a doc comment block as a quoted string
 /// That is, strips leading `///` (or `/**`, etc)
 /// and strips the ending `*/`
@@ -279,6 +320,8 @@ struct TtTreeSink<'a> {
     cursor: Cursor<'a>,
     text_pos: TextUnit,
     inner: SyntaxTreeBuilder,
+    range_marker: Option<(TextRange, tt::TokenId)>,
+    range_map: &'a mut ExpandedRangeMap,
 
     // Number of roots
     // Use for detect ill-form tree which is not single root
@@ -286,13 +329,15 @@ struct TtTreeSink<'a> {
 }
 
 impl<'a> TtTreeSink<'a> {
-    fn new(cursor: Cursor<'a>) -> Self {
+    fn new(cursor: Cursor<'a>, range_map: &'a mut ExpandedRangeMap) -> Self {
         TtTreeSink {
             buf: String::new(),
             cursor,
             text_pos: 0.into(),
             inner: SyntaxTreeBuilder::default(),
             roots: smallvec::SmallVec::new(),
+            range_map,
+            range_marker: None,
         }
     }
 }
@@ -317,6 +362,8 @@ impl<'a> TreeSink for TtTreeSink<'a> {
             return;
         }
 
+        let mut last_ident = None;
+
         for _ in 0..n_tokens {
             if self.cursor.eof() {
                 break;
@@ -326,6 +373,10 @@ impl<'a> TreeSink for TtTreeSink<'a> {
                 Some(tt::TokenTree::Leaf(leaf)) => {
                     self.cursor = self.cursor.bump();
                     self.buf += &format!("{}", leaf);
+
+                    if let tt::Leaf::Ident(ident) = leaf {
+                        last_ident = Some(ident);
+                    }
                 }
                 Some(tt::TokenTree::Subtree(subtree)) => {
                     self.cursor = self.cursor.subtree().unwrap();
@@ -345,6 +396,14 @@ impl<'a> TreeSink for TtTreeSink<'a> {
         self.buf.clear();
         self.inner.token(kind, text);
 
+        // Mark the range if needed
+        if let Some((range, token_id)) = self.range_marker.as_mut() {
+            if let Some(ident) = last_ident {
+                *range = TextRange::offset_len(range.start(), TextUnit::of_str(&ident.text));
+                *token_id = ident.id;
+            }
+        }
+
         // Add whitespace between adjoint puncts
         let next = self.cursor.bump();
         if let (
@@ -354,6 +413,7 @@ impl<'a> TreeSink for TtTreeSink<'a> {
         {
             if curr.spacing == tt::Spacing::Alone {
                 self.inner.token(WHITESPACE, " ".into());
+                self.text_pos += TextUnit::of_char(' ');
             }
         }
     }
@@ -361,6 +421,15 @@ impl<'a> TreeSink for TtTreeSink<'a> {
     fn start_node(&mut self, kind: SyntaxKind) {
         self.inner.start_node(kind);
 
+        self.range_marker = if kind == IDENT {
+            Some((
+                TextRange::offset_len(self.text_pos, TextUnit::from_usize(0)),
+                tt::TokenId::unspecified(),
+            ))
+        } else {
+            None
+        };
+
         match self.roots.last_mut() {
             None | Some(0) => self.roots.push(1),
             Some(ref mut n) => **n += 1,
@@ -370,6 +439,12 @@ impl<'a> TreeSink for TtTreeSink<'a> {
     fn finish_node(&mut self) {
         self.inner.finish_node();
         *self.roots.last_mut().unwrap() -= 1;
+
+        if let Some(range) = self.range_marker {
+            if range.1 != tt::TokenId::unspecified() {
+                self.range_map.set(range.0, &range.1)
+            }
+        }
     }
 
     fn error(&mut self, error: ParseError) {
diff --git a/crates/ra_mbe/src/tests.rs b/crates/ra_mbe/src/tests.rs
index a23e3afe3..a848ea334 100644
--- a/crates/ra_mbe/src/tests.rs
+++ b/crates/ra_mbe/src/tests.rs
@@ -126,7 +126,7 @@ fn test_expr_order() {
 "#,
     );
     let expanded = expand(&rules, "foo! { 1 + 1}");
-    let tree = token_tree_to_items(&expanded).unwrap().tree();
+    let tree = token_tree_to_items(&expanded).unwrap().0.tree();
 
     let dump = format!("{:#?}", tree.syntax());
     assert_eq_text!(
@@ -383,7 +383,7 @@ fn test_expand_to_item_list() {
             ",
     );
     let expansion = expand(&rules, "structs!(Foo, Bar);");
-    let tree = token_tree_to_items(&expansion).unwrap().tree();
+    let tree = token_tree_to_items(&expansion).unwrap().0.tree();
     assert_eq!(
         format!("{:#?}", tree.syntax()).trim(),
         r#"
@@ -501,7 +501,7 @@ fn test_tt_to_stmts() {
     );
 
     let expanded = expand(&rules, "foo!{}");
-    let stmts = token_tree_to_macro_stmts(&expanded).unwrap().tree();
+    let stmts = token_tree_to_macro_stmts(&expanded).unwrap().0.tree();
 
     assert_eq!(
         format!("{:#?}", stmts.syntax()).trim(),
@@ -946,7 +946,7 @@ fn test_vec() {
     );
 
     let expansion = expand(&rules, r#"vec![1u32,2];"#);
-    let tree = token_tree_to_expr(&expansion).unwrap().tree();
+    let tree = token_tree_to_expr(&expansion).unwrap().0.tree();
 
     assert_eq!(
         format!("{:#?}", tree.syntax()).trim(),
@@ -1436,8 +1436,8 @@ pub(crate) fn assert_expansion(
     };
     let (expanded_tree, expected_tree) = match kind {
         MacroKind::Items => {
-            let expanded_tree = token_tree_to_items(&expanded).unwrap().tree();
-            let expected_tree = token_tree_to_items(&expected).unwrap().tree();
+            let expanded_tree = token_tree_to_items(&expanded).unwrap().0.tree();
+            let expected_tree = token_tree_to_items(&expected).unwrap().0.tree();
 
             (
                 debug_dump_ignore_spaces(expanded_tree.syntax()).trim().to_string(),
@@ -1446,8 +1446,8 @@ pub(crate) fn assert_expansion(
         }
 
         MacroKind::Stmts => {
-            let expanded_tree = token_tree_to_macro_stmts(&expanded).unwrap().tree();
-            let expected_tree = token_tree_to_macro_stmts(&expected).unwrap().tree();
+            let expanded_tree = token_tree_to_macro_stmts(&expanded).unwrap().0.tree();
+            let expected_tree = token_tree_to_macro_stmts(&expected).unwrap().0.tree();
 
             (
                 debug_dump_ignore_spaces(expanded_tree.syntax()).trim().to_string(),