internal: Don't eagerly construct AstIdMaps

Author: Veykril

crates/hir-def/src/body.rs

@@ -50,7 +50,7 @@ pub struct Expander {
     cfg_expander: CfgExpander,
     def_map: Arc<DefMap>,
     current_file_id: HirFileId,
-    ast_id_map: Arc<AstIdMap>,
+    ast_id_map: Option<Arc<AstIdMap>>,
     module: LocalModuleId,
     recursion_limit: usize,
 }
@@ -80,12 +80,11 @@ impl Expander {
     pub fn new(db: &dyn DefDatabase, current_file_id: HirFileId, module: ModuleId) -> Expander {
         let cfg_expander = CfgExpander::new(db, current_file_id, module.krate);
         let def_map = module.def_map(db);
-        let ast_id_map = db.ast_id_map(current_file_id);
         Expander {
             cfg_expander,
             def_map,
             current_file_id,
-            ast_id_map,
+            ast_id_map: None,
             module: module.local_id,
             recursion_limit: 0,
         }
@@ -175,7 +174,7 @@ impl Expander {
         };
         self.cfg_expander.hygiene = Hygiene::new(db.upcast(), file_id);
         self.current_file_id = file_id;
-        self.ast_id_map = db.ast_id_map(file_id);
+        self.ast_id_map = None;
 
         ExpandResult { value: Some((mark, node)), err }
     }
@@ -213,8 +212,9 @@ impl Expander {
         self.def_map.resolve_path(db, self.module, path, BuiltinShadowMode::Other).0.take_macros()
     }
 
-    fn ast_id<N: AstNode>(&self, item: &N) -> AstId<N> {
-        let file_local_id = self.ast_id_map.ast_id(item);
+    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)
     }
 
@@ -233,7 +233,7 @@ impl Expander {
 #[derive(Debug)]
 pub struct Mark {
     file_id: HirFileId,
-    ast_id_map: Arc<AstIdMap>,
+    ast_id_map: Option<Arc<AstIdMap>>,
     bomb: DropBomb,
 }
 

crates/hir-def/src/body/lower.rs

@@ -11,6 +11,7 @@ use hir_expand::{
     ExpandError, HirFileId, InFile,
 };
 use la_arena::Arena;
+use once_cell::unsync::OnceCell;
 use profile::Count;
 use rustc_hash::FxHashMap;
 use syntax::{
@@ -41,38 +42,42 @@ use crate::{
 pub struct LowerCtx<'a> {
     pub db: &'a dyn DefDatabase,
     hygiene: Hygiene,
-    file_id: Option<HirFileId>,
-    source_ast_id_map: Option<Arc<AstIdMap>>,
+    ast_id_map: Option<(HirFileId, OnceCell<Arc<AstIdMap>>)>,
 }
 
 impl<'a> LowerCtx<'a> {
     pub fn new(db: &'a dyn DefDatabase, file_id: HirFileId) -> Self {
         LowerCtx {
             db,
             hygiene: Hygiene::new(db.upcast(), file_id),
-            file_id: Some(file_id),
-            source_ast_id_map: Some(db.ast_id_map(file_id)),
+            ast_id_map: Some((file_id, OnceCell::new())),
         }
     }
 
     pub fn with_hygiene(db: &'a dyn DefDatabase, hygiene: &Hygiene) -> Self {
-        LowerCtx { db, hygiene: hygiene.clone(), file_id: None, source_ast_id_map: None }
+        LowerCtx { db, hygiene: hygiene.clone(), ast_id_map: None }
     }
 
     pub(crate) fn hygiene(&self) -> &Hygiene {
         &self.hygiene
     }
 
     pub(crate) fn file_id(&self) -> HirFileId {
-        self.file_id.unwrap()
+        self.ast_id_map.as_ref().unwrap().0
     }
 
     pub(crate) fn lower_path(&self, ast: ast::Path) -> Option<Path> {
         Path::from_src(ast, self)
     }
 
-    pub(crate) fn ast_id<N: AstNode>(&self, item: &N) -> Option<FileAstId<N>> {
-        self.source_ast_id_map.as_ref().map(|ast_id_map| ast_id_map.ast_id(item))
+    pub(crate) fn ast_id<N: AstNode>(
+        &self,
+        db: &dyn DefDatabase,
+        item: &N,
+    ) -> Option<FileAstId<N>> {
+        let (file_id, ast_id_map) = self.ast_id_map.as_ref()?;
+        let ast_id_map = ast_id_map.get_or_init(|| db.ast_id_map(*file_id));
+        Some(ast_id_map.ast_id(item))
     }
 }
 
@@ -675,7 +680,7 @@ impl ExprCollector<'_> {
     }
 
     fn collect_block(&mut self, block: ast::BlockExpr) -> ExprId {
-        let ast_id = self.expander.ast_id(&block);
+        let ast_id = self.expander.ast_id(self.db, &block);
         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);

crates/hir-def/src/data.rs

@@ -550,14 +550,17 @@ impl<'a> AssocItemCollector<'a> {
                 AssocItem::MacroCall(call) => {
                     let call = &item_tree[call];
                     let ast_id_map = self.db.ast_id_map(self.expander.current_file_id());
-                    let root = self.db.parse_or_expand(self.expander.current_file_id()).unwrap();
-                    let call = ast_id_map.get(call.ast_id).to_node(&root);
-                    let _cx =
-                        stdx::panic_context::enter(format!("collect_items MacroCall: {}", call));
-                    let res = self.expander.enter_expand(self.db, call);
-
-                    if let Ok(ExpandResult { value: Some((mark, mac)), .. }) = res {
-                        self.collect_macro_items(mark, mac);
+                    if let Some(root) = self.db.parse_or_expand(self.expander.current_file_id()) {
+                        let call = ast_id_map.get(call.ast_id).to_node(&root);
+                        let _cx = stdx::panic_context::enter(format!(
+                            "collect_items MacroCall: {}",
+                            call
+                        ));
+                        let res = self.expander.enter_expand(self.db, call);
+
+                        if let Ok(ExpandResult { value: Some((mark, mac)), .. }) = res {
+                            self.collect_macro_items(mark, mac);
+                        }
                     }
                 }
             }

crates/hir-def/src/type_ref.rs

@@ -237,7 +237,7 @@ impl TypeRef {
             }
             ast::Type::MacroType(mt) => match mt.macro_call() {
                 Some(mc) => ctx
-                    .ast_id(&mc)
+                    .ast_id(ctx.db, &mc)
                     .map(|mc| TypeRef::Macro(InFile::new(ctx.file_id(), mc)))
                     .unwrap_or(TypeRef::Error),
                 None => TypeRef::Error,