2 files changed, 15 insertions, 24 deletions

diff --git a/crates/hir-def/src/body.rs b/crates/hir-def/src/body.rs
index 7e9acd3f86e..080a307b1f8 100644
--- a/crates/hir-def/src/body.rs
+++ b/crates/hir-def/src/body.rs
@@ -5,21 +5,18 @@ mod lower;
 mod tests;
 pub mod scope;
 
-use std::{mem, ops::Index, sync::Arc};
+use std::{ops::Index, sync::Arc};
 
 use base_db::CrateId;
 use cfg::{CfgExpr, CfgOptions};
 use drop_bomb::DropBomb;
 use either::Either;
-use hir_expand::{
-    ast_id_map::AstIdMap, hygiene::Hygiene, AstId, ExpandError, ExpandResult, HirFileId, InFile,
-    MacroCallId,
-};
+use hir_expand::{hygiene::Hygiene, ExpandError, ExpandResult, HirFileId, InFile, MacroCallId};
 use la_arena::{Arena, ArenaMap};
 use limit::Limit;
 use profile::Count;
 use rustc_hash::FxHashMap;
-use syntax::{ast, AstNode, AstPtr, SyntaxNodePtr};
+use syntax::{ast, AstPtr, SyntaxNodePtr};
 
 use crate::{
     attr::{Attrs, RawAttrs},
@@ -50,7 +47,6 @@ pub struct Expander {
     cfg_expander: CfgExpander,
     def_map: Arc<DefMap>,
     current_file_id: HirFileId,
-    ast_id_map: Option<Arc<AstIdMap>>,
     module: LocalModuleId,
     recursion_limit: usize,
 }
@@ -84,7 +80,6 @@ impl Expander {
             cfg_expander,
             def_map,
             current_file_id,
-            ast_id_map: None,
             module: module.local_id,
             recursion_limit: 0,
         }
@@ -167,14 +162,10 @@ impl Expander {
         tracing::debug!("macro expansion {:#?}", node.syntax());
 
         self.recursion_limit += 1;
-        let mark = Mark {
-            file_id: self.current_file_id,
-            ast_id_map: mem::take(&mut self.ast_id_map),
-            bomb: DropBomb::new("expansion mark dropped"),
-        };
+        let mark =
+            Mark { file_id: self.current_file_id, bomb: DropBomb::new("expansion mark dropped") };
         self.cfg_expander.hygiene = Hygiene::new(db.upcast(), file_id);
         self.current_file_id = file_id;
-        self.ast_id_map = None;
 
         ExpandResult { value: Some((mark, node)), err }
     }
@@ -182,7 +173,6 @@ impl Expander {
     pub fn exit(&mut self, db: &dyn DefDatabase, mut mark: Mark) {
         self.cfg_expander.hygiene = Hygiene::new(db.upcast(), mark.file_id);
         self.current_file_id = mark.file_id;
-        self.ast_id_map = mem::take(&mut mark.ast_id_map);
         self.recursion_limit -= 1;
         mark.bomb.defuse();
     }
@@ -212,12 +202,6 @@ impl Expander {
         self.def_map.resolve_path(db, self.module, path, BuiltinShadowMode::Other).0.take_macros()
     }
 
-    fn ast_id<N: AstNode>(&mut self, db: &dyn DefDatabase, item: &N) -> AstId<N> {
-        let file_local_id =
-            self.ast_id_map.get_or_insert_with(|| db.ast_id_map(self.current_file_id)).ast_id(item);
-        AstId::new(self.current_file_id, file_local_id)
-    }
-
     fn recursion_limit(&self, db: &dyn DefDatabase) -> Limit {
         let limit = db.crate_limits(self.cfg_expander.krate).recursion_limit as _;
 
@@ -233,7 +217,6 @@ impl Expander {
 #[derive(Debug)]
 pub struct Mark {
     file_id: HirFileId,
-    ast_id_map: Option<Arc<AstIdMap>>,
     bomb: DropBomb,
 }
 
diff --git a/crates/hir-def/src/body/lower.rs b/crates/hir-def/src/body/lower.rs
index 006376e362c..8cef31ed151 100644
--- a/crates/hir-def/src/body/lower.rs
+++ b/crates/hir-def/src/body/lower.rs
@@ -8,7 +8,7 @@ use hir_expand::{
     ast_id_map::{AstIdMap, FileAstId},
     hygiene::Hygiene,
     name::{name, AsName, Name},
-    ExpandError, HirFileId, InFile,
+    AstId, ExpandError, HirFileId, InFile,
 };
 use la_arena::Arena;
 use once_cell::unsync::OnceCell;
@@ -90,6 +90,7 @@ pub(super) fn lower(
     ExprCollector {
         db,
         source_map: BodySourceMap::default(),
+        ast_id_map: db.ast_id_map(expander.current_file_id),
         body: Body {
             exprs: Arena::default(),
             pats: Arena::default(),
@@ -110,6 +111,7 @@ pub(super) fn lower(
 struct ExprCollector<'a> {
     db: &'a dyn DefDatabase,
     expander: Expander,
+    ast_id_map: Arc<AstIdMap>,
     body: Body,
     source_map: BodySourceMap,
     // a poor-mans union-find?
@@ -591,8 +593,13 @@ impl ExprCollector<'_> {
         match res.value {
             Some((mark, expansion)) => {
                 self.source_map.expansions.insert(macro_call_ptr, self.expander.current_file_id);
+                let prev_ast_id_map = mem::replace(
+                    &mut self.ast_id_map,
+                    self.db.ast_id_map(self.expander.current_file_id),
+                );
 
                 let id = collector(self, Some(expansion));
+                self.ast_id_map = prev_ast_id_map;
                 self.expander.exit(self.db, mark);
                 id
             }
@@ -680,7 +687,8 @@ impl ExprCollector<'_> {
     }
 
     fn collect_block(&mut self, block: ast::BlockExpr) -> ExprId {
-        let ast_id = self.expander.ast_id(self.db, &block);
+        let file_local_id = self.ast_id_map.ast_id(&block);
+        let ast_id = AstId::new(self.expander.current_file_id, file_local_id);
         let block_loc =
             BlockLoc { ast_id, module: self.expander.def_map.module_id(self.expander.module) };
         let block_id = self.db.intern_block(block_loc);