1 files changed, 21 insertions, 20 deletions

diff --git a/crates/hir_def/src/item_tree.rs b/crates/hir_def/src/item_tree.rs
index 9a433b61c..b8d7608e7 100644
--- a/crates/hir_def/src/item_tree.rs
+++ b/crates/hir_def/src/item_tree.rs
@@ -21,6 +21,7 @@ use hir_expand::{
     HirFileId, InFile,
 };
 use la_arena::{Arena, Idx, RawIdx};
+use profile::Count;
 use rustc_hash::FxHashMap;
 use smallvec::SmallVec;
 use syntax::{ast, match_ast};
@@ -67,15 +68,16 @@ impl GenericParamsId {
 /// The item tree of a source file.
 #[derive(Debug, Eq, PartialEq)]
 pub struct ItemTree {
+    _c: Count<Self>,
+
     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
@@ -118,9 +120,9 @@ impl ItemTree {
 
     fn empty() -> Self {
         Self {
+            _c: Count::new(),
             top_level: Default::default(),
             attrs: Default::default(),
-            inner_items: Default::default(),
             data: Default::default(),
         }
     }
@@ -145,9 +147,9 @@ impl ItemTree {
                 macro_calls,
                 macro_rules,
                 macro_defs,
-                exprs,
                 vis,
                 generics,
+                inner_items,
             } = &mut **data;
 
             imports.shrink_to_fit();
@@ -167,10 +169,11 @@ impl ItemTree {
             macro_calls.shrink_to_fit();
             macro_rules.shrink_to_fit();
             macro_defs.shrink_to_fit();
-            exprs.shrink_to_fit();
 
             vis.arena.shrink_to_fit();
             generics.arena.shrink_to_fit();
+
+            inner_items.shrink_to_fit();
         }
     }
 
@@ -193,16 +196,18 @@ 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> + '_ {
+        match &self.data {
+            Some(data) => Some(data.inner_items.values().flatten().copied()).into_iter().flatten(),
+            None => None.into_iter().flatten(),
+        }
     }
 
-    pub fn all_inner_items(&self) -> impl Iterator<Item = ModItem> + '_ {
-        self.inner_items.values().flatten().copied()
+    pub fn inner_items_of_block(&self, block: FileAstId<ast::BlockExpr>) -> &[ModItem] {
+        match &self.data {
+            Some(data) => data.inner_items.get(&block).map(|it| &**it).unwrap_or(&[]),
+            None => &[],
+        }
     }
 
     pub fn source<S: ItemTreeNode>(&self, db: &dyn DefDatabase, of: ItemTreeId<S>) -> S::Source {
@@ -296,10 +301,11 @@ struct ItemTreeData {
     macro_calls: Arena<MacroCall>,
     macro_rules: Arena<MacroRules>,
     macro_defs: Arena<MacroDef>,
-    exprs: Arena<Expr>,
 
     vis: ItemVisibilities,
     generics: GenericParamsStorage,
+
+    inner_items: FxHashMap<FileAstId<ast::BlockExpr>, SmallVec<[ModItem; 1]>>,
 }
 
 #[derive(Debug, Eq, PartialEq, Hash)]
@@ -461,7 +467,7 @@ macro_rules! impl_index {
     };
 }
 
-impl_index!(fields: Field, variants: Variant, exprs: Expr);
+impl_index!(fields: Field, variants: Variant);
 
 impl Index<RawVisibilityId> for ItemTree {
     type Output = RawVisibility;
@@ -664,11 +670,6 @@ pub struct MacroDef {
     pub ast_id: FileAstId<ast::MacroDef>,
 }
 
-// NB: There's no `FileAstId` for `Expr`. The only case where this would be useful is for array
-// lengths, but we don't do much with them yet.
-#[derive(Debug, Clone, Eq, PartialEq)]
-pub struct Expr;
-
 macro_rules! impl_froms {
     ($e:ident { $($v:ident ($t:ty)),* $(,)? }) => {
         $(