Merge #7359

7359: ItemTree: store a mapping from blocks to inner items r=jonas-schievink a=jonas-schievink

To do name resolution within block expressions, we need to know which inner items are located inside each block expression. This adds such a mapping to `ItemTree`, replacing the previous one, which was seemingly unused other than to access all the inner items.

This also assigns `AstId`s to block expressions, which is needed to store the mapping in salsa.

Co-authored-by: Jonas Schievink <[email protected]>

3 files changed, 49 insertions, 32 deletions

diff --git a/crates/hir_def/src/item_tree.rs b/crates/hir_def/src/item_tree.rs
index ff62928df..6494cebd3 100644
--- a/crates/hir_def/src/item_tree.rs
+++ b/crates/hir_def/src/item_tree.rs
@@ -69,13 +69,12 @@ impl GenericParamsId {
 pub struct ItemTree {
     top_level: SmallVec<[ModItem; 1]>,
     attrs: FxHashMap<AttrOwner, RawAttrs>,
-    inner_items: FxHashMap<FileAstId<ast::Item>, SmallVec<[ModItem; 1]>>,
 
     data: Option<Box<ItemTreeData>>,
 }
 
 impl ItemTree {
-    pub fn item_tree_query(db: &dyn DefDatabase, file_id: HirFileId) -> Arc<ItemTree> {
+    pub(crate) fn item_tree_query(db: &dyn DefDatabase, file_id: HirFileId) -> Arc<ItemTree> {
         let _p = profile::span("item_tree_query").detail(|| format!("{:?}", file_id));
         let syntax = if let Some(node) = db.parse_or_expand(file_id) {
             node
@@ -117,12 +116,7 @@ impl ItemTree {
     }
 
     fn empty() -> Self {
-        Self {
-            top_level: Default::default(),
-            attrs: Default::default(),
-            inner_items: Default::default(),
-            data: Default::default(),
-        }
+        Self { top_level: Default::default(), attrs: Default::default(), data: Default::default() }
     }
 
     fn shrink_to_fit(&mut self) {
@@ -147,6 +141,7 @@ impl ItemTree {
                 macro_defs,
                 vis,
                 generics,
+                inner_items,
             } = &mut **data;
 
             imports.shrink_to_fit();
@@ -169,6 +164,8 @@ impl ItemTree {
 
             vis.arena.shrink_to_fit();
             generics.arena.shrink_to_fit();
+
+            inner_items.shrink_to_fit();
         }
     }
 
@@ -191,16 +188,11 @@ impl ItemTree {
         self.raw_attrs(of).clone().filter(db, krate)
     }
 
-    /// Returns the lowered inner items that `ast` corresponds to.
-    ///
-    /// Most AST items are lowered to a single `ModItem`, but some (eg. `use` items) may be lowered
-    /// to multiple items in the `ItemTree`.
-    pub fn inner_items(&self, ast: FileAstId<ast::Item>) -> &[ModItem] {
-        &self.inner_items[&ast]
-    }
-
     pub fn all_inner_items(&self) -> impl Iterator<Item = ModItem> + '_ {
-        self.inner_items.values().flatten().copied()
+        match &self.data {
+            Some(data) => Some(data.inner_items.values().flatten().copied()).into_iter().flatten(),
+            None => None.into_iter().flatten(),
+        }
     }
 
     pub fn source<S: ItemTreeNode>(&self, db: &dyn DefDatabase, of: ItemTreeId<S>) -> S::Source {
@@ -297,6 +289,8 @@ struct ItemTreeData {
 
     vis: ItemVisibilities,
     generics: GenericParamsStorage,
+
+    inner_items: FxHashMap<FileAstId<ast::BlockExpr>, SmallVec<[ModItem; 1]>>,
 }
 
 #[derive(Debug, Eq, PartialEq, Hash)]
diff --git a/crates/hir_def/src/item_tree/lower.rs b/crates/hir_def/src/item_tree/lower.rs
index 5e71ca42c..56fe569ff 100644
--- a/crates/hir_def/src/item_tree/lower.rs
+++ b/crates/hir_def/src/item_tree/lower.rs
@@ -6,7 +6,7 @@ use hir_expand::{ast_id_map::AstIdMap, hygiene::Hygiene, name::known, HirFileId}
 use smallvec::SmallVec;
 use syntax::{
     ast::{self, ModuleItemOwner},
-    SyntaxNode,
+    SyntaxNode, WalkEvent,
 };
 
 use crate::{
@@ -150,14 +150,29 @@ impl Ctx {
 
     fn collect_inner_items(&mut self, container: &SyntaxNode) {
         let forced_vis = self.forced_visibility.take();
-        let mut inner_items = mem::take(&mut self.tree.inner_items);
-        inner_items.extend(container.descendants().skip(1).filter_map(ast::Item::cast).filter_map(
-            |item| {
-                let ast_id = self.source_ast_id_map.ast_id(&item);
-                Some((ast_id, self.lower_mod_item(&item, true)?.0))
-            },
-        ));
-        self.tree.inner_items = inner_items;
+
+        let mut current_block = None;
+        for event in container.preorder().skip(1) {
+            if let WalkEvent::Enter(node) = event {
+                match_ast! {
+                    match node {
+                        ast::BlockExpr(block) => {
+                            current_block = Some(self.source_ast_id_map.ast_id(&block));
+                        },
+                        ast::Item(item) => {
+                            let mod_items = self.lower_mod_item(&item, true);
+                            if let (Some(mod_items), Some(block)) = (mod_items, current_block) {
+                                if !mod_items.0.is_empty() {
+                                    self.data().inner_items.entry(block).or_default().extend(mod_items.0.iter().copied());
+                                }
+                            }
+                        },
+                        _ => {}
+                    }
+                }
+            }
+        }
+
         self.forced_visibility = forced_vis;
     }
 
diff --git a/crates/hir_expand/src/ast_id_map.rs b/crates/hir_expand/src/ast_id_map.rs
index 2401b0cc5..0991fffd8 100644
--- a/crates/hir_expand/src/ast_id_map.rs
+++ b/crates/hir_expand/src/ast_id_map.rs
@@ -13,7 +13,7 @@ use std::{
 };
 
 use la_arena::{Arena, Idx};
-use syntax::{ast, AstNode, AstPtr, SyntaxNode, SyntaxNodePtr};
+use syntax::{ast, match_ast, AstNode, AstPtr, SyntaxNode, SyntaxNodePtr};
 
 /// `AstId` points to an AST node in a specific file.
 pub struct FileAstId<N: AstNode> {
@@ -72,12 +72,20 @@ impl AstIdMap {
         // get lower ids then children. That is, adding a new child does not
         // change parent's id. This means that, say, adding a new function to a
         // trait does not change ids of top-level items, which helps caching.
-        bdfs(node, |it| match ast::Item::cast(it) {
-            Some(module_item) => {
-                res.alloc(module_item.syntax());
-                true
+        bdfs(node, |it| {
+            match_ast! {
+                match it {
+                    ast::Item(module_item) => {
+                        res.alloc(module_item.syntax());
+                        true
+                    },
+                    ast::BlockExpr(block) => {
+                        res.alloc(block.syntax());
+                        true
+                    },
+                    _ => false,
+                }
             }
-            None => false,
         });
         res
     }