refactor function return infer

Fix #643
Fix #659

Author: CppCXY

crates/emmylua_code_analysis/src/compilation/analyzer/lua/closure.rs

@@ -1,3 +1,5 @@
+use std::ops::Deref;
+
 use emmylua_parser::{
     LuaAst, LuaAstNode, LuaCallArgList, LuaCallExpr, LuaClosureExpr, LuaFuncStat, LuaVarExpr,
 };
@@ -215,7 +217,7 @@ pub fn analyze_return_point(
         match point {
             LuaReturnPoint::Expr(expr) => {
                 let expr_type = infer_expr(db, cache, expr.clone())?;
-                return_type = TypeOps::Union.apply(db, &return_type, &expr_type);
+                return_type = union_return_expr(db, return_type, expr_type);
             }
             LuaReturnPoint::MuliExpr(exprs) => {
                 let mut multi_return = vec![];
@@ -224,10 +226,10 @@ pub fn analyze_return_point(
                     multi_return.push(expr_type);
                 }
                 let typ = LuaType::Variadic(VariadicType::Multi(multi_return).into());
-                return_type = TypeOps::Union.apply(db, &return_type, &typ);
+                return_type = union_return_expr(db, return_type, typ);
             }
             LuaReturnPoint::Nil => {
-                return_type = TypeOps::Union.apply(db, &return_type, &LuaType::Nil);
+                return_type = union_return_expr(db, return_type, LuaType::Nil);
             }
             _ => {}
         }
@@ -239,3 +241,92 @@ pub fn analyze_return_point(
         name: None,
     }])
 }
+
+fn union_return_expr(db: &DbIndex, left: LuaType, right: LuaType) -> LuaType {
+    if left == LuaType::Unknown {
+        return right;
+    }
+
+    match (&left, &right) {
+        (LuaType::Variadic(left_variadic), LuaType::Variadic(right_variadic)) => {
+            match (&left_variadic.deref(), &right_variadic.deref()) {
+                (VariadicType::Base(left_base), VariadicType::Base(right_base)) => {
+                    let union_base = TypeOps::Union.apply(db, left_base, right_base);
+                    LuaType::Variadic(VariadicType::Base(union_base).into())
+                }
+                (VariadicType::Multi(left_multi), VariadicType::Multi(right_multi)) => {
+                    let mut new_multi = vec![];
+                    let max_len = left_multi.len().max(right_multi.len());
+                    for i in 0..max_len {
+                        let left_type = left_multi.get(i).cloned().unwrap_or(LuaType::Nil);
+                        let right_type = right_multi.get(i).cloned().unwrap_or(LuaType::Nil);
+                        new_multi.push(TypeOps::Union.apply(db, &left_type, &right_type));
+                    }
+                    LuaType::Variadic(VariadicType::Multi(new_multi).into())
+                }
+                // difficult to merge the type, use let
+                _ => left.clone(),
+            }
+        }
+        (LuaType::Variadic(variadic), _) => {
+            let first_type = variadic.get_type(0).cloned().unwrap_or(LuaType::Unknown);
+            let first_union_type = TypeOps::Union.apply(db, &first_type, &right);
+
+            match variadic.deref() {
+                VariadicType::Base(base) => {
+                    let union_base = TypeOps::Union.apply(db, base, &LuaType::Nil);
+                    LuaType::Variadic(
+                        VariadicType::Multi(vec![
+                            first_union_type,
+                            LuaType::Variadic(VariadicType::Base(union_base).into()),
+                        ])
+                        .into(),
+                    )
+                }
+                VariadicType::Multi(multi) => {
+                    let mut new_multi = multi.clone();
+                    if new_multi.len() > 0 {
+                        new_multi[0] = first_union_type;
+                        for i in 1..new_multi.len() {
+                            new_multi[i] = TypeOps::Union.apply(db, &new_multi[i], &LuaType::Nil);
+                        }
+                    } else {
+                        new_multi.push(first_union_type);
+                    }
+
+                    LuaType::Variadic(VariadicType::Multi(new_multi).into())
+                }
+            }
+        }
+        (_, LuaType::Variadic(variadic)) => {
+            let first_type = variadic.get_type(0).cloned().unwrap_or(LuaType::Unknown);
+            let first_union_type = TypeOps::Union.apply(db, &left, &first_type);
+            match variadic.deref() {
+                VariadicType::Base(base) => {
+                    let union_base = TypeOps::Union.apply(db, base, &LuaType::Nil);
+                    LuaType::Variadic(
+                        VariadicType::Multi(vec![
+                            first_union_type,
+                            LuaType::Variadic(VariadicType::Base(union_base).into()),
+                        ])
+                        .into(),
+                    )
+                }
+                VariadicType::Multi(multi) => {
+                    let mut new_multi = multi.clone();
+                    if new_multi.len() > 0 {
+                        new_multi[0] = first_union_type;
+                        for i in 1..new_multi.len() {
+                            new_multi[i] = TypeOps::Union.apply(db, &new_multi[i], &LuaType::Nil);
+                        }
+                    } else {
+                        new_multi.push(first_union_type);
+                    }
+
+                    LuaType::Variadic(VariadicType::Multi(new_multi).into())
+                }
+            }
+        }
+        _ => TypeOps::Union.apply(db, &left, &right),
+    }
+}

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

@@ -1,6 +1,6 @@
 #[cfg(test)]
 mod test {
-    use crate::{LuaType, VirtualWorkspace};
+    use crate::{DiagnosticCode, LuaType, VirtualWorkspace};
 
     #[test]
     fn test_issue_376() {
@@ -41,4 +41,53 @@ mod test {
         "#,
         ));
     }
+
+    #[test]
+    fn test_issue_659() {
+        let mut ws = VirtualWorkspace::new_with_init_std_lib();
+
+        assert!(ws.check_code_for(
+            DiagnosticCode::ReturnTypeMismatch,
+            r#"
+        --- @async
+        --- @generic R
+        --- @param fn fun(): R...
+        --- @return R...
+        function wrap(fn) end
+
+        ---@async
+        --- @param a {}?
+        --- @return {}?
+        --- @return string? err
+        function get(a)
+            return wrap(function()
+                if not a then
+                    return nil, 'err'
+                end
+
+                return a
+            end)
+        end
+        "#,
+        ));
+    }
+
+    #[test]
+    fn test_issue_643() {
+        let mut ws = VirtualWorkspace::new_with_init_std_lib();
+
+        assert!(ws.check_code_for(
+            DiagnosticCode::AssignTypeMismatch,
+            r#"
+            local function foo(b)
+                if not b then
+                    return
+                end
+                return 'a', 1
+            end
+            --- @type 'a'?
+            local _ = foo()
+        "#,
+        ));
+    }
 }