Merge pull request #20106 from Veykril/push-pytuxksnntux

Salsa idiomize `VariantFields` query

Author: Veykril

crates/hir-def/src/db.rs

@@ -29,7 +29,6 @@ use crate::{
     signatures::{
         ConstSignature, EnumSignature, FunctionSignature, ImplSignature, StaticSignature,
         StructSignature, TraitAliasSignature, TraitSignature, TypeAliasSignature, UnionSignature,
-        VariantFields,
     },
     tt,
     visibility::{self, Visibility},
@@ -113,17 +112,6 @@ pub trait DefDatabase: InternDatabase + ExpandDatabase + SourceDatabase {
 
     // region:data
 
-    #[salsa::invoke(VariantFields::query)]
-    fn variant_fields_with_source_map(
-        &self,
-        id: VariantId,
-    ) -> (Arc<VariantFields>, Arc<ExpressionStoreSourceMap>);
-
-    #[salsa::tracked]
-    fn variant_fields(&self, id: VariantId) -> Arc<VariantFields> {
-        self.variant_fields_with_source_map(id).0
-    }
-
     #[salsa::tracked]
     fn trait_signature(&self, trait_: TraitId) -> Arc<TraitSignature> {
         self.trait_signature_with_source_map(trait_).0

crates/hir-def/src/expr_store.rs

@@ -9,7 +9,10 @@ pub mod scope;
 #[cfg(test)]
 mod tests;
 
-use std::ops::{Deref, Index};
+use std::{
+    ops::{Deref, Index},
+    sync::LazyLock,
+};
 
 use cfg::{CfgExpr, CfgOptions};
 use either::Either;
@@ -19,6 +22,7 @@ use rustc_hash::FxHashMap;
 use smallvec::SmallVec;
 use span::{Edition, SyntaxContext};
 use syntax::{AstPtr, SyntaxNodePtr, ast};
+use triomphe::Arc;
 use tt::TextRange;
 
 use crate::{
@@ -220,6 +224,12 @@ impl ExpressionStoreBuilder {
 }
 
 impl ExpressionStore {
+    pub fn empty_singleton() -> Arc<Self> {
+        static EMPTY: LazyLock<Arc<ExpressionStore>> =
+            LazyLock::new(|| Arc::new(ExpressionStoreBuilder::default().finish()));
+        EMPTY.clone()
+    }
+
     /// Returns an iterator over all block expressions in this store that define inner items.
     pub fn blocks<'a>(
         &'a self,
@@ -636,6 +646,12 @@ impl Index<PathId> for ExpressionStore {
 // FIXME: Change `node_` prefix to something more reasonable.
 // Perhaps `expr_syntax` and `expr_id`?
 impl ExpressionStoreSourceMap {
+    pub fn empty_singleton() -> Arc<Self> {
+        static EMPTY: LazyLock<Arc<ExpressionStoreSourceMap>> =
+            LazyLock::new(|| Arc::new(ExpressionStoreSourceMap::default()));
+        EMPTY.clone()
+    }
+
     pub fn expr_or_pat_syntax(&self, id: ExprOrPatId) -> Result<ExprOrPatSource, SyntheticSyntax> {
         match id {
             ExprOrPatId::ExprId(id) => self.expr_syntax(id),

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

@@ -2250,7 +2250,7 @@ impl ExprCollector<'_> {
                         Some(ModuleDefId::ConstId(_)) => (None, Pat::Path(name.into())),
                         Some(ModuleDefId::EnumVariantId(variant))
                         // FIXME: This can cause a cycle if the user is writing invalid code
-                            if self.db.variant_fields(variant.into()).shape != FieldsShape::Record =>
+                            if variant.fields(self.db).shape != FieldsShape::Record =>
                         {
                             (None, Pat::Path(name.into()))
                         }

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

@@ -121,7 +121,7 @@ pub fn print_variant_body_hir(db: &dyn DefDatabase, owner: VariantId, edition: E
         VariantId::UnionId(it) => format!("union {}", item_name(db, it, "<missing>")),
     };
 
-    let fields = db.variant_fields(owner);
+    let fields = owner.fields(db);
 
     let mut p = Printer {
         db,

crates/hir-def/src/lib.rs

@@ -87,6 +87,7 @@ use crate::{
     attr::Attrs,
     builtin_type::BuiltinType,
     db::DefDatabase,
+    expr_store::ExpressionStoreSourceMap,
     hir::generics::{LocalLifetimeParamId, LocalTypeOrConstParamId},
     nameres::{
         LocalDefMap,
@@ -254,9 +255,35 @@ impl_intern!(FunctionId, FunctionLoc, intern_function, lookup_intern_function);
 type StructLoc = ItemLoc<ast::Struct>;
 impl_intern!(StructId, StructLoc, intern_struct, lookup_intern_struct);
 
+impl StructId {
+    pub fn fields(self, db: &dyn DefDatabase) -> &VariantFields {
+        VariantFields::firewall(db, self.into())
+    }
+
+    pub fn fields_with_source_map(
+        self,
+        db: &dyn DefDatabase,
+    ) -> (Arc<VariantFields>, Arc<ExpressionStoreSourceMap>) {
+        VariantFields::query(db, self.into())
+    }
+}
+
 pub type UnionLoc = ItemLoc<ast::Union>;
 impl_intern!(UnionId, UnionLoc, intern_union, lookup_intern_union);
 
+impl UnionId {
+    pub fn fields(self, db: &dyn DefDatabase) -> &VariantFields {
+        VariantFields::firewall(db, self.into())
+    }
+
+    pub fn fields_with_source_map(
+        self,
+        db: &dyn DefDatabase,
+    ) -> (Arc<VariantFields>, Arc<ExpressionStoreSourceMap>) {
+        VariantFields::query(db, self.into())
+    }
+}
+
 pub type EnumLoc = ItemLoc<ast::Enum>;
 impl_intern!(EnumId, EnumLoc, intern_enum, lookup_intern_enum);
 
@@ -337,6 +364,20 @@ pub struct EnumVariantLoc {
 }
 impl_intern!(EnumVariantId, EnumVariantLoc, intern_enum_variant, lookup_intern_enum_variant);
 impl_loc!(EnumVariantLoc, id: Variant, parent: EnumId);
+
+impl EnumVariantId {
+    pub fn fields(self, db: &dyn DefDatabase) -> &VariantFields {
+        VariantFields::firewall(db, self.into())
+    }
+
+    pub fn fields_with_source_map(
+        self,
+        db: &dyn DefDatabase,
+    ) -> (Arc<VariantFields>, Arc<ExpressionStoreSourceMap>) {
+        VariantFields::query(db, self.into())
+    }
+}
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct Macro2Loc {
     pub container: ModuleId,
@@ -1024,8 +1065,15 @@ pub enum VariantId {
 impl_from!(EnumVariantId, StructId, UnionId for VariantId);
 
 impl VariantId {
-    pub fn variant_data(self, db: &dyn DefDatabase) -> Arc<VariantFields> {
-        db.variant_fields(self)
+    pub fn fields(self, db: &dyn DefDatabase) -> &VariantFields {
+        VariantFields::firewall(db, self)
+    }
+
+    pub fn fields_with_source_map(
+        self,
+        db: &dyn DefDatabase,
+    ) -> (Arc<VariantFields>, Arc<ExpressionStoreSourceMap>) {
+        VariantFields::query(db, self)
     }
 
     pub fn file_id(self, db: &dyn DefDatabase) -> HirFileId {

crates/hir-def/src/signatures.rs

@@ -1,6 +1,6 @@
 //! Item signature IR definitions
 
-use std::ops::Not as _;
+use std::{cell::LazyCell, ops::Not as _};
 
 use bitflags::bitflags;
 use cfg::{CfgExpr, CfgOptions};
@@ -731,29 +731,26 @@ pub struct VariantFields {
     pub store: Arc<ExpressionStore>,
     pub shape: FieldsShape,
 }
+
+#[salsa::tracked]
 impl VariantFields {
-    #[inline]
+    #[salsa::tracked(returns(clone))]
     pub(crate) fn query(
         db: &dyn DefDatabase,
         id: VariantId,
     ) -> (Arc<Self>, Arc<ExpressionStoreSourceMap>) {
-        let (shape, (fields, store, source_map)) = match id {
+        let (shape, result) = match id {
             VariantId::EnumVariantId(id) => {
                 let loc = id.lookup(db);
                 let parent = loc.parent.lookup(db);
                 let source = loc.source(db);
                 let shape = adt_shape(source.value.kind());
-                let span_map = db.span_map(source.file_id);
-                let override_visibility = visibility_from_ast(
-                    db,
-                    source.value.parent_enum().visibility(),
-                    &mut |range| span_map.span_for_range(range).ctx,
-                );
+                let enum_vis = Some(source.value.parent_enum().visibility());
                 let fields = lower_field_list(
                     db,
                     parent.container,
                     source.map(|src| src.field_list()),
-                    Some(override_visibility),
+                    enum_vis,
                 );
                 (shape, fields)
             }
@@ -777,10 +774,29 @@ impl VariantFields {
                 (FieldsShape::Record, fields)
             }
         };
+        match result {
+            Some((fields, store, source_map)) => (
+                Arc::new(VariantFields { fields, store: Arc::new(store), shape }),
+                Arc::new(source_map),
+            ),
+            None => (
+                Arc::new(VariantFields {
+                    fields: Arena::default(),
+                    store: ExpressionStore::empty_singleton(),
+                    shape,
+                }),
+                ExpressionStoreSourceMap::empty_singleton(),
+            ),
+        }
+    }
 
-        (Arc::new(VariantFields { fields, store: Arc::new(store), shape }), Arc::new(source_map))
+    #[salsa::tracked(returns(deref))]
+    pub(crate) fn firewall(db: &dyn DefDatabase, id: VariantId) -> Arc<Self> {
+        Self::query(db, id).0
     }
+}
 
+impl VariantFields {
     pub fn len(&self) -> usize {
         self.fields.len()
     }
@@ -798,31 +814,24 @@ fn lower_field_list(
     db: &dyn DefDatabase,
     module: ModuleId,
     fields: InFile<Option<ast::FieldList>>,
-    override_visibility: Option<RawVisibility>,
-) -> (Arena<FieldData>, ExpressionStore, ExpressionStoreSourceMap) {
+    override_visibility: Option<Option<ast::Visibility>>,
+) -> Option<(Arena<FieldData>, ExpressionStore, ExpressionStoreSourceMap)> {
     let file_id = fields.file_id;
-    match fields.value {
-        Some(ast::FieldList::RecordFieldList(fields)) => lower_fields(
+    match fields.value? {
+        ast::FieldList::RecordFieldList(fields) => lower_fields(
             db,
             module,
             InFile::new(file_id, fields.fields().map(|field| (field.ty(), field))),
             |_, field| as_name_opt(field.name()),
             override_visibility,
         ),
-        Some(ast::FieldList::TupleFieldList(fields)) => lower_fields(
+        ast::FieldList::TupleFieldList(fields) => lower_fields(
             db,
             module,
             InFile::new(file_id, fields.fields().map(|field| (field.ty(), field))),
             |idx, _| Name::new_tuple_field(idx),
             override_visibility,
         ),
-        None => lower_fields(
-            db,
-            module,
-            InFile::new(file_id, std::iter::empty::<(Option<ast::Type>, ast::RecordField)>()),
-            |_, _| Name::missing(),
-            None,
-        ),
     }
 }
 
@@ -831,22 +840,34 @@ fn lower_fields<Field: ast::HasAttrs + ast::HasVisibility>(
     module: ModuleId,
     fields: InFile<impl Iterator<Item = (Option<ast::Type>, Field)>>,
     mut field_name: impl FnMut(usize, &Field) -> Name,
-    override_visibility: Option<RawVisibility>,
-) -> (Arena<FieldData>, ExpressionStore, ExpressionStoreSourceMap) {
-    let mut arena = Arena::new();
+    override_visibility: Option<Option<ast::Visibility>>,
+) -> Option<(Arena<FieldData>, ExpressionStore, ExpressionStoreSourceMap)> {
     let cfg_options = module.krate.cfg_options(db);
     let mut col = ExprCollector::new(db, module, fields.file_id);
+    let override_visibility = override_visibility.map(|vis| {
+        LazyCell::new(|| {
+            let span_map = db.span_map(fields.file_id);
+            visibility_from_ast(db, vis, &mut |range| span_map.span_for_range(range).ctx)
+        })
+    });
+
+    let mut arena = Arena::new();
     let mut idx = 0;
+    let mut has_fields = false;
     for (ty, field) in fields.value {
+        has_fields = true;
         match Attrs::is_cfg_enabled_for(db, &field, col.span_map(), cfg_options) {
             Ok(()) => {
                 let type_ref =
                     col.lower_type_ref_opt(ty, &mut ExprCollector::impl_trait_error_allocator);
-                let visibility = override_visibility.clone().unwrap_or_else(|| {
-                    visibility_from_ast(db, field.visibility(), &mut |range| {
-                        col.span_map().span_for_range(range).ctx
-                    })
-                });
+                let visibility = override_visibility.as_ref().map_or_else(
+                    || {
+                        visibility_from_ast(db, field.visibility(), &mut |range| {
+                            col.span_map().span_for_range(range).ctx
+                        })
+                    },
+                    |it| RawVisibility::clone(it),
+                );
                 let is_unsafe = field
                     .syntax()
                     .children_with_tokens()
@@ -867,9 +888,12 @@ fn lower_fields<Field: ast::HasAttrs + ast::HasVisibility>(
             }
         }
     }
+    if !has_fields {
+        return None;
+    }
     let store = col.store.finish();
     arena.shrink_to_fit();
-    (arena, store, col.source_map)
+    Some((arena, store, col.source_map))
 }
 
 #[derive(Debug, PartialEq, Eq)]
@@ -948,7 +972,7 @@ impl EnumVariants {
         self.variants.iter().all(|&(v, _, _)| {
             // The condition check order is slightly modified from rustc
             // to improve performance by early returning with relatively fast checks
-            let variant = &db.variant_fields(v.into());
+            let variant = v.fields(db);
             if !variant.fields().is_empty() {
                 return false;
             }

crates/hir-def/src/visibility.rs

@@ -273,7 +273,7 @@ pub(crate) fn field_visibilities_query(
     db: &dyn DefDatabase,
     variant_id: VariantId,
 ) -> Arc<ArenaMap<LocalFieldId, Visibility>> {
-    let variant_fields = db.variant_fields(variant_id);
+    let variant_fields = variant_id.fields(db);
     let fields = variant_fields.fields();
     if fields.is_empty() {
         return Arc::default();

crates/hir-ty/src/chalk_db.rs

@@ -801,7 +801,7 @@ pub(crate) fn adt_datum_query(
 
     // this slows down rust-analyzer by quite a bit unfortunately, so enabling this is currently not worth it
     let _variant_id_to_fields = |id: VariantId| {
-        let variant_data = &id.variant_data(db);
+        let variant_data = &id.fields(db);
         let fields = if variant_data.fields().is_empty() {
             vec![]
         } else {

crates/hir-ty/src/diagnostics/decl_check.rs

@@ -307,7 +307,7 @@ impl<'a> DeclValidator<'a> {
 
     /// Check incorrect names for struct fields.
     fn validate_struct_fields(&mut self, struct_id: StructId) {
-        let data = self.db.variant_fields(struct_id.into());
+        let data = struct_id.fields(self.db);
         if data.shape != FieldsShape::Record {
             return;
         };
@@ -468,7 +468,7 @@ impl<'a> DeclValidator<'a> {
 
     /// Check incorrect names for fields of enum variant.
     fn validate_enum_variant_fields(&mut self, variant_id: EnumVariantId) {
-        let variant_data = self.db.variant_fields(variant_id.into());
+        let variant_data = variant_id.fields(self.db);
         if variant_data.shape != FieldsShape::Record {
             return;
         };