Idiomatic salsa use for enum variants query

Author: Veykril

Cargo.toml

@@ -134,7 +134,9 @@ pulldown-cmark-to-cmark = "10.0.4"
 pulldown-cmark = { version = "0.9.6", default-features = false }
 rayon = "1.10.0"
 rowan = "=0.15.15"
-salsa = { version = "0.22.0", default-features = false, features = ["rayon","salsa_unstable"] }
+# Ideally we'd not enable the macros feature but unfortunately the `tracked` attribute does not work
+# on impls without it
+salsa = { version = "0.22.0", default-features = true, features = ["rayon","salsa_unstable", "macros"] }
 salsa-macros = "0.22.0"
 semver = "1.0.26"
 serde = { version = "1.0.219" }

crates/hir-def/src/db.rs

@@ -8,7 +8,6 @@ use hir_expand::{
 use intern::{Symbol, sym};
 use la_arena::ArenaMap;
 use syntax::{AstPtr, ast};
-use thin_vec::ThinVec;
 use triomphe::Arc;
 
 use crate::{
@@ -32,9 +31,9 @@ use crate::{
         diagnostics::DefDiagnostics,
     },
     signatures::{
-        ConstSignature, EnumSignature, EnumVariants, FunctionSignature, ImplSignature,
-        InactiveEnumVariantCode, StaticSignature, StructSignature, TraitAliasSignature,
-        TraitSignature, TypeAliasSignature, UnionSignature, VariantFields,
+        ConstSignature, EnumSignature, FunctionSignature, ImplSignature, StaticSignature,
+        StructSignature, TraitAliasSignature, TraitSignature, TypeAliasSignature, UnionSignature,
+        VariantFields,
     },
     tt,
     visibility::{self, Visibility},
@@ -121,19 +120,6 @@ pub trait DefDatabase: InternDatabase + ExpandDatabase + SourceDatabase {
         id: VariantId,
     ) -> (Arc<VariantFields>, Arc<ExpressionStoreSourceMap>);
 
-    // FIXME: Should we make this transparent? The only unstable thing in `enum_variants_with_diagnostics()`
-    // is ast ids, and ast ids are pretty stable now.
-    #[salsa::tracked]
-    fn enum_variants(&self, id: EnumId) -> Arc<EnumVariants> {
-        self.enum_variants_with_diagnostics(id).0
-    }
-
-    #[salsa::invoke(EnumVariants::enum_variants_query)]
-    fn enum_variants_with_diagnostics(
-        &self,
-        id: EnumId,
-    ) -> (Arc<EnumVariants>, Option<Arc<ThinVec<InactiveEnumVariantCode>>>);
-
     #[salsa::transparent]
     #[salsa::invoke(ImplItems::impl_items_query)]
     fn impl_items(&self, e: ImplId) -> Arc<ImplItems>;

crates/hir-def/src/find_path.rs

@@ -137,7 +137,7 @@ fn find_path_inner(ctx: &FindPathCtx<'_>, item: ItemInNs, max_len: usize) -> Opt
         let loc = variant.lookup(ctx.db);
         if let Some(mut path) = find_path_inner(ctx, ItemInNs::Types(loc.parent.into()), max_len) {
             path.push_segment(
-                ctx.db.enum_variants(loc.parent).variants[loc.index as usize].1.clone(),
+                loc.parent.enum_variants(ctx.db).variants[loc.index as usize].1.clone(),
             );
             return Some(path);
         }

crates/hir-def/src/lang_item.rs

@@ -125,7 +125,7 @@ pub fn crate_lang_items(db: &dyn DefDatabase, krate: Crate) -> Option<Box<LangIt
                 }
                 ModuleDefId::AdtId(AdtId::EnumId(e)) => {
                     lang_items.collect_lang_item(db, e, LangItemTarget::EnumId);
-                    db.enum_variants(e).variants.iter().for_each(|&(id, _, _)| {
+                    e.enum_variants(db).variants.iter().for_each(|&(id, _, _)| {
                         lang_items.collect_lang_item(db, id, LangItemTarget::EnumVariant);
                     });
                 }

crates/hir-def/src/lib.rs

@@ -51,6 +51,7 @@ pub mod visibility;
 
 use intern::{Interned, sym};
 pub use rustc_abi as layout;
+use thin_vec::ThinVec;
 use triomphe::Arc;
 
 pub use crate::signatures::LocalFieldId;
@@ -88,7 +89,7 @@ use crate::{
     db::DefDatabase,
     hir::generics::{LocalLifetimeParamId, LocalTypeOrConstParamId},
     nameres::{LocalDefMap, block_def_map, crate_def_map, crate_local_def_map},
-    signatures::VariantFields,
+    signatures::{EnumVariants, InactiveEnumVariantCode, VariantFields},
 };
 
 type FxIndexMap<K, V> = indexmap::IndexMap<K, V, rustc_hash::FxBuildHasher>;
@@ -253,6 +254,21 @@ impl_intern!(UnionId, UnionLoc, intern_union, lookup_intern_union);
 pub type EnumLoc = ItemLoc<ast::Enum>;
 impl_intern!(EnumId, EnumLoc, intern_enum, lookup_intern_enum);
 
+impl EnumId {
+    #[inline]
+    pub fn enum_variants(self, db: &dyn DefDatabase) -> &EnumVariants {
+        &self.enum_variants_with_diagnostics(db).0
+    }
+
+    #[inline]
+    pub fn enum_variants_with_diagnostics(
+        self,
+        db: &dyn DefDatabase,
+    ) -> &(EnumVariants, Option<ThinVec<InactiveEnumVariantCode>>) {
+        EnumVariants::of(db, self)
+    }
+}
+
 type ConstLoc = AssocItemLoc<ast::Const>;
 impl_intern!(ConstId, ConstLoc, intern_const, lookup_intern_const);
 

crates/hir-def/src/nameres/collector.rs

@@ -995,9 +995,8 @@ impl<'db> DefCollector<'db> {
                     Some(ModuleDefId::AdtId(AdtId::EnumId(e))) => {
                         cov_mark::hit!(glob_enum);
                         // glob import from enum => just import all the variants
-                        let resolutions = self
-                            .db
-                            .enum_variants(e)
+                        let resolutions = e
+                            .enum_variants(self.db)
                             .variants
                             .iter()
                             .map(|&(variant, ref name, _)| {

crates/hir-def/src/nameres/path_resolution.rs

@@ -528,8 +528,8 @@ impl DefMap {
                     // enum variant
                     cov_mark::hit!(can_import_enum_variant);
 
-                    let res = db
-                        .enum_variants(e)
+                    let res = e
+                        .enum_variants(db)
                         .variants
                         .iter()
                         .find(|(_, name, _)| name == segment)

crates/hir-def/src/nameres/tests/incremental.rs

@@ -172,8 +172,7 @@ fn no() {}
                 "ast_id_map_shim",
                 "parse_shim",
                 "real_span_map_shim",
-                "enum_variants_shim",
-                "enum_variants_with_diagnostics_shim",
+                "of_",
             ]
         "#]],
         expect![[r#"
@@ -182,7 +181,7 @@ fn no() {}
                 "ast_id_map_shim",
                 "file_item_tree_query",
                 "real_span_map_shim",
-                "enum_variants_with_diagnostics_shim",
+                "of_",
             ]
         "#]],
     );

crates/hir-def/src/signatures.rs

@@ -884,11 +884,13 @@ pub struct EnumVariants {
     pub variants: Box<[(EnumVariantId, Name, FieldsShape)]>,
 }
 
+#[salsa::tracked]
 impl EnumVariants {
-    pub(crate) fn enum_variants_query(
+    #[salsa::tracked(returns(ref))]
+    pub(crate) fn of(
         db: &dyn DefDatabase,
         e: EnumId,
-    ) -> (Arc<EnumVariants>, Option<Arc<ThinVec<InactiveEnumVariantCode>>>) {
+    ) -> (EnumVariants, Option<ThinVec<InactiveEnumVariantCode>>) {
         let loc = e.lookup(db);
         let source = loc.source(db);
         let ast_id_map = db.ast_id_map(source.file_id);
@@ -898,7 +900,7 @@ impl EnumVariants {
         let cfg_options = loc.container.krate.cfg_options(db);
         let mut index = 0;
         let Some(variants) = source.value.variant_list() else {
-            return (Arc::new(EnumVariants { variants: Box::default() }), None);
+            return (EnumVariants { variants: Box::default() }, None);
         };
         let variants = variants
             .variants()
@@ -926,12 +928,11 @@ impl EnumVariants {
             })
             .collect();
 
-        (
-            Arc::new(EnumVariants { variants }),
-            diagnostics.is_empty().not().then(|| Arc::new(diagnostics)),
-        )
+        (EnumVariants { variants }, diagnostics.is_empty().not().then_some(diagnostics))
     }
+}
 
+impl EnumVariants {
     pub fn variant(&self, name: &Name) -> Option<EnumVariantId> {
         self.variants.iter().find_map(|(v, n, _)| if n == name { Some(*v) } else { None })
     }

crates/hir-ty/src/chalk_db.rs

@@ -823,8 +823,8 @@ pub(crate) fn adt_datum_query(
             (rust_ir::AdtKind::Struct, vec![variant_id_to_fields(id.into())])
         }
         hir_def::AdtId::EnumId(id) => {
-            let variants = db
-                .enum_variants(id)
+            let variants = id
+                .enum_variants(db)
                 .variants
                 .iter()
                 .map(|&(variant_id, _, _)| variant_id_to_fields(variant_id.into()))