Merge pull request #21900 from Veykril/push-ylmpryvvpsym

Fully implement `VariantFields `expression support

Author: Veykril

crates/hir-def/src/db.rs

@@ -4,16 +4,15 @@ use hir_expand::{
     EditionedFileId, HirFileId, InFile, Lookup, MacroCallId, MacroDefId, MacroDefKind,
     db::ExpandDatabase,
 };
-use la_arena::ArenaMap;
 use triomphe::Arc;
 
 use crate::{
     AnonConstId, AnonConstLoc, AssocItemId, AttrDefId, BlockId, BlockLoc, ConstId, ConstLoc,
     EnumId, EnumLoc, EnumVariantId, EnumVariantLoc, ExternBlockId, ExternBlockLoc, ExternCrateId,
-    ExternCrateLoc, FunctionId, FunctionLoc, 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,
+    ExternCrateLoc, FunctionId, FunctionLoc, ImplId, ImplLoc, Macro2Id, Macro2Loc, MacroExpander,
+    MacroId, MacroRulesId, MacroRulesLoc, MacroRulesLocFlags, ProcMacroId, ProcMacroLoc, StaticId,
+    StaticLoc, StructId, StructLoc, TraitId, TraitLoc, TypeAliasId, TypeAliasLoc, UnionId,
+    UnionLoc, UseId, UseLoc,
     attrs::AttrFlags,
     item_tree::{ItemTree, file_item_tree_query},
     nameres::crate_def_map,
@@ -98,9 +97,6 @@ pub trait DefDatabase: InternDatabase + ExpandDatabase + SourceDatabase {
 
     // region:visibilities
 
-    #[salsa::invoke(visibility::field_visibilities_query)]
-    fn field_visibilities(&self, var: VariantId) -> Arc<ArenaMap<LocalFieldId, Visibility>>;
-
     #[salsa::invoke(visibility::assoc_visibility_query)]
     fn assoc_visibility(&self, def: AssocItemId) -> Visibility;
 

crates/hir-def/src/expr_store.rs

@@ -31,6 +31,7 @@ use crate::{
         PatId, RecordFieldPat, RecordSpread, Statement,
     },
     nameres::{DefMap, block_def_map},
+    signatures::VariantFields,
     type_ref::{LifetimeRef, LifetimeRefId, PathId, TypeRef, TypeRefId},
 };
 
@@ -95,14 +96,16 @@ pub type LifetimeSource = InFile<LifetimePtr>;
 /// Used by signature/body inference to determine the expected type for each
 /// const expression root.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-pub enum ConstExprOrigin {
+pub enum RootExprOrigin {
     /// Array length expression: `[T; <expr>]` — expected type is `usize`.
     ArrayLength,
     /// Const parameter default value: `const N: usize = <expr>`.
     ConstParam(crate::hir::generics::LocalTypeOrConstParamId),
     /// Const generic argument in a path: `SomeType::<{ <expr> }>` or `some_fn::<{ <expr> }>()`.
     /// Determining the expected type requires path resolution, so it is deferred.
     GenericArgsPath,
+    /// The root expression of a body.
+    BodyRoot,
 }
 
 // We split the store into types-only and expressions, because most stores (e.g. generics)
@@ -125,11 +128,8 @@ struct ExpressionOnlyStore {
     /// to variables and have hygiene (some refer to items, we don't know at this stage).
     ident_hygiene: FxHashMap<ExprOrPatId, HygieneId>,
 
-    /// Maps const expression roots to their origin.
-    ///
-    /// Populated during lowering. Used by signature inference to determine expected types,
-    /// and by `signature_const_expr_roots()` to enumerate roots for scope computation.
-    const_expr_origins: ThinVec<(ExprId, ConstExprOrigin)>,
+    /// Maps expression roots to their origin.
+    expr_roots: SmallVec<[(ExprId, RootExprOrigin); 1]>,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -243,7 +243,7 @@ pub struct ExpressionStoreBuilder {
     pub types: Arena<TypeRef>,
     block_scopes: Vec<BlockId>,
     ident_hygiene: FxHashMap<ExprOrPatId, HygieneId>,
-    pub const_expr_origins: Option<ThinVec<(ExprId, ConstExprOrigin)>>,
+    pub inference_roots: Option<SmallVec<[(ExprId, RootExprOrigin); 1]>>,
 
     // AST expressions can create patterns in destructuring assignments. Therefore, `ExprSource` can also map
     // to `PatId`, and `PatId` can also map to `ExprSource` (the other way around is unaffected).
@@ -315,7 +315,7 @@ impl ExpressionStoreBuilder {
             mut bindings,
             mut binding_owners,
             mut ident_hygiene,
-            mut const_expr_origins,
+            inference_roots: mut expr_roots,
             mut types,
             mut lifetimes,
 
@@ -375,7 +375,7 @@ impl ExpressionStoreBuilder {
 
         let store = {
             let expr_only = if has_exprs {
-                if let Some(const_expr_origins) = &mut const_expr_origins {
+                if let Some(const_expr_origins) = &mut expr_roots {
                     const_expr_origins.shrink_to_fit();
                 }
                 Some(Box::new(ExpressionOnlyStore {
@@ -386,7 +386,7 @@ impl ExpressionStoreBuilder {
                     binding_owners,
                     block_scopes: block_scopes.into_boxed_slice(),
                     ident_hygiene,
-                    const_expr_origins: const_expr_origins.unwrap_or_default(),
+                    expr_roots: expr_roots.unwrap_or_default(),
                 }))
             } else {
                 None
@@ -448,6 +448,9 @@ impl ExpressionStore {
                 }
             }
             ExpressionStoreOwnerId::Body(body) => &Body::of(db, body).store,
+            ExpressionStoreOwnerId::VariantFields(variant_id) => {
+                &VariantFields::of(db, variant_id).store
+            }
         }
     }
 
@@ -505,30 +508,27 @@ impl ExpressionStore {
                 let (store, sm) = Body::with_source_map(db, body);
                 (&store.store, &sm.store)
             }
+            ExpressionStoreOwnerId::VariantFields(variant_id) => {
+                let (store, sm) = VariantFields::with_source_map(db, variant_id);
+                (&store.store, sm)
+            }
         }
     }
 
-    /// Returns all const expression root `ExprId`s found in this store.
-    ///
-    /// Used to compute expression scopes for signature stores.
-    pub fn signature_const_expr_roots(&self) -> impl Iterator<Item = ExprId> {
+    /// Returns all expression root `ExprId`s found in this store.
+    pub fn expr_roots(&self) -> impl Iterator<Item = ExprId> {
         self.const_expr_origins().iter().map(|&(id, _)| id)
     }
 
     /// Like [`Self::signature_const_expr_roots`], but also returns the origin
-    /// of each const expression.
-    ///
-    /// This is used by signature inference to determine the expected type for
-    /// each root expression.
-    pub fn signature_const_expr_roots_with_origins(
-        &self,
-    ) -> impl Iterator<Item = (ExprId, ConstExprOrigin)> {
+    /// of each expression.
+    pub fn expr_roots_with_origins(&self) -> impl Iterator<Item = (ExprId, RootExprOrigin)> {
         self.const_expr_origins().iter().map(|&(id, origin)| (id, origin))
     }
 
     /// Returns the map of const expression roots to their origins.
-    pub fn const_expr_origins(&self) -> &[(ExprId, ConstExprOrigin)] {
-        self.expr_only.as_ref().map_or(&[], |it| &it.const_expr_origins)
+    pub fn const_expr_origins(&self) -> &[(ExprId, RootExprOrigin)] {
+        self.expr_only.as_ref().map_or(&[], |it| &it.expr_roots)
     }
 
     /// Returns an iterator over all block expressions in this store that define inner items.

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

@@ -29,8 +29,6 @@ pub struct Body {
     /// empty.
     pub params: Box<[PatId]>,
     pub self_param: Option<BindingId>,
-    /// The `ExprId` of the actual body expression.
-    pub body_expr: ExprId,
 }
 
 impl ops::Deref for Body {
@@ -115,6 +113,10 @@ impl Body {
 }
 
 impl Body {
+    pub fn root_expr(&self) -> ExprId {
+        self.store.expr_roots().next().unwrap()
+    }
+
     pub fn pretty_print(
         &self,
         db: &dyn DefDatabase,

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

@@ -18,6 +18,7 @@ use hir_expand::{
 };
 use intern::{Symbol, sym};
 use rustc_hash::FxHashMap;
+use smallvec::smallvec;
 use stdx::never;
 use syntax::{
     AstNode, AstPtr, SyntaxNodePtr,
@@ -36,9 +37,9 @@ use crate::{
     attrs::AttrFlags,
     db::DefDatabase,
     expr_store::{
-        Body, BodySourceMap, ConstExprOrigin, ExprPtr, ExpressionStore, ExpressionStoreBuilder,
+        Body, BodySourceMap, ExprPtr, ExpressionStore, ExpressionStoreBuilder,
         ExpressionStoreDiagnostics, ExpressionStoreSourceMap, HygieneId, LabelPtr, LifetimePtr,
-        PatPtr, TypePtr,
+        PatPtr, RootExprOrigin, TypePtr,
         expander::Expander,
         lower::generics::ImplTraitLowerFn,
         path::{AssociatedTypeBinding, GenericArg, GenericArgs, GenericArgsParentheses, Path},
@@ -117,9 +118,10 @@ pub(super) fn lower_body(
             params = (0..count).map(|_| collector.missing_pat()).collect();
         };
         let body_expr = collector.missing_expr();
+        collector.store.inference_roots = Some(smallvec![(body_expr, RootExprOrigin::BodyRoot)]);
         let (store, source_map) = collector.store.finish();
         return (
-            Body { store, params: params.into_boxed_slice(), self_param, body_expr },
+            Body { store, params: params.into_boxed_slice(), self_param },
             BodySourceMap { self_param: source_map_self_param, store: source_map },
         );
     }
@@ -173,10 +175,11 @@ pub(super) fn lower_body(
             }
         },
     );
+    collector.store.inference_roots = Some(smallvec![(body_expr, RootExprOrigin::BodyRoot)]);
 
     let (store, source_map) = collector.store.finish();
     (
-        Body { store, params: params.into_boxed_slice(), self_param, body_expr },
+        Body { store, params: params.into_boxed_slice(), self_param },
         BodySourceMap { self_param: source_map_self_param, store: source_map },
     )
 }
@@ -535,7 +538,7 @@ impl<'db> ExprCollector<'db> {
         current_file_id: HirFileId,
     ) -> ExprCollector<'_> {
         let mut this = Self::body(db, module, current_file_id);
-        this.store.const_expr_origins = Some(Default::default());
+        this.store.inference_roots = Some(Default::default());
         this
     }
 
@@ -635,8 +638,8 @@ impl<'db> ExprCollector<'db> {
             }
             ast::Type::ArrayType(inner) => {
                 let len = self.lower_const_arg_opt(inner.const_arg());
-                if let Some(const_expr_origins) = &mut self.store.const_expr_origins {
-                    const_expr_origins.push((len.expr, ConstExprOrigin::ArrayLength));
+                if let Some(const_expr_origins) = &mut self.store.inference_roots {
+                    const_expr_origins.push((len.expr, RootExprOrigin::ArrayLength));
                 }
                 TypeRef::Array(ArrayType {
                     ty: self.lower_type_ref_opt(inner.ty(), impl_trait_lower_fn),
@@ -922,8 +925,8 @@ impl<'db> ExprCollector<'db> {
                 }
                 ast::GenericArg::ConstArg(arg) => {
                     let arg = self.lower_const_arg(arg);
-                    if let Some(const_expr_origins) = &mut self.store.const_expr_origins {
-                        const_expr_origins.push((arg.expr, ConstExprOrigin::GenericArgsPath));
+                    if let Some(const_expr_origins) = &mut self.store.inference_roots {
+                        const_expr_origins.push((arg.expr, RootExprOrigin::GenericArgsPath));
                     }
                     args.push(GenericArg::Const(arg))
                 }
@@ -1065,16 +1068,16 @@ impl<'db> ExprCollector<'db> {
     }
 
     fn lower_const_arg_opt(&mut self, arg: Option<ast::ConstArg>) -> ConstRef {
-        let const_expr_origins = self.store.const_expr_origins.take();
+        let const_expr_origins = self.store.inference_roots.take();
         let r = ConstRef { expr: self.collect_expr_opt(arg.and_then(|it| it.expr())) };
-        self.store.const_expr_origins = const_expr_origins;
+        self.store.inference_roots = const_expr_origins;
         r
     }
 
     pub fn lower_const_arg(&mut self, arg: ast::ConstArg) -> ConstRef {
-        let const_expr_origins = self.store.const_expr_origins.take();
+        let const_expr_origins = self.store.inference_roots.take();
         let r = ConstRef { expr: self.collect_expr_opt(arg.expr()) };
-        self.store.const_expr_origins = const_expr_origins;
+        self.store.inference_roots = const_expr_origins;
         r
     }
 

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

@@ -130,11 +130,11 @@ impl GenericParamsCollector {
                     let param = ConstParamData { name, ty, default };
                     let idx = self.type_or_consts.alloc(param.into());
                     if let Some(default) = default
-                        && let Some(const_expr_origins) = &mut ec.store.const_expr_origins
+                        && let Some(const_expr_origins) = &mut ec.store.inference_roots
                     {
                         const_expr_origins.push((
                             default.expr,
-                            crate::expr_store::ConstExprOrigin::ConstParam(idx),
+                            crate::expr_store::RootExprOrigin::ConstParam(idx),
                         ));
                     }
                 }

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

@@ -105,7 +105,7 @@ pub fn print_body_hir(
         p.buf.push(')');
         p.buf.push(' ');
     }
-    p.print_expr(body.body_expr);
+    p.print_expr(body.root_expr());
     if matches!(owner, DefWithBodyId::StaticId(_) | DefWithBodyId::ConstId(_)) {
         p.buf.push(';');
     }

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

@@ -3,13 +3,14 @@ use hir_expand::{MacroDefId, name::Name};
 use la_arena::{Arena, ArenaMap, Idx, IdxRange, RawIdx};
 
 use crate::{
-    BlockId, DefWithBodyId, ExpressionStoreOwnerId, GenericDefId,
+    BlockId, DefWithBodyId, ExpressionStoreOwnerId, GenericDefId, VariantId,
     db::DefDatabase,
     expr_store::{Body, ExpressionStore, HygieneId},
     hir::{
         Binding, BindingId, Expr, ExprId, Item, LabelId, Pat, PatId, Statement,
         generics::GenericParams,
     },
+    signatures::VariantFields,
 };
 
 pub type ScopeId = Idx<ScopeData>;
@@ -65,11 +66,20 @@ impl ExprScopes {
     #[salsa::tracked(returns(ref))]
     pub fn sig_expr_scopes(db: &dyn DefDatabase, def: GenericDefId) -> ExprScopes {
         let (_, store) = GenericParams::with_store(db, def);
-        let roots = store.signature_const_expr_roots();
+        let roots = store.expr_roots();
         let mut scopes = ExprScopes::new_store(store, roots);
         scopes.shrink_to_fit();
         scopes
     }
+
+    #[salsa::tracked(returns(ref))]
+    pub fn variant_scopes(db: &dyn DefDatabase, def: VariantId) -> ExprScopes {
+        let fields = VariantFields::of(db, def);
+        let roots = fields.store.expr_roots();
+        let mut scopes = ExprScopes::new_store(&fields.store, roots);
+        scopes.shrink_to_fit();
+        scopes
+    }
 }
 
 impl ExprScopes {
@@ -78,6 +88,9 @@ impl ExprScopes {
         match def.into() {
             ExpressionStoreOwnerId::Body(def) => Self::body_expr_scopes(db, def),
             ExpressionStoreOwnerId::Signature(def) => Self::sig_expr_scopes(db, def),
+            ExpressionStoreOwnerId::VariantFields(variant_id) => {
+                Self::variant_scopes(db, variant_id)
+            }
         }
     }
 
@@ -138,7 +151,7 @@ impl ExprScopes {
             scopes.add_bindings(body, root, self_param, body.binding_hygiene(self_param));
         }
         scopes.add_params_bindings(body, root, &body.params);
-        compute_expr_scopes(body.body_expr, body, &mut scopes, &mut root);
+        compute_expr_scopes(body.root_expr(), body, &mut scopes, &mut root);
         scopes
     }