Infer the value of an index expression where the member key is keyof.

Author: xuhuanzy

crates/emmylua_code_analysis/src/compilation/test/member_infer_test.rs

@@ -102,4 +102,28 @@ mod test {
         assert_eq!(e_ty, LuaType::Integer);
         assert_eq!(f_ty, LuaType::Integer);
     }
+
+    #[test]
+    fn test_keyof() {
+        let mut ws = VirtualWorkspace::new();
+
+        ws.def(
+            r#"
+        ---@class SuiteHooks
+        ---@field beforeAll string
+        ---@field afterAll number
+
+        ---@type SuiteHooks
+        local hooks = {}
+
+        ---@type keyof SuiteHooks
+        local name = "beforeAll"
+
+        A = hooks[name]
+        "#,
+        );
+
+        let ty = ws.expr_ty("A");
+        assert_eq!(ws.humanize_type(ty), "(number|string)");
+    }
 }

crates/emmylua_code_analysis/src/diagnostic/checker/check_field.rs

@@ -388,7 +388,7 @@ fn get_key_types(db: &DbIndex, typ: &LuaType) -> HashSet<LuaType> {
                 type_set.insert(current_type);
             }
             LuaType::Call(alias_call) => {
-                if let Some(key_types) = get_key_of_keys(db, alias_call) {
+                if let Some(key_types) = get_keyof_keys(db, alias_call) {
                     for t in key_types {
                         stack.push(t.clone());
                     }
@@ -547,7 +547,7 @@ pub fn parse_require_expr_module_info<'a>(
         .find_module(&module_path)
 }
 
-fn get_key_of_keys(db: &DbIndex, alias_call: &LuaAliasCallType) -> Option<Vec<LuaType>> {
+fn get_keyof_keys(db: &DbIndex, alias_call: &LuaAliasCallType) -> Option<Vec<LuaType>> {
     if alias_call.get_call_kind() != LuaAliasCallKind::KeyOf {
         return None;
     }

crates/emmylua_code_analysis/src/semantic/infer/infer_index.rs

@@ -9,13 +9,13 @@ use rowan::TextRange;
 use smol_str::SmolStr;
 
 use crate::{
-    CacheEntry, GenericTpl, InFiled, InferGuardRef, LuaArrayLen, LuaArrayType, LuaDeclOrMemberId,
-    LuaInferCache, LuaInstanceType, LuaMemberOwner, LuaOperatorOwner, TypeOps,
+    CacheEntry, GenericTpl, InFiled, InferGuardRef, LuaAliasCallKind, LuaArrayLen, LuaArrayType,
+    LuaDeclOrMemberId, LuaInferCache, LuaInstanceType, LuaMemberOwner, LuaOperatorOwner, TypeOps,
     db_index::{
         DbIndex, LuaGenericType, LuaIntersectionType, LuaMemberKey, LuaObjectType,
         LuaOperatorMetaMethod, LuaTupleType, LuaType, LuaTypeDeclId, LuaUnionType,
     },
-    enum_variable_is_param, get_tpl_ref_extend_type,
+    enum_variable_is_param, get_keyof_members, get_tpl_ref_extend_type,
     semantic::{
         InferGuard,
         generic::{TypeSubstitutor, instantiate_type_generic},
@@ -420,26 +420,9 @@ fn infer_custom_type_member(
         return member_item.resolve_type(db);
     }
 
-    if type_decl.is_class()
-        && let Some(super_types) = type_index.get_super_types(&prefix_type_id)
-    {
-        for super_type in super_types {
-            let result =
-                infer_member_by_member_key(db, cache, &super_type, index_expr.clone(), infer_guard);
-
-            match result {
-                Ok(member_type) => {
-                    return Ok(member_type);
-                }
-                Err(InferFailReason::FieldNotFound) => {}
-                Err(err) => return Err(err),
-            }
-        }
-    }
-
     // 解决`key`为表达式的情况
     if let LuaIndexKey::Expr(expr) = index_key
-        && let Some(keys) = expr_to_member_key(db, cache, &expr)
+        && let Some(keys) = get_expr_member_key(db, cache, &expr)
     {
         let mut result_types = Vec::new();
         for key in keys {
@@ -462,6 +445,23 @@ fn infer_custom_type_member(
         }
     }
 
+    if type_decl.is_class()
+        && let Some(super_types) = type_index.get_super_types(&prefix_type_id)
+    {
+        for super_type in super_types {
+            let result =
+                infer_member_by_member_key(db, cache, &super_type, index_expr.clone(), infer_guard);
+
+            match result {
+                Ok(member_type) => {
+                    return Ok(member_type);
+                }
+                Err(InferFailReason::FieldNotFound) => {}
+                Err(err) => return Err(err),
+            }
+        }
+    }
+
     Err(InferFailReason::FieldNotFound)
 }
 
@@ -1233,46 +1233,76 @@ fn infer_namespace_member(
     ))
 }
 
-fn expr_to_member_key(
+fn get_expr_member_key(
     db: &DbIndex,
     cache: &mut LuaInferCache,
     expr: &LuaExpr,
-) -> Option<HashSet<LuaMemberKey>> {
+) -> Option<Vec<LuaMemberKey>> {
     let expr_type = infer_expr(db, cache, expr.clone()).ok()?;
     let mut keys: HashSet<LuaMemberKey> = HashSet::new();
     let mut stack = vec![expr_type.clone()];
     let mut visited = HashSet::new();
 
     while let Some(current_type) = stack.pop() {
-        if visited.contains(&current_type) {
+        if !visited.insert(current_type.clone()) {
             continue;
         }
-        visited.insert(current_type.clone());
         match &current_type {
             LuaType::StringConst(name) | LuaType::DocStringConst(name) => {
-                keys.insert(name.as_ref().to_string().into());
+                keys.insert(LuaMemberKey::Name((**name).clone()));
             }
             LuaType::IntegerConst(i) | LuaType::DocIntegerConst(i) => {
-                keys.insert((*i).into());
+                keys.insert(LuaMemberKey::Integer(*i));
+            }
+            LuaType::Call(alias_call) => {
+                if alias_call.get_call_kind() == LuaAliasCallKind::KeyOf {
+                    let operands = alias_call.get_operands();
+                    if operands.len() == 1 {
+                        if let Some(members) = get_keyof_members(db, &operands[0]) {
+                            keys.extend(members.into_iter().map(|member| member.key));
+                        }
+                    }
+                }
+            }
+            LuaType::MultiLineUnion(multi_union) => {
+                for (typ, _) in multi_union.get_unions() {
+                    if !visited.contains(typ) {
+                        stack.push(typ.clone());
+                    }
+                }
             }
             LuaType::Union(union_typ) => {
                 for t in union_typ.into_vec() {
-                    stack.push(t.clone());
+                    if !visited.contains(&t) {
+                        stack.push(t.clone());
+                    }
                 }
             }
             LuaType::TableConst(_) | LuaType::Tuple(_) => {
-                keys.insert(LuaMemberKey::ExprType(expr_type.clone()));
+                keys.insert(LuaMemberKey::ExprType(current_type.clone()));
             }
             LuaType::Ref(id) => {
-                if let Some(type_decl) = db.get_type_index().get_type_decl(id)
-                    && (type_decl.is_enum() || type_decl.is_alias())
-                {
-                    keys.insert(LuaMemberKey::ExprType(current_type.clone()));
+                if let Some(type_decl) = db.get_type_index().get_type_decl(id) {
+                    if type_decl.is_alias() {
+                        if let Some(origin_type) = type_decl.get_alias_origin(db, None) {
+                            if !visited.contains(&origin_type) {
+                                stack.push(origin_type);
+                            }
+                            continue;
+                        }
+                    }
+                    if type_decl.is_enum() || type_decl.is_alias() {
+                        keys.insert(LuaMemberKey::ExprType(current_type.clone()));
+                    }
                 }
             }
             _ => {}
         }
     }
+
+    // 转换为 Vec 并排序以确保顺序确定性
+    let mut keys: Vec<_> = keys.into_iter().collect();
+    keys.sort();
     Some(keys)
 }
 

crates/emmylua_code_analysis/src/semantic/type_check/complex_type/call_type_check.rs

@@ -32,10 +32,10 @@ pub fn check_call_type_compact(
                     }
 
                     let source_key_types = LuaType::Union(Arc::new(LuaUnionType::from_vec(
-                        get_key_of_keys(context, &source_operands[0]),
+                        get_keyof_keys(context, &source_operands[0]),
                     )));
                     let compact_key_types = LuaType::Union(Arc::new(LuaUnionType::from_vec(
-                        get_key_of_keys(context, &compact_operands[0]),
+                        get_keyof_keys(context, &compact_operands[0]),
                     )));
                     return check_general_type_compact(
                         context,
@@ -46,7 +46,7 @@ pub fn check_call_type_compact(
                 }
             }
             _ => {
-                let key_types = get_key_of_keys(context, &source_operands[0]);
+                let key_types = get_keyof_keys(context, &source_operands[0]);
                 for key_type in &key_types {
                     match check_general_type_compact(
                         context,
@@ -68,7 +68,7 @@ pub fn check_call_type_compact(
     Ok(())
 }
 
-fn get_key_of_keys(context: &TypeCheckContext, prefix_type: &LuaType) -> Vec<LuaType> {
+fn get_keyof_keys(context: &TypeCheckContext, prefix_type: &LuaType) -> Vec<LuaType> {
     let members = get_keyof_members(context.db, prefix_type).unwrap_or_default();
     let key_types = members
         .iter()