3 files changed, 222 insertions, 212 deletions

diff --git a/crates/hir_expand/src/db.rs b/crates/hir_expand/src/db.rs
index 1e4b0cc19..3e9abd8a1 100644
--- a/crates/hir_expand/src/db.rs
+++ b/crates/hir_expand/src/db.rs
@@ -3,14 +3,14 @@
 use std::sync::Arc;
 
 use base_db::{salsa, SourceDatabase};
-use mbe::{ExpandError, ExpandResult, MacroDef, MacroRules};
+use mbe::{ExpandError, ExpandResult};
 use parser::FragmentKind;
 use syntax::{
     algo::diff,
-    ast::{MacroStmts, NameOwner},
+    ast::{self, NameOwner},
     AstNode, GreenNode, Parse,
     SyntaxKind::*,
-    SyntaxNode,
+    SyntaxNode, SyntaxToken,
 };
 
 use crate::{
@@ -27,23 +27,28 @@ const TOKEN_LIMIT: usize = 524288;
 
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub enum TokenExpander {
-    MacroRules(mbe::MacroRules),
-    MacroDef(mbe::MacroDef),
+    /// Old-style `macro_rules`.
+    MacroRules { mac: mbe::MacroRules, def_site_token_map: mbe::TokenMap },
+    /// AKA macros 2.0.
+    MacroDef { mac: mbe::MacroDef, def_site_token_map: mbe::TokenMap },
+    /// Stuff like `line!` and `file!`.
     Builtin(BuiltinFnLikeExpander),
+    /// `derive(Copy)` and such.
     BuiltinDerive(BuiltinDeriveExpander),
+    /// The thing we love the most here in rust-analyzer -- procedural macros.
     ProcMacro(ProcMacroExpander),
 }
 
 impl TokenExpander {
-    pub fn expand(
+    fn expand(
         &self,
         db: &dyn AstDatabase,
         id: LazyMacroId,
         tt: &tt::Subtree,
     ) -> mbe::ExpandResult<tt::Subtree> {
         match self {
-            TokenExpander::MacroRules(it) => it.expand(tt),
-            TokenExpander::MacroDef(it) => it.expand(tt),
+            TokenExpander::MacroRules { mac, .. } => mac.expand(tt),
+            TokenExpander::MacroDef { mac, .. } => mac.expand(tt),
             TokenExpander::Builtin(it) => it.expand(db, id, tt),
             // FIXME switch these to ExpandResult as well
             TokenExpander::BuiltinDerive(it) => it.expand(db, id, tt).into(),
@@ -56,23 +61,23 @@ impl TokenExpander {
         }
     }
 
-    pub fn map_id_down(&self, id: tt::TokenId) -> tt::TokenId {
+    pub(crate) fn map_id_down(&self, id: tt::TokenId) -> tt::TokenId {
         match self {
-            TokenExpander::MacroRules(it) => it.map_id_down(id),
-            TokenExpander::MacroDef(it) => it.map_id_down(id),
-            TokenExpander::Builtin(..) => id,
-            TokenExpander::BuiltinDerive(..) => id,
-            TokenExpander::ProcMacro(..) => id,
+            TokenExpander::MacroRules { mac, .. } => mac.map_id_down(id),
+            TokenExpander::MacroDef { mac, .. } => mac.map_id_down(id),
+            TokenExpander::Builtin(..)
+            | TokenExpander::BuiltinDerive(..)
+            | TokenExpander::ProcMacro(..) => id,
         }
     }
 
-    pub fn map_id_up(&self, id: tt::TokenId) -> (tt::TokenId, mbe::Origin) {
+    pub(crate) fn map_id_up(&self, id: tt::TokenId) -> (tt::TokenId, mbe::Origin) {
         match self {
-            TokenExpander::MacroRules(it) => it.map_id_up(id),
-            TokenExpander::MacroDef(it) => it.map_id_up(id),
-            TokenExpander::Builtin(..) => (id, mbe::Origin::Call),
-            TokenExpander::BuiltinDerive(..) => (id, mbe::Origin::Call),
-            TokenExpander::ProcMacro(..) => (id, mbe::Origin::Call),
+            TokenExpander::MacroRules { mac, .. } => mac.map_id_up(id),
+            TokenExpander::MacroDef { mac, .. } => mac.map_id_up(id),
+            TokenExpander::Builtin(..)
+            | TokenExpander::BuiltinDerive(..)
+            | TokenExpander::ProcMacro(..) => (id, mbe::Origin::Call),
         }
     }
 }
@@ -82,28 +87,48 @@ impl TokenExpander {
 pub trait AstDatabase: SourceDatabase {
     fn ast_id_map(&self, file_id: HirFileId) -> Arc<AstIdMap>;
 
+    /// Main public API -- parsis a hir file, not caring whether it's a real
+    /// file or a macro expansion.
     #[salsa::transparent]
     fn parse_or_expand(&self, file_id: HirFileId) -> Option<SyntaxNode>;
-
-    #[salsa::interned]
-    fn intern_macro(&self, macro_call: MacroCallLoc) -> LazyMacroId;
-    fn macro_arg_text(&self, id: MacroCallId) -> Option<GreenNode>;
-    #[salsa::transparent]
-    fn macro_arg(&self, id: MacroCallId) -> Option<Arc<(tt::Subtree, mbe::TokenMap)>>;
-    fn macro_def(&self, id: MacroDefId) -> Option<Arc<(TokenExpander, mbe::TokenMap)>>;
+    /// Implementation for the macro case.
     fn parse_macro_expansion(
         &self,
         macro_file: MacroFile,
     ) -> ExpandResult<Option<(Parse<SyntaxNode>, Arc<mbe::TokenMap>)>>;
-    fn macro_expand(&self, macro_call: MacroCallId) -> ExpandResult<Option<Arc<tt::Subtree>>>;
-
-    /// Firewall query that returns the error from the `macro_expand` query.
-    fn macro_expand_error(&self, macro_call: MacroCallId) -> Option<ExpandError>;
 
+    /// Macro ids. That's probably the tricksiest bit in rust-analyzer, and the
+    /// reason why we use salsa at all.
+    ///
+    /// We encode macro definitions into ids of macro calls, this what allows us
+    /// to be incremental.
+    #[salsa::interned]
+    fn intern_macro(&self, macro_call: MacroCallLoc) -> LazyMacroId;
+    /// Certain built-in macros are eager (`format!(concat!("file: ", file!(), "{}"")), 92`).
+    /// For them, we actually want to encode the whole token tree as an argument.
     #[salsa::interned]
     fn intern_eager_expansion(&self, eager: EagerCallLoc) -> EagerMacroId;
 
+    /// Lowers syntactic macro call to a token tree representation.
+    #[salsa::transparent]
+    fn macro_arg(&self, id: MacroCallId) -> Option<Arc<(tt::Subtree, mbe::TokenMap)>>;
+    /// Extracts syntax node, corresponding to a macro call. That's a firewall
+    /// query, only typing in the macro call itself changes the returned
+    /// subtree.
+    fn macro_arg_text(&self, id: MacroCallId) -> Option<GreenNode>;
+    /// Gets the expander for this macro. This compiles declarative macros, and
+    /// just fetches procedural ones.
+    fn macro_def(&self, id: MacroDefId) -> Option<Arc<TokenExpander>>;
+
+    /// Expand macro call to a token tree. This query is LRUed (we keep 128 or so results in memory)
+    fn macro_expand(&self, macro_call: MacroCallId) -> ExpandResult<Option<Arc<tt::Subtree>>>;
+    /// Special case of the previous query for procedural macros. We can't LRU
+    /// proc macros, since they are not deterministic in general, and
+    /// non-determinism breaks salsa in a very, very, very bad way. @edwin0cheng
+    /// heroically debugged this once!
     fn expand_proc_macro(&self, call: MacroCallId) -> Result<tt::Subtree, mbe::ExpandError>;
+    /// Firewall query that returns the error from the `macro_expand` query.
+    fn macro_expand_error(&self, macro_call: MacroCallId) -> Option<ExpandError>;
 
     fn hygiene_frame(&self, file_id: HirFileId) -> Arc<HygieneFrame>;
 }
@@ -115,36 +140,159 @@ pub trait AstDatabase: SourceDatabase {
 pub fn expand_hypothetical(
     db: &dyn AstDatabase,
     actual_macro_call: MacroCallId,
-    hypothetical_args: &syntax::ast::TokenTree,
-    token_to_map: syntax::SyntaxToken,
-) -> Option<(SyntaxNode, syntax::SyntaxToken)> {
-    let macro_file = MacroFile { macro_call_id: actual_macro_call };
+    hypothetical_args: &ast::TokenTree,
+    token_to_map: SyntaxToken,
+) -> Option<(SyntaxNode, SyntaxToken)> {
     let (tt, tmap_1) = mbe::syntax_node_to_token_tree(hypothetical_args.syntax());
     let range =
         token_to_map.text_range().checked_sub(hypothetical_args.syntax().text_range().start())?;
     let token_id = tmap_1.token_by_range(range)?;
-    let macro_def = expander(db, actual_macro_call)?;
+
+    let lazy_id = match actual_macro_call {
+        MacroCallId::LazyMacro(id) => id,
+        MacroCallId::EagerMacro(_) => return None,
+    };
+
+    let macro_def = {
+        let loc = db.lookup_intern_macro(lazy_id);
+        db.macro_def(loc.def)?
+    };
+
+    let hypothetical_expansion = macro_def.expand(db, lazy_id, &tt);
+
+    let fragment_kind = to_fragment_kind(db, actual_macro_call);
+
     let (node, tmap_2) =
-        parse_macro_with_arg(db, macro_file, Some(std::sync::Arc::new((tt, tmap_1)))).value?;
-    let token_id = macro_def.0.map_id_down(token_id);
+        mbe::token_tree_to_syntax_node(&hypothetical_expansion.value, fragment_kind).ok()?;
+
+    let token_id = macro_def.map_id_down(token_id);
     let range = tmap_2.range_by_token(token_id)?.by_kind(token_to_map.kind())?;
     let token = node.syntax_node().covering_element(range).into_token()?;
     Some((node.syntax_node(), token))
 }
 
 fn ast_id_map(db: &dyn AstDatabase, file_id: HirFileId) -> Arc<AstIdMap> {
-    let map =
-        db.parse_or_expand(file_id).map_or_else(AstIdMap::default, |it| AstIdMap::from_source(&it));
+    let map = db.parse_or_expand(file_id).map(|it| AstIdMap::from_source(&it)).unwrap_or_default();
     Arc::new(map)
 }
 
-fn macro_def(db: &dyn AstDatabase, id: MacroDefId) -> Option<Arc<(TokenExpander, mbe::TokenMap)>> {
+fn parse_or_expand(db: &dyn AstDatabase, file_id: HirFileId) -> Option<SyntaxNode> {
+    match file_id.0 {
+        HirFileIdRepr::FileId(file_id) => Some(db.parse(file_id).tree().syntax().clone()),
+        HirFileIdRepr::MacroFile(macro_file) => {
+            db.parse_macro_expansion(macro_file).value.map(|(it, _)| it.syntax_node())
+        }
+    }
+}
+
+fn parse_macro_expansion(
+    db: &dyn AstDatabase,
+    macro_file: MacroFile,
+) -> ExpandResult<Option<(Parse<SyntaxNode>, Arc<mbe::TokenMap>)>> {
+    let _p = profile::span("parse_macro_expansion");
+    let result = db.macro_expand(macro_file.macro_call_id);
+
+    if let Some(err) = &result.err {
+        // Note:
+        // The final goal we would like to make all parse_macro success,
+        // such that the following log will not call anyway.
+        match macro_file.macro_call_id {
+            MacroCallId::LazyMacro(id) => {
+                let loc: MacroCallLoc = db.lookup_intern_macro(id);
+                let node = loc.kind.node(db);
+
+                // collect parent information for warning log
+                let parents = std::iter::successors(loc.kind.file_id().call_node(db), |it| {
+                    it.file_id.call_node(db)
+                })
+                .map(|n| format!("{:#}", n.value))
+                .collect::<Vec<_>>()
+                .join("\n");
+
+                log::warn!(
+                    "fail on macro_parse: (reason: {:?} macro_call: {:#}) parents: {}",
+                    err,
+                    node.value,
+                    parents
+                );
+            }
+            _ => {
+                log::warn!("fail on macro_parse: (reason: {:?})", err);
+            }
+        }
+    }
+    let tt = match result.value {
+        Some(tt) => tt,
+        None => return ExpandResult { value: None, err: result.err },
+    };
+
+    let fragment_kind = to_fragment_kind(db, macro_file.macro_call_id);
+
+    log::debug!("expanded = {}", tt.as_debug_string());
+    log::debug!("kind = {:?}", fragment_kind);
+
+    let (parse, rev_token_map) = match mbe::token_tree_to_syntax_node(&tt, fragment_kind) {
+        Ok(it) => it,
+        Err(err) => {
+            log::debug!(
+                "failed to parse expanstion to {:?} = {}",
+                fragment_kind,
+                tt.as_debug_string()
+            );
+            return ExpandResult::only_err(err);
+        }
+    };
+
+    match result.err {
+        Some(err) => {
+            // Safety check for recursive identity macro.
+            let node = parse.syntax_node();
+            let file: HirFileId = macro_file.into();
+            let call_node = match file.call_node(db) {
+                Some(it) => it,
+                None => {
+                    return ExpandResult::only_err(err);
+                }
+            };
+            if is_self_replicating(&node, &call_node.value) {
+                return ExpandResult::only_err(err);
+            } else {
+                ExpandResult { value: Some((parse, Arc::new(rev_token_map))), err: Some(err) }
+            }
+        }
+        None => {
+            log::debug!("parse = {:?}", parse.syntax_node().kind());
+            ExpandResult { value: Some((parse, Arc::new(rev_token_map))), err: None }
+        }
+    }
+}
+
+fn macro_arg(db: &dyn AstDatabase, id: MacroCallId) -> Option<Arc<(tt::Subtree, mbe::TokenMap)>> {
+    let arg = db.macro_arg_text(id)?;
+    let (tt, tmap) = mbe::syntax_node_to_token_tree(&SyntaxNode::new_root(arg));
+    Some(Arc::new((tt, tmap)))
+}
+
+fn macro_arg_text(db: &dyn AstDatabase, id: MacroCallId) -> Option<GreenNode> {
+    let id = match id {
+        MacroCallId::LazyMacro(id) => id,
+        MacroCallId::EagerMacro(_id) => {
+            // FIXME: support macro_arg for eager macro
+            return None;
+        }
+    };
+    let loc = db.lookup_intern_macro(id);
+    let arg = loc.kind.arg(db)?;
+    Some(arg.green())
+}
+
+fn macro_def(db: &dyn AstDatabase, id: MacroDefId) -> Option<Arc<TokenExpander>> {
     match id.kind {
         MacroDefKind::Declarative(ast_id) => match ast_id.to_node(db) {
-            syntax::ast::Macro::MacroRules(macro_rules) => {
+            ast::Macro::MacroRules(macro_rules) => {
                 let arg = macro_rules.token_tree()?;
-                let (tt, tmap) = mbe::ast_to_token_tree(&arg);
-                let rules = match MacroRules::parse(&tt) {
+                let (tt, def_site_token_map) = mbe::ast_to_token_tree(&arg);
+                let mac = match mbe::MacroRules::parse(&tt) {
                     Ok(it) => it,
                     Err(err) => {
                         let name = macro_rules.name().map(|n| n.to_string()).unwrap_or_default();
@@ -152,12 +300,12 @@ fn macro_def(db: &dyn AstDatabase, id: MacroDefId) -> Option<Arc<(TokenExpander,
                         return None;
                     }
                 };
-                Some(Arc::new((TokenExpander::MacroRules(rules), tmap)))
+                Some(Arc::new(TokenExpander::MacroRules { mac, def_site_token_map }))
             }
-            syntax::ast::Macro::MacroDef(macro_def) => {
+            ast::Macro::MacroDef(macro_def) => {
                 let arg = macro_def.body()?;
-                let (tt, tmap) = mbe::ast_to_token_tree(&arg);
-                let rules = match MacroDef::parse(&tt) {
+                let (tt, def_site_token_map) = mbe::ast_to_token_tree(&arg);
+                let mac = match mbe::MacroDef::parse(&tt) {
                     Ok(it) => it,
                     Err(err) => {
                         let name = macro_def.name().map(|n| n.to_string()).unwrap_or_default();
@@ -165,41 +313,18 @@ fn macro_def(db: &dyn AstDatabase, id: MacroDefId) -> Option<Arc<(TokenExpander,
                         return None;
                     }
                 };
-                Some(Arc::new((TokenExpander::MacroDef(rules), tmap)))
+                Some(Arc::new(TokenExpander::MacroDef { mac, def_site_token_map }))
             }
         },
-        MacroDefKind::BuiltIn(expander, _) => {
-            Some(Arc::new((TokenExpander::Builtin(expander), mbe::TokenMap::default())))
-        }
+        MacroDefKind::BuiltIn(expander, _) => Some(Arc::new(TokenExpander::Builtin(expander))),
         MacroDefKind::BuiltInDerive(expander, _) => {
-            Some(Arc::new((TokenExpander::BuiltinDerive(expander), mbe::TokenMap::default())))
+            Some(Arc::new(TokenExpander::BuiltinDerive(expander)))
         }
         MacroDefKind::BuiltInEager(..) => None,
-        MacroDefKind::ProcMacro(expander, ..) => {
-            Some(Arc::new((TokenExpander::ProcMacro(expander), mbe::TokenMap::default())))
-        }
+        MacroDefKind::ProcMacro(expander, ..) => Some(Arc::new(TokenExpander::ProcMacro(expander))),
     }
 }
 
-fn macro_arg_text(db: &dyn AstDatabase, id: MacroCallId) -> Option<GreenNode> {
-    let id = match id {
-        MacroCallId::LazyMacro(id) => id,
-        MacroCallId::EagerMacro(_id) => {
-            // FIXME: support macro_arg for eager macro
-            return None;
-        }
-    };
-    let loc = db.lookup_intern_macro(id);
-    let arg = loc.kind.arg(db)?;
-    Some(arg.green())
-}
-
-fn macro_arg(db: &dyn AstDatabase, id: MacroCallId) -> Option<Arc<(tt::Subtree, mbe::TokenMap)>> {
-    let arg = db.macro_arg_text(id)?;
-    let (tt, tmap) = mbe::syntax_node_to_token_tree(&SyntaxNode::new_root(arg));
-    Some(Arc::new((tt, tmap)))
-}
-
 fn macro_expand(db: &dyn AstDatabase, id: MacroCallId) -> ExpandResult<Option<Arc<tt::Subtree>>> {
     macro_expand_with_arg(db, id, None)
 }
@@ -208,19 +333,6 @@ fn macro_expand_error(db: &dyn AstDatabase, macro_call: MacroCallId) -> Option<E
     db.macro_expand(macro_call).err
 }
 
-fn expander(db: &dyn AstDatabase, id: MacroCallId) -> Option<Arc<(TokenExpander, mbe::TokenMap)>> {
-    let lazy_id = match id {
-        MacroCallId::LazyMacro(id) => id,
-        MacroCallId::EagerMacro(_id) => {
-            return None;
-        }
-    };
-
-    let loc = db.lookup_intern_macro(lazy_id);
-    let macro_rules = db.macro_def(loc.def)?;
-    Some(macro_rules)
-}
-
 fn macro_expand_with_arg(
     db: &dyn AstDatabase,
     id: MacroCallId,
@@ -254,7 +366,7 @@ fn macro_expand_with_arg(
         Some(it) => it,
         None => return ExpandResult::str_err("Fail to find macro definition".into()),
     };
-    let ExpandResult { value: tt, err } = macro_rules.0.expand(db, lazy_id, &macro_arg.0);
+    let ExpandResult { value: tt, err } = macro_rules.expand(db, lazy_id, &macro_arg.0);
     // Set a hard limit for the expanded tt
     let count = tt.count();
     if count > TOKEN_LIMIT {
@@ -294,116 +406,11 @@ fn expand_proc_macro(
     expander.expand(db, loc.krate, &macro_arg.0)
 }
 
-fn parse_or_expand(db: &dyn AstDatabase, file_id: HirFileId) -> Option<SyntaxNode> {
-    match file_id.0 {
-        HirFileIdRepr::FileId(file_id) => Some(db.parse(file_id).tree().syntax().clone()),
-        HirFileIdRepr::MacroFile(macro_file) => {
-            db.parse_macro_expansion(macro_file).value.map(|(it, _)| it.syntax_node())
-        }
-    }
-}
-
-fn parse_macro_expansion(
-    db: &dyn AstDatabase,
-    macro_file: MacroFile,
-) -> ExpandResult<Option<(Parse<SyntaxNode>, Arc<mbe::TokenMap>)>> {
-    parse_macro_with_arg(db, macro_file, None)
-}
-
-fn parse_macro_with_arg(
-    db: &dyn AstDatabase,
-    macro_file: MacroFile,
-    arg: Option<Arc<(tt::Subtree, mbe::TokenMap)>>,
-) -> ExpandResult<Option<(Parse<SyntaxNode>, Arc<mbe::TokenMap>)>> {
-    let macro_call_id = macro_file.macro_call_id;
-    let result = if let Some(arg) = arg {
-        macro_expand_with_arg(db, macro_call_id, Some(arg))
-    } else {
-        db.macro_expand(macro_call_id)
-    };
-
-    let _p = profile::span("parse_macro_expansion");
-
-    if let Some(err) = &result.err {
-        // Note:
-        // The final goal we would like to make all parse_macro success,
-        // such that the following log will not call anyway.
-        match macro_call_id {
-            MacroCallId::LazyMacro(id) => {
-                let loc: MacroCallLoc = db.lookup_intern_macro(id);
-                let node = loc.kind.node(db);
-
-                // collect parent information for warning log
-                let parents = std::iter::successors(loc.kind.file_id().call_node(db), |it| {
-                    it.file_id.call_node(db)
-                })
-                .map(|n| format!("{:#}", n.value))
-                .collect::<Vec<_>>()
-                .join("\n");
-
-                log::warn!(
-                    "fail on macro_parse: (reason: {:?} macro_call: {:#}) parents: {}",
-                    err,
-                    node.value,
-                    parents
-                );
-            }
-            _ => {
-                log::warn!("fail on macro_parse: (reason: {:?})", err);
-            }
-        }
-    }
-    let tt = match result.value {
-        Some(tt) => tt,
-        None => return ExpandResult { value: None, err: result.err },
-    };
-
-    let fragment_kind = to_fragment_kind(db, macro_call_id);
-
-    log::debug!("expanded = {}", tt.as_debug_string());
-    log::debug!("kind = {:?}", fragment_kind);
-
-    let (parse, rev_token_map) = match mbe::token_tree_to_syntax_node(&tt, fragment_kind) {
-        Ok(it) => it,
-        Err(err) => {
-            log::debug!(
-                "failed to parse expanstion to {:?} = {}",
-                fragment_kind,
-                tt.as_debug_string()
-            );
-            return ExpandResult::only_err(err);
-        }
-    };
-
-    match result.err {
-        Some(err) => {
-            // Safety check for recursive identity macro.
-            let node = parse.syntax_node();
-            let file: HirFileId = macro_file.into();
-            let call_node = match file.call_node(db) {
-                Some(it) => it,
-                None => {
-                    return ExpandResult::only_err(err);
-                }
-            };
-            if is_self_replicating(&node, &call_node.value) {
-                return ExpandResult::only_err(err);
-            } else {
-                ExpandResult { value: Some((parse, Arc::new(rev_token_map))), err: Some(err) }
-            }
-        }
-        None => {
-            log::debug!("parse = {:?}", parse.syntax_node().kind());
-            ExpandResult { value: Some((parse, Arc::new(rev_token_map))), err: None }
-        }
-    }
-}
-
 fn is_self_replicating(from: &SyntaxNode, to: &SyntaxNode) -> bool {
     if diff(from, to).is_empty() {
         return true;
     }
-    if let Some(stmts) = MacroStmts::cast(from.clone()) {
+    if let Some(stmts) = ast::MacroStmts::cast(from.clone()) {
         if stmts.statements().any(|stmt| diff(stmt.syntax(), to).is_empty()) {
             return true;
         }
diff --git a/crates/hir_expand/src/hygiene.rs b/crates/hir_expand/src/hygiene.rs
index 779725629..ed61ebca3 100644
--- a/crates/hir_expand/src/hygiene.rs
+++ b/crates/hir_expand/src/hygiene.rs
@@ -5,6 +5,7 @@
 use std::sync::Arc;
 
 use base_db::CrateId;
+use db::TokenExpander;
 use either::Either;
 use mbe::Origin;
 use parser::SyntaxKind;
@@ -115,7 +116,7 @@ struct HygieneInfo {
     /// The `macro_rules!` arguments.
     def_start: Option<InFile<TextSize>>,
 
-    macro_def: Arc<(db::TokenExpander, mbe::TokenMap)>,
+    macro_def: Arc<TokenExpander>,
     macro_arg: Arc<(tt::Subtree, mbe::TokenMap)>,
     exp_map: Arc<mbe::TokenMap>,
 }
@@ -124,13 +125,16 @@ impl HygieneInfo {
     fn map_ident_up(&self, token: TextRange) -> Option<(InFile<TextRange>, Origin)> {
         let token_id = self.exp_map.token_by_range(token)?;
 
-        let (token_id, origin) = self.macro_def.0.map_id_up(token_id);
+        let (token_id, origin) = self.macro_def.map_id_up(token_id);
         let (token_map, tt) = match origin {
             mbe::Origin::Call => (&self.macro_arg.1, self.arg_start),
-            mbe::Origin::Def => (
-                &self.macro_def.1,
-                *self.def_start.as_ref().expect("`Origin::Def` used with non-`macro_rules!` macro"),
-            ),
+            mbe::Origin::Def => match (&*self.macro_def, self.def_start) {
+                (TokenExpander::MacroDef { def_site_token_map, .. }, Some(tt))
+                | (TokenExpander::MacroRules { def_site_token_map, .. }, Some(tt)) => {
+                    (def_site_token_map, tt)
+                }
+                _ => panic!("`Origin::Def` used with non-`macro_rules!` macro"),
+            },
         };
 
         let range = token_map.range_by_token(token_id)?.by_kind(SyntaxKind::IDENT)?;
diff --git a/crates/hir_expand/src/lib.rs b/crates/hir_expand/src/lib.rs
index a0e6aec62..0402640de 100644
--- a/crates/hir_expand/src/lib.rs
+++ b/crates/hir_expand/src/lib.rs
@@ -351,7 +351,7 @@ pub struct ExpansionInfo {
     /// The `macro_rules!` arguments.
     def: Option<InFile<ast::TokenTree>>,
 
-    macro_def: Arc<(db::TokenExpander, mbe::TokenMap)>,
+    macro_def: Arc<db::TokenExpander>,
     macro_arg: Arc<(tt::Subtree, mbe::TokenMap)>,
     exp_map: Arc<mbe::TokenMap>,
 }
@@ -368,7 +368,7 @@ impl ExpansionInfo {
         assert_eq!(token.file_id, self.arg.file_id);
         let range = token.value.text_range().checked_sub(self.arg.value.text_range().start())?;
         let token_id = self.macro_arg.1.token_by_range(range)?;
-        let token_id = self.macro_def.0.map_id_down(token_id);
+        let token_id = self.macro_def.map_id_down(token_id);
 
         let range = self.exp_map.range_by_token(token_id)?.by_kind(token.value.kind())?;
 
@@ -383,17 +383,16 @@ impl ExpansionInfo {
     ) -> Option<(InFile<SyntaxToken>, Origin)> {
         let token_id = self.exp_map.token_by_range(token.value.text_range())?;
 
-        let (token_id, origin) = self.macro_def.0.map_id_up(token_id);
+        let (token_id, origin) = self.macro_def.map_id_up(token_id);
         let (token_map, tt) = match origin {
             mbe::Origin::Call => (&self.macro_arg.1, self.arg.clone()),
-            mbe::Origin::Def => (
-                &self.macro_def.1,
-                self.def
-                    .as_ref()
-                    .expect("`Origin::Def` used with non-`macro_rules!` macro")
-                    .as_ref()
-                    .map(|tt| tt.syntax().clone()),
-            ),
+            mbe::Origin::Def => match (&*self.macro_def, self.def.as_ref()) {
+                (db::TokenExpander::MacroRules { def_site_token_map, .. }, Some(tt))
+                | (db::TokenExpander::MacroDef { def_site_token_map, .. }, Some(tt)) => {
+                    (def_site_token_map, tt.as_ref().map(|tt| tt.syntax().clone()))
+                }
+                _ => panic!("`Origin::Def` used with non-`macro_rules!` macro"),
+            },
         };
 
         let range = token_map.range_by_token(token_id)?.by_kind(token.value.kind())?;