Return item tree query results by ref

4 files changed, 82 insertions, 99 deletions

diff --git a/src/tools/rust-analyzer/crates/hir-def/src/attr.rs b/src/tools/rust-analyzer/crates/hir-def/src/attr.rs
index d63bd391fd5..b509e69b0d3 100644
--- a/src/tools/rust-analyzer/crates/hir-def/src/attr.rs
+++ b/src/tools/rust-analyzer/crates/hir-def/src/attr.rs
@@ -26,6 +26,7 @@ use crate::{
     AdtId, AstIdLoc, AttrDefId, GenericParamId, HasModule, LocalFieldId, Lookup, MacroId,
     VariantId,
     db::DefDatabase,
+    item_tree::block_item_tree_query,
     lang_item::LangItem,
     nameres::{ModuleOrigin, ModuleSource},
     src::{HasChildSource, HasSource},
@@ -539,7 +540,7 @@ impl AttrsWithOwner {
                         definition_tree_id.item_tree(db).raw_attrs(definition.upcast()).clone()
                     }
                     ModuleOrigin::BlockExpr { id, .. } => {
-                        let tree = db.block_item_tree(id);
+                        let tree = block_item_tree_query(db, id);
                         tree.top_level_raw_attrs().clone()
                     }
                 };
diff --git a/src/tools/rust-analyzer/crates/hir-def/src/db.rs b/src/tools/rust-analyzer/crates/hir-def/src/db.rs
index 5f996931396..27fe62c4f0d 100644
--- a/src/tools/rust-analyzer/crates/hir-def/src/db.rs
+++ b/src/tools/rust-analyzer/crates/hir-def/src/db.rs
@@ -24,7 +24,7 @@ use crate::{
     },
     hir::generics::GenericParams,
     import_map::ImportMap,
-    item_tree::ItemTree,
+    item_tree::{ItemTree, file_item_tree_query},
     lang_item::{self, LangItem},
     nameres::{
         assoc::{ImplItems, TraitItems},
@@ -108,11 +108,9 @@ pub trait DefDatabase: InternDatabase + ExpandDatabase + SourceDatabase {
     fn expand_proc_attr_macros(&self) -> bool;
 
     /// Computes an [`ItemTree`] for the given file or macro expansion.
-    #[salsa::invoke(ItemTree::file_item_tree_query)]
-    fn file_item_tree(&self, file_id: HirFileId) -> Arc<ItemTree>;
-
-    #[salsa::invoke(ItemTree::block_item_tree_query)]
-    fn block_item_tree(&self, block_id: BlockId) -> Arc<ItemTree>;
+    #[salsa::invoke(file_item_tree_query)]
+    #[salsa::transparent]
+    fn file_item_tree(&self, file_id: HirFileId) -> &ItemTree;
 
     /// Turns a MacroId into a MacroDefId, describing the macro's definition post name resolution.
     #[salsa::invoke(macro_def)]
diff --git a/src/tools/rust-analyzer/crates/hir-def/src/item_tree.rs b/src/tools/rust-analyzer/crates/hir-def/src/item_tree.rs
index 0836d5758be..90a5dd2dc7d 100644
--- a/src/tools/rust-analyzer/crates/hir-def/src/item_tree.rs
+++ b/src/tools/rust-analyzer/crates/hir-def/src/item_tree.rs
@@ -85,6 +85,74 @@ impl fmt::Debug for RawVisibilityId {
     }
 }
 
+#[salsa_macros::tracked(returns(ref))]
+pub(crate) fn file_item_tree_query(db: &dyn DefDatabase, file_id: HirFileId) -> ItemTree {
+    let _p = tracing::info_span!("file_item_tree_query", ?file_id).entered();
+
+    let ctx = lower::Ctx::new(db, file_id);
+    let syntax = db.parse_or_expand(file_id);
+    let mut item_tree = match_ast! {
+        match syntax {
+            ast::SourceFile(file) => {
+                let top_attrs = RawAttrs::new(db, &file, ctx.span_map());
+                let mut item_tree = ctx.lower_module_items(&file);
+                item_tree.top_attrs = top_attrs;
+                item_tree
+            },
+            ast::MacroItems(items) => {
+                ctx.lower_module_items(&items)
+            },
+            ast::MacroStmts(stmts) => {
+                // The produced statements can include items, which should be added as top-level
+                // items.
+                ctx.lower_macro_stmts(stmts)
+            },
+            _ => {
+                if never!(syntax.kind() == SyntaxKind::ERROR, "{:?} from {:?} {}", file_id, syntax, syntax) {
+                    return Default::default();
+                }
+                panic!("cannot create item tree for file {file_id:?} from {syntax:?} {syntax}");
+            },
+        }
+    };
+
+    item_tree.shrink_to_fit();
+    item_tree
+}
+
+#[salsa_macros::tracked(returns(ref))]
+pub(crate) fn block_item_tree_query(db: &dyn DefDatabase, block: BlockId) -> Arc<ItemTree> {
+    let _p = tracing::info_span!("block_item_tree_query", ?block).entered();
+    // Blocks have a tendency to be empty due to macro calls that do not expand to items,
+    // so deduplicate this case via `Arc` to reduce the size of the query storage here.
+    static EMPTY: OnceLock<Arc<ItemTree>> = OnceLock::new();
+
+    let loc = block.lookup(db);
+    let block = loc.ast_id.to_node(db);
+
+    let ctx = lower::Ctx::new(db, loc.ast_id.file_id);
+    let mut item_tree = ctx.lower_block(&block);
+    if item_tree.data.is_empty()
+        && item_tree.top_level.is_empty()
+        && item_tree.attrs.is_empty()
+        && item_tree.top_attrs.is_empty()
+    {
+        EMPTY
+            .get_or_init(|| {
+                Arc::new(ItemTree {
+                    top_level: Box::new([]),
+                    attrs: FxHashMap::default(),
+                    data: FxHashMap::default(),
+                    top_attrs: RawAttrs::EMPTY,
+                    vis: ItemVisibilities { arena: Box::new([]) },
+                })
+            })
+            .clone()
+    } else {
+        item_tree.shrink_to_fit();
+        Arc::new(item_tree)
+    }
+}
 /// The item tree of a source file.
 #[derive(Debug, Default, Eq, PartialEq)]
 pub struct ItemTree {
@@ -97,90 +165,6 @@ pub struct ItemTree {
 }
 
 impl ItemTree {
-    pub(crate) fn file_item_tree_query(db: &dyn DefDatabase, file_id: HirFileId) -> Arc<ItemTree> {
-        let _p = tracing::info_span!("file_item_tree_query", ?file_id).entered();
-        static EMPTY: OnceLock<Arc<ItemTree>> = OnceLock::new();
-
-        let ctx = lower::Ctx::new(db, file_id);
-        let syntax = db.parse_or_expand(file_id);
-        let mut item_tree = match_ast! {
-            match syntax {
-                ast::SourceFile(file) => {
-                    let top_attrs = RawAttrs::new(db, &file, ctx.span_map());
-                    let mut item_tree = ctx.lower_module_items(&file);
-                    item_tree.top_attrs = top_attrs;
-                    item_tree
-                },
-                ast::MacroItems(items) => {
-                    ctx.lower_module_items(&items)
-                },
-                ast::MacroStmts(stmts) => {
-                    // The produced statements can include items, which should be added as top-level
-                    // items.
-                    ctx.lower_macro_stmts(stmts)
-                },
-                _ => {
-                    if never!(syntax.kind() == SyntaxKind::ERROR, "{:?} from {:?} {}", file_id, syntax, syntax) {
-                        return Default::default();
-                    }
-                    panic!("cannot create item tree for file {file_id:?} from {syntax:?} {syntax}");
-                },
-            }
-        };
-
-        if item_tree.data.is_empty()
-            && item_tree.top_level.is_empty()
-            && item_tree.attrs.is_empty()
-            && item_tree.top_attrs.is_empty()
-        {
-            EMPTY
-                .get_or_init(|| {
-                    Arc::new(ItemTree {
-                        top_level: Box::new([]),
-                        attrs: FxHashMap::default(),
-                        data: FxHashMap::default(),
-                        top_attrs: RawAttrs::EMPTY,
-                        vis: ItemVisibilities { arena: Box::new([]) },
-                    })
-                })
-                .clone()
-        } else {
-            item_tree.shrink_to_fit();
-            Arc::new(item_tree)
-        }
-    }
-
-    pub(crate) fn block_item_tree_query(db: &dyn DefDatabase, block: BlockId) -> Arc<ItemTree> {
-        let _p = tracing::info_span!("block_item_tree_query", ?block).entered();
-        static EMPTY: OnceLock<Arc<ItemTree>> = OnceLock::new();
-
-        let loc = block.lookup(db);
-        let block = loc.ast_id.to_node(db);
-
-        let ctx = lower::Ctx::new(db, loc.ast_id.file_id);
-        let mut item_tree = ctx.lower_block(&block);
-        if item_tree.data.is_empty()
-            && item_tree.top_level.is_empty()
-            && item_tree.attrs.is_empty()
-            && item_tree.top_attrs.is_empty()
-        {
-            EMPTY
-                .get_or_init(|| {
-                    Arc::new(ItemTree {
-                        top_level: Box::new([]),
-                        attrs: FxHashMap::default(),
-                        data: FxHashMap::default(),
-                        top_attrs: RawAttrs::EMPTY,
-                        vis: ItemVisibilities { arena: Box::new([]) },
-                    })
-                })
-                .clone()
-        } else {
-            item_tree.shrink_to_fit();
-            Arc::new(item_tree)
-        }
-    }
-
     /// Returns an iterator over all items located at the top level of the `HirFileId` this
     /// `ItemTree` was created from.
     pub(crate) fn top_level_items(&self) -> &[ModItemId] {
@@ -297,10 +281,10 @@ impl TreeId {
         Self { file, block }
     }
 
-    pub(crate) fn item_tree(&self, db: &dyn DefDatabase) -> Arc<ItemTree> {
+    pub(crate) fn item_tree<'db>(&self, db: &'db dyn DefDatabase) -> &'db ItemTree {
         match self.block {
-            Some(block) => db.block_item_tree(block),
-            None => db.file_item_tree(self.file),
+            Some(block) => block_item_tree_query(db, block),
+            None => file_item_tree_query(db, self.file),
         }
     }
 
diff --git a/src/tools/rust-analyzer/crates/hir-def/src/nameres/collector.rs b/src/tools/rust-analyzer/crates/hir-def/src/nameres/collector.rs
index 4d19df01952..3b2488c8ea6 100644
--- a/src/tools/rust-analyzer/crates/hir-def/src/nameres/collector.rs
+++ b/src/tools/rust-analyzer/crates/hir-def/src/nameres/collector.rs
@@ -357,7 +357,7 @@ impl DefCollector<'_> {
             macro_depth: 0,
             module_id: DefMap::ROOT,
             tree_id: TreeId::new(file_id.into(), None),
-            item_tree: &item_tree,
+            item_tree,
             mod_dir: ModDir::root(),
         }
         .collect_in_top_module(item_tree.top_level_items());
@@ -378,7 +378,7 @@ impl DefCollector<'_> {
                 macro_depth: 0,
                 module_id: DefMap::ROOT,
                 tree_id,
-                item_tree: &item_tree,
+                item_tree,
                 mod_dir: ModDir::root(),
             }
             .collect_in_top_module(item_tree.top_level_items());
@@ -1381,7 +1381,7 @@ impl DefCollector<'_> {
                             macro_depth: directive.depth,
                             module_id: directive.module_id,
                             tree_id: *tree,
-                            item_tree: &item_tree,
+                            item_tree,
                             mod_dir,
                         }
                         .collect(&[*mod_item], directive.container);
@@ -1570,7 +1570,7 @@ impl DefCollector<'_> {
             macro_depth: depth,
             tree_id: TreeId::new(file_id, None),
             module_id,
-            item_tree: &item_tree,
+            item_tree,
             mod_dir,
         }
         .collect(item_tree.top_level_items(), container);
@@ -2142,7 +2142,7 @@ impl ModCollector<'_, '_> {
                                     macro_depth: self.macro_depth,
                                     module_id,
                                     tree_id: TreeId::new(file_id.into(), None),
-                                    item_tree: &item_tree,
+                                    item_tree,
                                     mod_dir,
                                 }
                                 .collect_in_top_module(item_tree.top_level_items());