Turn BlockLoc into a tracked struct

Author: Veykril

crates/hir-def/src/db.rs

@@ -8,12 +8,12 @@ use la_arena::ArenaMap;
 use triomphe::Arc;
 
 use crate::{
-    AssocItemId, AttrDefId, BlockId, BlockLoc, ConstId, ConstLoc, DefWithBodyId, EnumId, EnumLoc,
-    EnumVariantId, EnumVariantLoc, ExternBlockId, ExternBlockLoc, ExternCrateId, ExternCrateLoc,
-    FunctionId, FunctionLoc, GenericDefId, ImplId, ImplLoc, LocalFieldId, Macro2Id, Macro2Loc,
-    MacroExpander, MacroId, MacroRulesId, MacroRulesLoc, MacroRulesLocFlags, ProcMacroId,
-    ProcMacroLoc, StaticId, StaticLoc, StructId, StructLoc, TraitId, TraitLoc, TypeAliasId,
-    TypeAliasLoc, UnionId, UnionLoc, UseId, UseLoc, VariantId,
+    AssocItemId, AttrDefId, ConstId, ConstLoc, DefWithBodyId, EnumId, EnumLoc, EnumVariantId,
+    EnumVariantLoc, ExternBlockId, ExternBlockLoc, ExternCrateId, ExternCrateLoc, FunctionId,
+    FunctionLoc, GenericDefId, ImplId, ImplLoc, LocalFieldId, Macro2Id, Macro2Loc, MacroExpander,
+    MacroId, MacroRulesId, MacroRulesLoc, MacroRulesLocFlags, ProcMacroId, ProcMacroLoc, StaticId,
+    StaticLoc, StructId, StructLoc, TraitId, TraitLoc, TypeAliasId, TypeAliasLoc, UnionId,
+    UnionLoc, UseId, UseLoc, VariantId,
     attrs::AttrFlags,
     expr_store::{
         Body, BodySourceMap, ExpressionStore, ExpressionStoreSourceMap, scope::ExprScopes,
@@ -82,9 +82,6 @@ pub trait InternDatabase: RootQueryDb {
     #[salsa::interned]
     fn intern_macro_rules(&self, loc: MacroRulesLoc) -> MacroRulesId;
     // endregion: items
-
-    #[salsa::interned]
-    fn intern_block(&self, loc: BlockLoc) -> BlockId;
 }
 
 #[query_group::query_group]

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

@@ -31,7 +31,7 @@ use triomphe::Arc;
 use tt::TextRange;
 
 use crate::{
-    AdtId, BlockId, BlockLoc, DefWithBodyId, FunctionId, GenericDefId, ImplId, MacroId,
+    AdtId, BlockId, BlockIdLt, DefWithBodyId, FunctionId, GenericDefId, ImplId, MacroId,
     ModuleDefId, ModuleId, TraitId, TypeAliasId, UnresolvedMacro,
     attrs::AttrFlags,
     builtin_type::BuiltinUint,
@@ -2114,7 +2114,7 @@ impl<'db> ExprCollector<'db> {
     ) -> ExprId {
         let block_id = self.expander.ast_id_map().ast_id_for_block(&block).map(|file_local_id| {
             let ast_id = self.expander.in_file(file_local_id);
-            self.db.intern_block(BlockLoc { ast_id, module: self.module })
+            unsafe { BlockIdLt::new(self.db, ast_id, self.module).to_static() }
         });
 
         let (module, def_map) =

crates/hir-def/src/find_path.rs

@@ -12,7 +12,7 @@ use intern::sym;
 use rustc_hash::FxHashSet;
 
 use crate::{
-    FindPathConfig, ModuleDefId, ModuleId,
+    FindPathConfig, ModuleDefId, ModuleIdLt,
     db::DefDatabase,
     item_scope::ItemInNs,
     nameres::DefMap,
@@ -24,7 +24,7 @@ use crate::{
 pub fn find_path(
     db: &dyn DefDatabase,
     item: ItemInNs,
-    from: ModuleId,
+    from: ModuleIdLt<'_>,
     mut prefix_kind: PrefixKind,
     ignore_local_imports: bool,
     mut cfg: FindPathConfig,
@@ -102,14 +102,14 @@ struct FindPathCtx<'db> {
     cfg: FindPathConfig,
     ignore_local_imports: bool,
     is_std_item: bool,
-    from: ModuleId,
+    from: ModuleIdLt<'db>,
     from_crate: Crate,
-    crate_root: ModuleId,
+    crate_root: ModuleIdLt<'db>,
     from_def_map: &'db DefMap,
     fuel: Cell<usize>,
 }
 
-/// Attempts to find a path to refer to the given `item` visible from the `from` ModuleId
+/// Attempts to find a path to refer to the given `item` visible from the `from` ModuleIdLt<'_>
 fn find_path_inner(ctx: &FindPathCtx<'_>, item: ItemInNs, max_len: usize) -> Option<ModPath> {
     // - if the item is a module, jump straight to module search
     if !ctx.is_std_item
@@ -157,10 +157,10 @@ fn find_path_inner(ctx: &FindPathCtx<'_>, item: ItemInNs, max_len: usize) -> Opt
 }
 
 #[tracing::instrument(skip_all)]
-fn find_path_for_module(
-    ctx: &FindPathCtx<'_>,
-    visited_modules: &mut FxHashSet<(ItemInNs, ModuleId)>,
-    module_id: ModuleId,
+fn find_path_for_module<'db>(
+    ctx: &'db FindPathCtx<'db>,
+    visited_modules: &mut FxHashSet<(ItemInNs, ModuleIdLt<'db>)>,
+    module_id: ModuleIdLt<'db>,
     maybe_extern: bool,
     max_len: usize,
 ) -> Option<Choice> {
@@ -217,7 +217,7 @@ fn find_path_for_module(
             ctx.db,
             ctx.from_def_map,
             ctx.from,
-            ItemInNs::Types(module_id.into()),
+            ItemInNs::Types(unsafe { module_id.to_static() }.into()),
             ctx.ignore_local_imports,
         );
         if let Some(scope_name) = scope_name {
@@ -244,7 +244,7 @@ fn find_path_for_module(
     }
 
     // - if the module is in the prelude, return it by that path
-    let item = ItemInNs::Types(module_id.into());
+    let item = ItemInNs::Types(unsafe { module_id.to_static() }.into());
     if let Some(choice) = find_in_prelude(ctx.db, ctx.from_def_map, item, ctx.from) {
         return Some(choice);
     }
@@ -257,10 +257,10 @@ fn find_path_for_module(
     best_choice
 }
 
-fn find_in_scope(
-    db: &dyn DefDatabase,
+fn find_in_scope<'db>(
+    db: &'db dyn DefDatabase,
     def_map: &DefMap,
-    from: ModuleId,
+    from: ModuleIdLt<'db>,
     item: ItemInNs,
     ignore_local_imports: bool,
 ) -> Option<Name> {
@@ -278,7 +278,7 @@ fn find_in_prelude(
     db: &dyn DefDatabase,
     local_def_map: &DefMap,
     item: ItemInNs,
-    from: ModuleId,
+    from: ModuleIdLt<'_>,
 ) -> Option<Choice> {
     let (prelude_module, _) = local_def_map.prelude()?;
     let prelude_def_map = prelude_module.def_map(db);
@@ -310,8 +310,8 @@ fn find_in_prelude(
 fn is_kw_kind_relative_to_from(
     db: &dyn DefDatabase,
     def_map: &DefMap,
-    item: ModuleId,
-    from: ModuleId,
+    item: ModuleIdLt<'_>,
+    from: ModuleIdLt<'_>,
 ) -> Option<PathKind> {
     if item.krate(db) != from.krate(db) || item.block(db).is_some() || from.block(db).is_some() {
         return None;
@@ -332,9 +332,9 @@ fn is_kw_kind_relative_to_from(
 }
 
 #[tracing::instrument(skip_all)]
-fn calculate_best_path(
-    ctx: &FindPathCtx<'_>,
-    visited_modules: &mut FxHashSet<(ItemInNs, ModuleId)>,
+fn calculate_best_path<'db>(
+    ctx: &'db FindPathCtx<'db>,
+    visited_modules: &mut FxHashSet<(ItemInNs, ModuleIdLt<'db>)>,
     item: ItemInNs,
     max_len: usize,
     best_choice: &mut Option<Choice>,
@@ -372,9 +372,9 @@ fn calculate_best_path(
     }
 }
 
-fn find_in_sysroot(
-    ctx: &FindPathCtx<'_>,
-    visited_modules: &mut FxHashSet<(ItemInNs, ModuleId)>,
+fn find_in_sysroot<'db>(
+    ctx: &'db FindPathCtx<'db>,
+    visited_modules: &mut FxHashSet<(ItemInNs, ModuleIdLt<'db>)>,
     item: ItemInNs,
     max_len: usize,
     best_choice: &mut Option<Choice>,
@@ -418,9 +418,9 @@ fn find_in_sysroot(
         });
 }
 
-fn find_in_dep(
-    ctx: &FindPathCtx<'_>,
-    visited_modules: &mut FxHashSet<(ItemInNs, ModuleId)>,
+fn find_in_dep<'db>(
+    ctx: &'db FindPathCtx<'db>,
+    visited_modules: &mut FxHashSet<(ItemInNs, ModuleIdLt<'db>)>,
     item: ItemInNs,
     max_len: usize,
     best_choice: &mut Option<Choice>,
@@ -461,9 +461,9 @@ fn find_in_dep(
     }
 }
 
-fn calculate_best_path_local(
-    ctx: &FindPathCtx<'_>,
-    visited_modules: &mut FxHashSet<(ItemInNs, ModuleId)>,
+fn calculate_best_path_local<'db>(
+    ctx: &'db FindPathCtx<'db>,
+    visited_modules: &mut FxHashSet<(ItemInNs, ModuleIdLt<'db>)>,
     item: ItemInNs,
     max_len: usize,
     best_choice: &mut Option<Choice>,
@@ -558,11 +558,11 @@ fn path_kind_len(kind: PathKind) -> usize {
 }
 
 /// Finds locations in `from.krate` from which `item` can be imported by `from`.
-fn find_local_import_locations(
-    ctx: &FindPathCtx<'_>,
+fn find_local_import_locations<'db>(
+    ctx: &'db FindPathCtx<'db>,
     item: ItemInNs,
-    visited_modules: &mut FxHashSet<(ItemInNs, ModuleId)>,
-    mut cb: impl FnMut(&mut FxHashSet<(ItemInNs, ModuleId)>, &Name, ModuleId),
+    visited_modules: &mut FxHashSet<(ItemInNs, ModuleIdLt<'db>)>,
+    mut cb: impl FnMut(&mut FxHashSet<(ItemInNs, ModuleIdLt<'db>)>, &Name, ModuleIdLt<'db>),
 ) {
     let _p = tracing::info_span!("find_local_import_locations").entered();
     let db = ctx.db;

crates/hir-def/src/import_map.rs

@@ -496,7 +496,7 @@ mod tests {
     use expect_test::{Expect, expect};
     use test_fixture::WithFixture;
 
-    use crate::{ItemContainerId, Lookup, nameres::assoc::TraitItems, test_db::TestDB};
+    use crate::{ItemContainerId, Lookup, ModuleIdLt, nameres::assoc::TraitItems, test_db::TestDB};
 
     use super::*;
 
@@ -628,8 +628,8 @@ mod tests {
         expect.assert_eq(&actual)
     }
 
-    fn render_path(db: &dyn DefDatabase, info: &ImportInfo) -> String {
-        let mut module = info.container;
+    fn render_path<'db>(db: &'db dyn DefDatabase, info: &ImportInfo) -> String {
+        let mut module: ModuleIdLt<'db> = info.container;
         let mut segments = vec![&info.name];
 
         let def_map = module.def_map(db);

crates/hir-def/src/item_tree.rs

@@ -58,7 +58,7 @@ use syntax::{SyntaxKind, ast, match_ast};
 use thin_vec::ThinVec;
 use triomphe::Arc;
 
-use crate::{BlockId, Lookup, db::DefDatabase};
+use crate::{BlockId, db::DefDatabase};
 
 pub(crate) use crate::item_tree::{
     attrs::*,
@@ -150,10 +150,10 @@ pub(crate) fn block_item_tree_query(db: &dyn DefDatabase, block: BlockId) -> Arc
     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 ast_id = block.ast_id(db);
+    let block = ast_id.to_node(db);
 
-    let ctx = lower::Ctx::new(db, loc.ast_id.file_id);
+    let ctx = lower::Ctx::new(db, ast_id.file_id);
     let mut item_tree = ctx.lower_block(&block);
     let ItemTree { top_level, top_attrs, attrs, vis, big_data, small_data } = &item_tree;
     if small_data.is_empty()

crates/hir-def/src/lib.rs

@@ -420,13 +420,31 @@ pub struct ProcMacroLoc {
 impl_intern!(ProcMacroId, ProcMacroLoc, intern_proc_macro, lookup_intern_proc_macro);
 impl_loc!(ProcMacroLoc, id: Fn, container: ModuleId);
 
-#[derive(Debug, Hash, PartialEq, Eq, Clone)]
-pub struct BlockLoc {
+#[salsa_macros::tracked(debug)]
+#[derive(PartialOrd, Ord)]
+pub struct BlockIdLt<'db> {
     pub ast_id: AstId<ast::BlockExpr>,
     /// The containing module.
-    pub module: ModuleId,
+    pub module: ModuleIdLt<'db>,
+}
+pub type BlockId = BlockIdLt<'static>;
+
+impl BlockIdLt<'_> {
+    /// # Safety
+    ///
+    /// The caller must ensure that the `ModuleId` is not leaked outside of query computations.
+    pub unsafe fn to_static(self) -> BlockId {
+        unsafe { std::mem::transmute(self) }
+    }
+}
+impl BlockId {
+    /// # Safety
+    ///
+    /// The caller must ensure that the `BlockId` comes from the given database.
+    pub unsafe fn to_db<'db>(self, _db: &'db dyn DefDatabase) -> BlockIdLt<'db> {
+        unsafe { std::mem::transmute(self) }
+    }
 }
-impl_intern!(BlockId, BlockLoc, intern_block, lookup_intern_block);
 
 #[salsa_macros::tracked(debug)]
 #[derive(PartialOrd, Ord)]
@@ -436,34 +454,26 @@ pub struct ModuleIdLt<'db> {
     /// If this `ModuleId` was derived from a `DefMap` for a block expression, this stores the
     /// `BlockId` of that block expression. If `None`, this module is part of the crate-level
     /// `DefMap` of `krate`.
-    pub block: Option<BlockId>,
+    pub block: Option<BlockIdLt<'db>>,
 }
 pub type ModuleId = ModuleIdLt<'static>;
 
-impl ModuleIdLt<'_> {
+impl<'db> ModuleIdLt<'db> {
     /// # Safety
     ///
     /// The caller must ensure that the `ModuleId` is not leaked outside of query computations.
     pub unsafe fn to_static(self) -> ModuleId {
         unsafe { std::mem::transmute(self) }
     }
-}
-impl ModuleId {
-    /// # Safety
-    ///
-    /// The caller must ensure that the `ModuleId` comes from the given database.
-    pub unsafe fn to_db<'db>(self, _db: &'db dyn DefDatabase) -> ModuleIdLt<'db> {
-        unsafe { std::mem::transmute(self) }
-    }
 
-    pub fn def_map(self, db: &dyn DefDatabase) -> &DefMap {
+    pub fn def_map(self, db: &'db dyn DefDatabase) -> &'db DefMap {
         match self.block(db) {
             Some(block) => block_def_map(db, block),
             None => crate_def_map(db, self.krate(db)),
         }
     }
 
-    pub(crate) fn local_def_map(self, db: &dyn DefDatabase) -> (&DefMap, &LocalDefMap) {
+    pub(crate) fn local_def_map(self, db: &'db dyn DefDatabase) -> (&'db DefMap, &'db LocalDefMap) {
         match self.block(db) {
             Some(block) => (block_def_map(db, block), self.only_local_def_map(db)),
             None => {
@@ -473,15 +483,15 @@ impl ModuleId {
         }
     }
 
-    pub(crate) fn only_local_def_map(self, db: &dyn DefDatabase) -> &LocalDefMap {
+    pub(crate) fn only_local_def_map(self, db: &'db dyn DefDatabase) -> &'db LocalDefMap {
         crate_local_def_map(db, self.krate(db)).local(db)
     }
 
-    pub fn crate_def_map(self, db: &dyn DefDatabase) -> &DefMap {
+    pub fn crate_def_map(self, db: &'db dyn DefDatabase) -> &'db DefMap {
         crate_def_map(db, self.krate(db))
     }
 
-    pub fn name(self, db: &dyn DefDatabase) -> Option<Name> {
+    pub fn name(self, db: &'db dyn DefDatabase) -> Option<Name> {
         let def_map = self.def_map(db);
         let parent = def_map[self].parent?;
         def_map[parent].children.iter().find_map(|(name, module_id)| {
@@ -491,15 +501,24 @@ impl ModuleId {
 
     /// Returns the module containing `self`, either the parent `mod`, or the module (or block) containing
     /// the block, if `self` corresponds to a block expression.
-    pub fn containing_module(self, db: &dyn DefDatabase) -> Option<ModuleId> {
+    pub fn containing_module(self, db: &'db dyn DefDatabase) -> Option<ModuleIdLt<'db>> {
         self.def_map(db).containing_module(self)
     }
 
-    pub fn is_block_module(self, db: &dyn DefDatabase) -> bool {
+    pub fn is_block_module(self, db: &'db dyn DefDatabase) -> bool {
         self.block(db).is_some() && self.def_map(db).root_module_id() == self
     }
 }
 
+impl ModuleId {
+    /// # Safety
+    ///
+    /// The caller must ensure that the `ModuleId` comes from the given database.
+    pub unsafe fn to_db<'db>(self, _db: &'db dyn DefDatabase) -> ModuleIdLt<'db> {
+        unsafe { std::mem::transmute(self) }
+    }
+}
+
 impl HasModule for ModuleId {
     #[inline]
     fn module(&self, _db: &dyn DefDatabase) -> ModuleId {