optimize infer_union

Author: xuhuanzy

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

@@ -205,4 +205,22 @@ mod tests {
             "#
         ));
     }
+
+    #[test]
+    fn test_1() {
+        let mut ws = VirtualWorkspace::new();
+
+        assert!(ws.check_code_for(
+            DiagnosticCode::ReturnTypeMismatch,
+            r#"
+            ---@return string?
+            local function a()
+                ---@type int?
+                local ccc
+                return ccc and a() or nil
+            end
+
+            "#
+        ));
+    }
 }

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

@@ -1,3 +1,5 @@
+use std::collections::HashSet;
+
 use emmylua_parser::{BinaryOperator, LuaBinaryExpr};
 use smol_str::SmolStr;
 
@@ -61,17 +63,35 @@ fn infer_binary_expr_type(
 }
 
 fn infer_union(db: &DbIndex, u: &LuaUnionType, right: &LuaType, op: BinaryOperator) -> InferResult {
-    let mut union_types = vec![];
+    let mut unique_union_types = HashSet::new();
+
     for ty in u.get_types() {
-        let ty = infer_binary_expr_type(db, ty.clone(), right.clone(), op)?;
-        union_types.push(ty);
+        let inferred_ty = infer_binary_expr_type(db, ty.clone(), right.clone(), op)?;
+        flatten_and_insert(inferred_ty, &mut unique_union_types);
     }
-    union_types.dedup();
 
-    match union_types.len() {
+    match unique_union_types.len() {
         0 => Some(LuaType::Unknown),
-        1 => Some(union_types[0].clone()),
-        _ => Some(LuaType::Union(LuaUnionType::new(union_types).into())),
+        1 => Some(unique_union_types.into_iter().next().unwrap()),
+        _ => Some(LuaType::Union(
+            LuaUnionType::new(unique_union_types.into_iter().collect()).into(),
+        )),
+    }
+}
+
+fn flatten_and_insert(ty: LuaType, unique_union_types: &mut HashSet<LuaType>) {
+    let mut stack = vec![ty];
+    while let Some(current_ty) = stack.pop() {
+        match current_ty {
+            LuaType::Union(u) => {
+                for inner_ty in u.get_types() {
+                    stack.push(inner_ty.clone());
+                }
+            }
+            _ => {
+                unique_union_types.insert(current_ty);
+            }
+        }
     }
 }