update flow

Author: xuhuanzy

crates/emmylua_code_analysis/resources/std/builtin.lua

@@ -149,6 +149,8 @@
 --- Get the parameters of a constructor as a tuple
 ---@alias ConstructorParameters<T> T extends new (fun(...: infer P): any) and P or never
 
+---@alias ReturnType<T extends function> T extends (fun(...: any): infer R) and R or any
+
 --- Make all properties in T optional
 ---@alias Partial<T> { [P in keyof T]?: T[P]; }
 

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

@@ -726,18 +726,48 @@ mod test {
         ws.def(
             r#"
             ---@alias ConstructorParameters<T> T extends new (fun(...: infer P): any) and P or never
+
             ---@alias Parameters<T extends function> T extends (fun(...: infer P): any) and P or never
 
+            ---@alias ReturnType<T extends function> T extends (fun(...: any): infer R) and R or any
+
             ---@alias Procedure fun(...: any[]): any
 
-            ---@alias MockParameters<T> T extends table and ConstructorParameters<T> or T extends Procedure and Parameters<T> or never
+            ---@alias MockParameters<T> T extends Procedure and Parameters<T> or never
+
+            ---@alias MockReturnType<T> T extends Procedure and ReturnType<T> or never
 
             ---@class Mock<T>
             ---@field calls MockParameters<T>[]
+            ---@overload fun(...: MockParameters<T>...): MockReturnType<T>
             "#,
         );
-        ws.def(
-            r#"
+        {
+            ws.def(
+                r#"
+                ---@generic T: Procedure
+                ---@param a T
+                ---@return Mock<T>
+                local function fn(a)
+                end
+
+                local sum = fn(function(a, b)
+                    return a + b
+                end)
+                A = sum
+            "#,
+            );
+
+            let result_ty = ws.expr_ty("A");
+            assert_eq!(
+                ws.humanize_type_detailed(result_ty),
+                "Mock<fun(a, b) -> any>"
+            );
+        }
+
+        {
+            ws.def(
+                r#"
                 ---@generic T: Procedure
                 ---@param a T?
                 ---@return Mock<T>
@@ -746,9 +776,10 @@ mod test {
 
                 result = fn().calls
             "#,
-        );
+            );
 
-        let result_ty = ws.expr_ty("result");
-        assert_eq!(ws.humanize_type(result_ty), "any[][]");
+            let result_ty = ws.expr_ty("result");
+            assert_eq!(ws.humanize_type(result_ty), "any[][]");
+        }
     }
 }

crates/emmylua_code_analysis/src/semantic/infer/narrow/get_type_at_flow.rs

@@ -1,4 +1,4 @@
-use emmylua_parser::{LuaAssignStat, LuaAstNode, LuaChunk, LuaVarExpr};
+use emmylua_parser::{LuaAssignStat, LuaAstNode, LuaChunk, LuaExpr, LuaVarExpr};
 
 use crate::{
     CacheEntry, DbIndex, FlowId, FlowNode, FlowNodeKind, FlowTree, InferFailReason, LuaDeclId,
@@ -61,8 +61,23 @@ pub fn get_type_at_flow(
             }
             FlowNodeKind::DeclPosition(position) => {
                 if *position <= var_ref_id.get_position() {
-                    result_type = get_var_ref_type(db, cache, var_ref_id)?;
-                    break;
+                    match get_var_ref_type(db, cache, var_ref_id) {
+                        Ok(var_type) => {
+                            result_type = var_type;
+                            break;
+                        }
+                        Err(err) => {
+                            // 尝试推断声明位置的类型, 如果发生错误则返回初始错误, 否则返回当前推断错误
+                            if let Some(init_type) =
+                                try_infer_decl_initializer_type(db, cache, root, var_ref_id)?
+                            {
+                                result_type = init_type;
+                                break;
+                            }
+
+                            return Err(err);
+                        }
+                    }
                 } else {
                     antecedent_flow_id = get_single_antecedent(tree, flow_node)?;
                 }
@@ -258,3 +273,39 @@ fn get_type_at_assign_stat(
 
     Ok(ResultTypeOrContinue::Continue)
 }
+
+fn try_infer_decl_initializer_type(
+    db: &DbIndex,
+    cache: &mut LuaInferCache,
+    root: &LuaChunk,
+    var_ref_id: &VarRefId,
+) -> Result<Option<LuaType>, InferFailReason> {
+    let Some(decl_id) = var_ref_id.get_decl_id_ref() else {
+        return Ok(None);
+    };
+
+    let decl = db
+        .get_decl_index()
+        .get_decl(&decl_id)
+        .ok_or(InferFailReason::None)?;
+
+    let Some(value_syntax_id) = decl.get_value_syntax_id() else {
+        return Ok(None);
+    };
+
+    let Some(node) = value_syntax_id.to_node_from_root(root.syntax()) else {
+        return Ok(None);
+    };
+
+    let Some(expr) = LuaExpr::cast(node) else {
+        return Ok(None);
+    };
+
+    let expr_type = infer_expr(db, cache, expr.clone())?;
+    let init_type = match expr_type {
+        LuaType::Variadic(variadic) => variadic.get_type(0).cloned(),
+        ty => Some(ty),
+    };
+
+    Ok(init_type)
+}