fix type_check simple_type enum

Author: xuhuanzy

crates/emmylua_code_analysis/src/db_index/type/type_decl.rs

@@ -4,9 +4,11 @@ use rowan::TextRange;
 use serde::{Deserialize, Deserializer, Serialize, Serializer};
 use smol_str::SmolStr;
 
-use crate::{instantiate_type_generic, DbIndex, FileId, TypeSubstitutor};
+use crate::{
+    instantiate_type_generic, DbIndex, FileId, LuaMemberKey, LuaMemberOwner, TypeSubstitutor,
+};
 
-use super::LuaType;
+use super::{LuaType, LuaUnionType};
 
 #[derive(Debug, Eq, PartialEq, Hash, Clone, Copy)]
 pub enum LuaDeclTypeKind {
@@ -173,6 +175,46 @@ impl LuaTypeDecl {
     pub fn merge_decl(&mut self, other: LuaTypeDecl) {
         self.locations.extend(other.locations);
     }
+
+    /// 获取枚举字段的类型
+    pub fn get_enum_field_type(&self, db: &DbIndex) -> Option<LuaType> {
+        if !self.is_enum() {
+            return None;
+        }
+
+        let enum_member_owner = LuaMemberOwner::Type(self.get_id());
+        let enum_members = db.get_member_index().get_members(&enum_member_owner)?;
+
+        let mut union_types = Vec::new();
+        if self.is_enum_key() {
+            for enum_member in enum_members {
+                let member_key = enum_member.get_key();
+                let fake_type = match member_key {
+                    LuaMemberKey::Name(name) => LuaType::DocStringConst(name.clone().into()),
+                    LuaMemberKey::Integer(i) => LuaType::IntegerConst(i.clone()),
+                    LuaMemberKey::None | LuaMemberKey::Expr(_) => continue,
+                };
+
+                union_types.push(fake_type);
+            }
+        } else {
+            for member in enum_members {
+                if let Some(type_cache) =
+                    db.get_type_index().get_type_cache(&member.get_id().into())
+                {
+                    let member_fake_type = match type_cache.as_type() {
+                        LuaType::StringConst(s) => LuaType::DocStringConst(s.clone().into()),
+                        LuaType::IntegerConst(i) => LuaType::DocIntegerConst(i.clone()),
+                        _ => type_cache.as_type().clone(),
+                    };
+
+                    union_types.push(member_fake_type);
+                }
+            }
+        }
+
+        return Some(LuaType::Union(LuaUnionType::new(union_types).into()));
+    }
 }
 
 #[derive(Debug, Eq, PartialEq, Hash, Clone)]

crates/emmylua_code_analysis/src/diagnostic/test/param_type_check_test.rs

@@ -649,6 +649,41 @@ mod test {
         ));
     }
 
+    #[test]
+    fn test_alias_union_enum_2() {
+        let mut ws = VirtualWorkspace::new();
+        assert!(!ws.check_code_for(
+            DiagnosticCode::ParamTypeNotMatch,
+            r#"
+                ---@alias EventType
+                ---| GlobalEventType
+                ---| UIEventType
+
+                ---@enum UIEventType
+                local UIEventType = {
+                    ['UI_CREATE'] = "ET_UI_PREFAB_CREATE_EVENT",
+                    ['UI_DELETE'] = "ET_UI_PREFAB_DEL_EVENT",
+                }
+
+                ---@enum GlobalEventType
+                local GlobalEventType = {
+                    ['GAME_INIT'] = 1,
+                    ['GAME_PAUSE'] = "ET_GAME_PAUSE",
+                }
+
+                ---@param event_name string
+                local function get_py_event_name(event_name)
+                end
+
+                ---@param a EventType
+                local function test(a)
+                    get_py_event_name(a)
+                end
+
+        "#
+        ));
+    }
+
     #[test]
     fn test_empty_class() {
         let mut ws = VirtualWorkspace::new();

crates/emmylua_code_analysis/src/semantic/type_check/simple_type.rs

@@ -1,7 +1,7 @@
-use crate::{semantic::type_check::is_sub_type_of, DbIndex, LuaType};
+use crate::{semantic::type_check::is_sub_type_of, DbIndex, LuaType, LuaTypeDeclId};
 
 use super::{
-    check_general_type_compact, ref_type::check_ref_type_compact, sub_type::get_base_type_id,
+    check_general_type_compact, sub_type::get_base_type_id,
     type_check_fail_reason::TypeCheckFailReason, type_check_guard::TypeCheckGuard, TypeCheckResult,
 };
 
@@ -263,14 +263,8 @@ fn check_base_type_for_ref_compact(
                 }
                 if let Some(decl) = db.get_type_index().get_type_decl(type_decl_id) {
                     if decl.is_enum() {
-                        // TODO: 优化, 不经过`check_ref_type_compact`
-                        if check_ref_type_compact(
-                            db,
-                            type_decl_id,
-                            compact_type,
-                            check_guard.next_level()?,
-                        )
-                        .is_ok()
+                        if check_enum_fields_match_source(db, source, type_decl_id, check_guard)
+                            .is_ok()
                         {
                             return Ok(());
                         }
@@ -282,3 +276,27 @@ fn check_base_type_for_ref_compact(
     }
     Err(TypeCheckFailReason::TypeNotMatch)
 }
+
+/// 检查`enum`的所有字段是否匹配`source`
+fn check_enum_fields_match_source(
+    db: &DbIndex,
+    source: &LuaType,
+    enum_type_decl_id: &LuaTypeDeclId,
+    check_guard: TypeCheckGuard,
+) -> TypeCheckResult {
+    if let Some(decl) = db.get_type_index().get_type_decl(enum_type_decl_id) {
+        if let Some(LuaType::Union(enum_fields)) = decl.get_enum_field_type(db) {
+            let is_match = enum_fields.get_types().iter().all(|field| {
+                let next_guard = check_guard.next_level();
+                if next_guard.is_err() {
+                    return false;
+                }
+                check_general_type_compact(db, source, field, next_guard.unwrap()).is_ok()
+            });
+            if is_match {
+                return Ok(());
+            }
+        }
+    }
+    Err(TypeCheckFailReason::TypeNotMatch)
+}