3 files changed, 84 insertions, 47 deletions

diff --git a/crates/ra_hir_expand/src/db.rs b/crates/ra_hir_expand/src/db.rs
index 5eadee9c2..9de7c1ea8 100644
--- a/crates/ra_hir_expand/src/db.rs
+++ b/crates/ra_hir_expand/src/db.rs
@@ -32,10 +32,17 @@ impl TokenExpander {
         }
     }
 
-    pub fn shift(&self) -> u32 {
+    pub fn map_id_down(&self, id: tt::TokenId) -> tt::TokenId {
         match self {
-            TokenExpander::MacroRules(it) => it.shift(),
-            TokenExpander::Builtin(_) => 0,
+            TokenExpander::MacroRules(it) => it.map_id_down(id),
+            TokenExpander::Builtin(..) => id,
+        }
+    }
+
+    pub fn map_id_up(&self, id: tt::TokenId) -> (tt::TokenId, mbe::Origin) {
+        match self {
+            TokenExpander::MacroRules(it) => it.map_id_up(id),
+            TokenExpander::Builtin(..) => (id, mbe::Origin::Def),
         }
     }
 }
diff --git a/crates/ra_hir_expand/src/lib.rs b/crates/ra_hir_expand/src/lib.rs
index 1774cb58b..5927a03ba 100644
--- a/crates/ra_hir_expand/src/lib.rs
+++ b/crates/ra_hir_expand/src/lib.rs
@@ -161,7 +161,7 @@ impl ExpansionInfo {
     pub fn translate_offset(&self, offset: TextUnit) -> Option<TextUnit> {
         let offset = offset.checked_sub(self.arg_start.1)?;
         let token_id = self.macro_arg.1.token_by_offset(offset)?;
-        let token_id = tt::TokenId(token_id.0 + self.macro_def.0.shift());
+        let token_id = self.macro_def.0.map_id_down(token_id);
 
         let (r, _) = self.exp_map.ranges.iter().find(|(_, tid)| *tid == token_id)?;
         Some(r.start())
@@ -170,11 +170,11 @@ impl ExpansionInfo {
     pub fn find_range(&self, from: TextRange) -> Option<(HirFileId, TextRange)> {
         let token_id = look_in_rev_map(&self.exp_map, from)?;
 
-        let shift = self.macro_def.0.shift();
-        let (token_map, (file_id, start_offset), token_id) = if token_id.0 >= shift {
-            (&self.macro_arg.1, self.arg_start, tt::TokenId(token_id.0 - shift).into())
-        } else {
-            (&self.macro_def.1, self.def_start, token_id)
+        let (token_id, origin) = self.macro_def.0.map_id_up(token_id);
+
+        let (token_map, (file_id, start_offset)) = match origin {
+            mbe::Origin::Call => (&self.macro_arg.1, self.arg_start),
+            mbe::Origin::Def => (&self.macro_def.1, self.def_start),
         };
 
         let range = token_map.relative_range_of(token_id)?;
diff --git a/crates/ra_mbe/src/lib.rs b/crates/ra_mbe/src/lib.rs
index 8a31d1c36..58ca95368 100644
--- a/crates/ra_mbe/src/lib.rs
+++ b/crates/ra_mbe/src/lib.rs
@@ -40,49 +40,75 @@ pub use crate::syntax_bridge::{
 /// and `$()*` have special meaning (see `Var` and `Repeat` data structures)
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct MacroRules {
-    pub(crate) rules: Vec<Rule>,
+    rules: Vec<Rule>,
     /// Highest id of the token we have in TokenMap
-    pub(crate) shift: u32,
+    shift: Shift,
 }
 
 #[derive(Clone, Debug, PartialEq, Eq)]
-pub(crate) struct Rule {
-    pub(crate) lhs: tt::Subtree,
-    pub(crate) rhs: tt::Subtree,
+struct Rule {
+    lhs: tt::Subtree,
+    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)
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+struct Shift(u32);
+
+impl Shift {
+    fn new(tt: &tt::Subtree) -> Shift {
+        // Note that TokenId is started from zero,
+        // We have to add 1 to prevent duplication.
+        let value = max_id(tt).map_or(0, |it| it + 1);
+        return Shift(value);
+
+        // 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_all(self, tt: &mut tt::Subtree) {
+        for t in tt.token_trees.iter_mut() {
+            match t {
+                tt::TokenTree::Leaf(leaf) => match leaf {
+                    tt::Leaf::Ident(ident) => ident.id = self.shift(ident.id),
+                    _ => (),
+                },
+                tt::TokenTree::Subtree(tt) => self.shift_all(tt),
             }
-            _ => 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),
+    fn shift(self, id: tt::TokenId) -> tt::TokenId {
+        if id == tt::TokenId::unspecified() {
+            return id;
         }
+        tt::TokenId(id.0 + self.0)
+    }
+
+    fn unshift(self, id: tt::TokenId) -> Option<tt::TokenId> {
+        id.0.checked_sub(self.0).map(tt::TokenId)
     }
 }
 
+pub enum Origin {
+    Def,
+    Call,
+}
+
 impl MacroRules {
     pub fn parse(tt: &tt::Subtree) -> Result<MacroRules, ParseError> {
         // Note: this parsing can be implemented using mbe machinery itself, by
@@ -105,21 +131,25 @@ impl MacroRules {
             validate(&rule.lhs)?;
         }
 
-        // 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 })
+        Ok(MacroRules { rules, shift: Shift::new(tt) })
     }
 
     pub fn expand(&self, tt: &tt::Subtree) -> Result<tt::Subtree, ExpandError> {
         // apply shift
         let mut tt = tt.clone();
-        shift_subtree(&mut tt, self.shift);
+        self.shift.shift_all(&mut tt);
         mbe_expander::expand(self, &tt)
     }
 
-    pub fn shift(&self) -> u32 {
-        self.shift
+    pub fn map_id_down(&self, id: tt::TokenId) -> tt::TokenId {
+        self.shift.shift(id)
+    }
+
+    pub fn map_id_up(&self, id: tt::TokenId) -> (tt::TokenId, Origin) {
+        match self.shift.unshift(id) {
+            Some(id) => (id, Origin::Call),
+            None => (id, Origin::Def),
+        }
     }
 }