Merge #2173

2173: MBE: Add TokenId shift in macro_rules r=matklad a=edwin0cheng

As discussed in #2169 , for fixing duplication TokenId during expansion :

> What we can do here is to re-number the tokens during expansion. Specifically:
> * when we create macro_rules, we note the highest id of the token we have as shift>
> * when we expand macro rules, if we need to output a token from definition, we just re-use its id
> * if we need to output a token from the argument, we increase its id by shift (so it's guaranteed to not to collide with anything from the definition)
> * finally, when we have a HirFileId of the expansion, we can look up the original value of shift and classify node to the arg/def by comparing it's id with shift.
> 

This PR implement first 3 points of above solution. 

Co-authored-by: Edwin Cheng <[email protected]>

2 files changed, 70 insertions, 2 deletions

diff --git a/crates/ra_mbe/src/lib.rs b/crates/ra_mbe/src/lib.rs
index afdbee84e..15f000175 100644
--- a/crates/ra_mbe/src/lib.rs
+++ b/crates/ra_mbe/src/lib.rs
@@ -42,6 +42,8 @@ pub use crate::syntax_bridge::{
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct MacroRules {
     pub(crate) rules: Vec<Rule>,
+    /// Highest id of the token we have in TokenMap
+    pub(crate) shift: u32,
 }
 
 #[derive(Clone, Debug, PartialEq, Eq)]
@@ -50,6 +52,38 @@ pub(crate) struct Rule {
     pub(crate) rhs: tt::Subtree,
 }
 
+// Find the max token id inside a subtree
+fn max_id(subtree: &tt::Subtree) -> Option<u32> {
+    subtree
+        .token_trees
+        .iter()
+        .filter_map(|tt| match tt {
+            tt::TokenTree::Subtree(subtree) => max_id(subtree),
+            tt::TokenTree::Leaf(tt::Leaf::Ident(ident))
+                if ident.id != tt::TokenId::unspecified() =>
+            {
+                Some(ident.id.0)
+            }
+            _ => None,
+        })
+        .max()
+}
+
+/// Shift given TokenTree token id
+fn shift_subtree(tt: &mut tt::Subtree, shift: u32) {
+    for t in tt.token_trees.iter_mut() {
+        match t {
+            tt::TokenTree::Leaf(leaf) => match leaf {
+                tt::Leaf::Ident(ident) if ident.id != tt::TokenId::unspecified() => {
+                    ident.id.0 += shift;
+                }
+                _ => (),
+            },
+            tt::TokenTree::Subtree(tt) => shift_subtree(tt, shift),
+        }
+    }
+}
+
 impl MacroRules {
     pub fn parse(tt: &tt::Subtree) -> Result<MacroRules, ParseError> {
         // Note: this parsing can be implemented using mbe machinery itself, by
@@ -72,10 +106,17 @@ impl MacroRules {
             validate(&rule.lhs)?;
         }
 
-        Ok(MacroRules { rules })
+        // Note that TokenId is started from zero,
+        // We have to add 1 to prevent duplication.
+        let shift = max_id(tt).map_or(0, |it| it + 1);
+        Ok(MacroRules { rules, shift })
     }
+
     pub fn expand(&self, tt: &tt::Subtree) -> Result<tt::Subtree, ExpandError> {
-        mbe_expander::expand(self, tt)
+        // apply shift
+        let mut tt = tt.clone();
+        shift_subtree(&mut tt, self.shift);
+        mbe_expander::expand(self, &tt)
     }
 }
 
diff --git a/crates/ra_mbe/src/tests.rs b/crates/ra_mbe/src/tests.rs
index f34ab52e1..a23e3afe3 100644
--- a/crates/ra_mbe/src/tests.rs
+++ b/crates/ra_mbe/src/tests.rs
@@ -59,6 +59,33 @@ mod rule_parsing {
 // * Make it pass :-)
 
 #[test]
+fn test_token_id_shift() {
+    let macro_definition = r#"
+macro_rules! foobar {
+    ($e:ident) => { foo bar $e }
+}
+"#;
+    let rules = create_rules(macro_definition);
+    let expansion = expand(&rules, "foobar!(baz);");
+
+    fn get_id(t: &tt::TokenTree) -> Option<u32> {
+        if let tt::TokenTree::Leaf(tt::Leaf::Ident(ident)) = t {
+            return Some(ident.id.0);
+        }
+        None
+    }
+
+    assert_eq!(expansion.token_trees.len(), 3);
+    // ($e:ident) => { foo bar $e }
+    //   0 1            2   3   4
+    assert_eq!(get_id(&expansion.token_trees[0]), Some(2));
+    assert_eq!(get_id(&expansion.token_trees[1]), Some(3));
+
+    // So baz should be 5
+    assert_eq!(get_id(&expansion.token_trees[2]), Some(5));
+}
+
+#[test]
 fn test_convert_tt() {
     let macro_definition = r#"
 macro_rules! impl_froms {