fix remove_type

Author: xuhuanzy

crates/emmylua_code_analysis/src/db_index/type/type_ops/remove_type.rs

@@ -9,55 +9,55 @@ pub fn remove_type(db: &DbIndex, source: LuaType, removed_type: LuaType) -> Opti
         }
     }
 
-    let source = get_real_type(db, &source).unwrap_or(&source);
+    let real_type = get_real_type(db, &source).unwrap_or(&source);
 
     match &removed_type {
         LuaType::Nil => {
-            if source.is_nil() {
+            if real_type.is_nil() {
                 return None;
             }
         }
         LuaType::Boolean => {
-            if source.is_boolean() {
+            if real_type.is_boolean() {
                 return None;
             }
         }
         LuaType::Integer => {
-            if source.is_integer() {
+            if real_type.is_integer() {
                 return None;
             }
         }
         LuaType::Number => {
-            if source.is_number() {
+            if real_type.is_number() {
                 return None;
             }
         }
         LuaType::String => {
-            if source.is_string() {
+            if real_type.is_string() {
                 return None;
             }
         }
         LuaType::Io => {
-            if source.is_io() {
+            if real_type.is_io() {
                 return None;
             }
         }
         LuaType::Function => {
-            if source.is_function() {
+            if real_type.is_function() {
                 return None;
             }
         }
         LuaType::Thread => {
-            if source.is_thread() {
+            if real_type.is_thread() {
                 return None;
             }
         }
         LuaType::Userdata => {
-            if source.is_userdata() {
+            if real_type.is_userdata() {
                 return None;
             }
         }
-        LuaType::Table => match &source {
+        LuaType::Table => match &real_type {
             LuaType::TableConst(_)
             | LuaType::Table
             | LuaType::Userdata
@@ -74,7 +74,7 @@ pub fn remove_type(db: &DbIndex, source: LuaType, removed_type: LuaType) -> Opti
                     return Some(source.clone());
                 }
                 if type_decl.is_alias() {
-                    if let Some(alias_ref) = get_real_type(db, &source) {
+                    if let Some(alias_ref) = get_real_type(db, &real_type) {
                         return remove_type(db, alias_ref.clone(), removed_type);
                     }
                 }
@@ -91,7 +91,7 @@ pub fn remove_type(db: &DbIndex, source: LuaType, removed_type: LuaType) -> Opti
             }
             _ => {}
         },
-        LuaType::DocStringConst(s) | LuaType::StringConst(s) => match &source {
+        LuaType::DocStringConst(s) | LuaType::StringConst(s) => match &real_type {
             LuaType::DocStringConst(s2) => {
                 if s == s2 {
                     return None;
@@ -104,7 +104,7 @@ pub fn remove_type(db: &DbIndex, source: LuaType, removed_type: LuaType) -> Opti
             }
             _ => {}
         },
-        LuaType::DocIntegerConst(i) | LuaType::IntegerConst(i) => match &source {
+        LuaType::DocIntegerConst(i) | LuaType::IntegerConst(i) => match &real_type {
             LuaType::DocIntegerConst(i2) => {
                 if i == i2 {
                     return None;
@@ -117,7 +117,7 @@ pub fn remove_type(db: &DbIndex, source: LuaType, removed_type: LuaType) -> Opti
             }
             _ => {}
         },
-        LuaType::DocBooleanConst(b) | LuaType::BooleanConst(b) => match &source {
+        LuaType::DocBooleanConst(b) | LuaType::BooleanConst(b) => match &real_type {
             LuaType::DocBooleanConst(b2) => {
                 if b == b2 {
                     return None;
@@ -133,7 +133,7 @@ pub fn remove_type(db: &DbIndex, source: LuaType, removed_type: LuaType) -> Opti
         _ => {}
     }
 
-    if let LuaType::Union(u) = &source {
+    if let LuaType::Union(u) = &real_type {
         let types = u
             .into_vec()
             .iter()
@@ -144,7 +144,7 @@ pub fn remove_type(db: &DbIndex, source: LuaType, removed_type: LuaType) -> Opti
         let types = u
             .into_vec()
             .iter()
-            .filter_map(|t| remove_type(db, source.clone(), t.clone()))
+            .filter_map(|t| remove_type(db, real_type.clone(), t.clone()))
             .collect::<Vec<_>>();
         return Some(LuaType::from_vec(types));
     }

crates/emmylua_code_analysis/src/db_index/type/type_ops/union_type.rs

@@ -48,6 +48,28 @@ pub fn union_type(source: LuaType, target: LuaType) -> LuaType {
                 LuaType::from_vec(vec![source.clone(), target.clone()])
             }
         }
+        (LuaType::MultiLineUnion(left), right) => {
+            let include = match right {
+                LuaType::StringConst(v) => {
+                    left.get_unions().iter().any(|(t, _)| match (t, right) {
+                        (LuaType::DocStringConst(a), _) => a == v,
+                        _ => false,
+                    })
+                }
+                LuaType::IntegerConst(v) => {
+                    left.get_unions().iter().any(|(t, _)| match (t, right) {
+                        (LuaType::DocIntegerConst(a), _) => a == v,
+                        _ => false,
+                    })
+                }
+                _ => false,
+            };
+
+            if include {
+                return source;
+            }
+            LuaType::from_vec(vec![source, target])
+        }
         // union
         (LuaType::Union(left), right) if !right.is_union() => {
             let left = left.deref().clone();

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

@@ -427,6 +427,10 @@ impl LuaType {
         matches!(self, LuaType::TypeGuard(_))
     }
 
+    pub fn is_multi_line_union(&self) -> bool {
+        matches!(self, LuaType::MultiLineUnion(_))
+    }
+
     pub fn from_vec(types: Vec<LuaType>) -> Self {
         return match types.len() {
             0 => LuaType::Nil,

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

@@ -1279,4 +1279,33 @@ mod test {
             "#
         ));
     }
+
+    #[test]
+    fn test_alias_branch_label_flow() {
+        let mut ws = VirtualWorkspace::new();
+        ws.def_file(
+            "test.lua",
+            r#"
+            ---@alias EditorAttrTypeAlias
+            ---| 'ATTR_BASE'
+            ---| 'ATTR_BASE_RATIO'
+            ---| 'ATTR_ALL_RATIO'
+
+            ---@param attr_element string
+            function test(attr_element) end
+        "#,
+        );
+
+        assert!(ws.check_code_for_namespace(
+            DiagnosticCode::ParamTypeNotMatch,
+            r#"
+            ---@param attr_type EditorAttrTypeAlias
+            function add_attr(attr_type)
+                if attr_type ~= 'ATTR_BASE' then
+                end
+                test(attr_type)
+            end
+        "#
+        ));
+    }
 }