feat(code-analysis): support @return_overload flow narrowing

crates/emmylua_code_analysis/resources/std/global.lua

@@ -253,10 +253,11 @@ function pairs(t) end
 --- boolean), which is true if the call succeeds without errors. In such case,
 --- `pcall` also returns all results from the call, after this first result. In
 --- case of any error, `pcall` returns **false** plus the error message.
----@generic T, R, R1
----@param f sync fun(...: T...): R1, R...
+---@generic T, R
+---@param f sync fun(...: T...): R...
 ---@param ... T...
----@return boolean, R1|string, R...
+---@return_overload true, R...
+---@return_overload false, string
 function pcall(f, ...) end
 
 ---

crates/emmylua_code_analysis/src/compilation/analyzer/doc/attribute_tags.rs

@@ -34,6 +34,9 @@ pub fn analyze_tag_attribute_use(
             (LuaAst::LuaDocTagReturn(_), LuaSemanticDeclId::Signature(_)) => {
                 return Some(());
             }
+            (LuaAst::LuaDocTagReturnOverload(_), LuaSemanticDeclId::Signature(_)) => {
+                return Some(());
+            }
             _ => {}
         }
     }
@@ -147,7 +150,8 @@ fn attribute_find_doc(comment: &LuaSyntaxNode) -> Option<LuaSyntaxNode> {
                 LuaKind::Syntax(
                     LuaSyntaxKind::DocTagField
                     | LuaSyntaxKind::DocTagParam
-                    | LuaSyntaxKind::DocTagReturn,
+                    | LuaSyntaxKind::DocTagReturn
+                    | LuaSyntaxKind::DocTagReturnOverload,
                 ) => {
                     if let Some(node) = sibling.as_node() {
                         return Some(node.clone());

crates/emmylua_code_analysis/src/compilation/analyzer/doc/tags.rs

@@ -22,7 +22,7 @@ use super::{
     type_def_tags::{analyze_alias, analyze_class, analyze_enum, analyze_func_generic},
     type_ref_tags::{
         analyze_as, analyze_cast, analyze_module, analyze_other, analyze_overload, analyze_param,
-        analyze_return, analyze_return_cast, analyze_see, analyze_type,
+        analyze_return, analyze_return_cast, analyze_return_overload, analyze_see, analyze_type,
     },
 };
 
@@ -55,6 +55,9 @@ pub fn analyze_tag(analyzer: &mut DocAnalyzer, tag: LuaDocTag) -> Option<()> {
         LuaDocTag::Return(return_tag) => {
             analyze_return(analyzer, return_tag)?;
         }
+        LuaDocTag::ReturnOverload(return_overload_tag) => {
+            analyze_return_overload(analyzer, return_overload_tag)?;
+        }
         LuaDocTag::ReturnCast(return_cast) => {
             analyze_return_cast(analyzer, return_cast)?;
         }

crates/emmylua_code_analysis/src/compilation/analyzer/doc/type_ref_tags.rs

@@ -1,8 +1,8 @@
 use emmylua_parser::{
     LuaAst, LuaAstNode, LuaAstToken, LuaBlock, LuaDocDescriptionOwner, LuaDocTagAs, LuaDocTagCast,
     LuaDocTagModule, LuaDocTagOther, LuaDocTagOverload, LuaDocTagParam, LuaDocTagReturn,
-    LuaDocTagReturnCast, LuaDocTagSchema, LuaDocTagSee, LuaDocTagType, LuaExpr, LuaLocalName,
-    LuaTokenKind, LuaVarExpr,
+    LuaDocTagReturnCast, LuaDocTagReturnOverload, LuaDocTagSchema, LuaDocTagSee, LuaDocTagType,
+    LuaExpr, LuaLocalName, LuaTokenKind, LuaVarExpr,
 };
 
 use super::{
@@ -16,8 +16,8 @@ use crate::{
     SignatureReturnStatus, TypeOps,
     compilation::analyzer::common::bind_type,
     db_index::{
-        LuaDeclId, LuaDocParamInfo, LuaDocReturnInfo, LuaMemberId, LuaOperator, LuaSemanticDeclId,
-        LuaSignatureId, LuaType,
+        LuaDeclId, LuaDocParamInfo, LuaDocReturnInfo, LuaDocReturnOverloadInfo, LuaMemberId,
+        LuaOperator, LuaSemanticDeclId, LuaSignatureId, LuaType,
     },
 };
 use crate::{
@@ -248,29 +248,59 @@ pub fn analyze_return(analyzer: &mut DocAnalyzer, tag: LuaDocTagReturn) -> Optio
     let description = tag
         .get_description()
         .map(|des| preprocess_description(&des.get_description_text(), None));
+    let return_infos = tag
+        .get_info_list()
+        .into_iter()
+        .map(|(doc_type, name_token)| LuaDocReturnInfo {
+            name: name_token.map(|name| name.get_name_text().to_string()),
+            type_ref: infer_type(analyzer, doc_type),
+            description: description.clone(),
+            attributes: None,
+        })
+        .collect::<Vec<_>>();
+
+    bind_signature_return_docs(analyzer, &tag, |signature| {
+        signature.return_docs.extend(return_infos);
+    })
+}
 
-    if let Some(closure) = find_owner_closure_or_report(analyzer, &tag) {
-        let signature_id = LuaSignatureId::from_closure(analyzer.file_id, &closure);
-        let returns = tag.get_info_list();
-        for (doc_type, name_token) in returns {
-            let name = name_token.map(|name| name.get_name_text().to_string());
-
-            let type_ref = infer_type(analyzer, doc_type);
-            let return_info = LuaDocReturnInfo {
-                name,
-                type_ref,
-                description: description.clone(),
-                attributes: None,
-            };
-
-            let signature = analyzer
-                .db
-                .get_signature_index_mut()
-                .get_or_create(signature_id);
-            signature.return_docs.push(return_info);
-            signature.resolve_return = SignatureReturnStatus::DocResolve;
-        }
+pub fn analyze_return_overload(
+    analyzer: &mut DocAnalyzer,
+    tag: LuaDocTagReturnOverload,
+) -> Option<()> {
+    let description = tag
+        .get_description()
+        .map(|des| preprocess_description(&des.get_description_text(), None))
+        .filter(|des| !des.is_empty());
+    let overload_info = LuaDocReturnOverloadInfo {
+        type_refs: tag
+            .get_types()
+            .map(|doc_type| infer_type(analyzer, doc_type))
+            .collect(),
+        description,
+    };
+    if overload_info.type_refs.is_empty() {
+        return Some(());
     }
+
+    bind_signature_return_docs(analyzer, &tag, |signature| {
+        signature.return_overloads.push(overload_info);
+    })
+}
+
+fn bind_signature_return_docs(
+    analyzer: &mut DocAnalyzer,
+    tag: &impl LuaAstNode,
+    bind: impl FnOnce(&mut crate::LuaSignature),
+) -> Option<()> {
+    let closure = find_owner_closure_or_report(analyzer, tag)?;
+    let signature_id = LuaSignatureId::from_closure(analyzer.file_id, &closure);
+    let signature = analyzer
+        .db
+        .get_signature_index_mut()
+        .get_or_create(signature_id);
+    bind(signature);
+    signature.resolve_return = SignatureReturnStatus::DocResolve;
     Some(())
 }
 

crates/emmylua_code_analysis/src/compilation/analyzer/flow/bind_analyze/stats.rs

@@ -1,11 +1,13 @@
 use emmylua_parser::{
     BinaryOperator, LuaAssignStat, LuaAst, LuaAstNode, LuaBlock, LuaBreakStat, LuaCallArgList,
     LuaCallExprStat, LuaDoStat, LuaExpr, LuaForRangeStat, LuaForStat, LuaFuncStat, LuaGotoStat,
-    LuaIfStat, LuaLabelStat, LuaLocalStat, LuaRepeatStat, LuaReturnStat, LuaVarExpr, LuaWhileStat,
+    LuaIfStat, LuaLabelStat, LuaLocalName, LuaLocalStat, LuaRepeatStat, LuaReturnStat, LuaVarExpr,
+    LuaWhileStat,
 };
 
 use crate::{
-    AnalyzeError, DiagnosticCode, FlowId, FlowNodeKind, LuaClosureId, LuaDeclId,
+    AnalyzeError, DeclMultiReturnRef, DeclMultiReturnRefAt, DiagnosticCode, FlowId, FlowNodeKind,
+    LuaClosureId, LuaDeclId,
     compilation::analyzer::flow::{
         bind_analyze::{
             bind_block, bind_each_child, bind_node,
@@ -33,13 +35,20 @@ pub fn bind_local_stat(
         }
     }
 
-    for value in values {
+    for value in &values {
         // If there are more values than names, we still need to bind the values
         bind_expr(binder, value.clone(), current);
     }
 
     let local_flow_id = binder.create_decl(local_stat.get_position());
     binder.add_antecedent(local_flow_id, current);
+    bind_multi_return_refs(
+        binder,
+        &get_local_decl_ids(binder, &local_names),
+        &values,
+        local_stat.get_position(),
+        local_flow_id,
+    );
     local_flow_id
 }
 
@@ -69,6 +78,27 @@ fn check_value_expr_is_check_expr(value_expr: LuaExpr) -> bool {
     }
 }
 
+fn get_local_decl_ids(
+    binder: &FlowBinder<'_>,
+    local_names: &[LuaLocalName],
+) -> Vec<Option<LuaDeclId>> {
+    local_names
+        .iter()
+        .map(|name| Some(LuaDeclId::new(binder.file_id, name.get_position())))
+        .collect()
+}
+
+fn get_var_decl_ids(binder: &FlowBinder<'_>, vars: &[LuaVarExpr]) -> Vec<Option<LuaDeclId>> {
+    vars.iter()
+        .map(|var| {
+            binder
+                .db
+                .get_reference_index()
+                .get_var_reference_decl(&binder.file_id, var.get_range())
+        })
+        .collect()
+}
+
 pub fn bind_assign_stat(
     binder: &mut FlowBinder,
     assign_stat: LuaAssignStat,
@@ -91,10 +121,57 @@ pub fn bind_assign_stat(
     let assignment_kind = FlowNodeKind::Assignment(assign_stat.to_ptr());
     let flow_id = binder.create_node(assignment_kind);
     binder.add_antecedent(flow_id, current);
+    bind_multi_return_refs(
+        binder,
+        &get_var_decl_ids(binder, &vars),
+        &values,
+        assign_stat.get_position(),
+        flow_id,
+    );
 
     flow_id
 }
 
+fn bind_multi_return_refs(
+    binder: &mut FlowBinder,
+    decl_ids: &[Option<LuaDeclId>],
+    values: &[LuaExpr],
+    position: rowan::TextSize,
+    flow_id: FlowId,
+) {
+    let tail_call = values.last().and_then(|value| match value {
+        LuaExpr::CallExpr(call_expr) => Some((values.len() - 1, call_expr.to_ptr())),
+        _ => None,
+    });
+
+    for (i, decl_id) in decl_ids.iter().enumerate() {
+        let Some(decl_id) = decl_id else {
+            continue;
+        };
+
+        let reference = tail_call.as_ref().and_then(|(last_value_idx, call_expr)| {
+            if i < *last_value_idx {
+                return None;
+            }
+
+            Some(DeclMultiReturnRef {
+                call_expr: call_expr.clone(),
+                return_index: i - *last_value_idx,
+            })
+        });
+
+        binder
+            .decl_multi_return_ref
+            .entry(*decl_id)
+            .or_default()
+            .push(DeclMultiReturnRefAt {
+                position,
+                flow_id,
+                reference,
+            });
+    }
+}
+
 pub fn bind_call_expr_stat(
     binder: &mut FlowBinder,
     call_expr_stat: LuaCallExprStat,

crates/emmylua_code_analysis/src/compilation/analyzer/flow/binder.rs

@@ -6,15 +6,16 @@ use rowan::TextSize;
 use smol_str::SmolStr;
 
 use crate::{
-    AnalyzeError, DbIndex, FileId, FlowAntecedent, FlowId, FlowNode, FlowNodeKind, FlowTree,
-    LuaClosureId, LuaDeclId,
+    AnalyzeError, DbIndex, DeclMultiReturnRefAt, FileId, FlowAntecedent, FlowId, FlowNode,
+    FlowNodeKind, FlowTree, LuaClosureId, LuaDeclId,
 };
 
 #[derive(Debug)]
 pub struct FlowBinder<'a> {
     pub db: &'a mut DbIndex,
     pub file_id: FileId,
     pub decl_bind_expr_ref: HashMap<LuaDeclId, LuaAstPtr<LuaExpr>>,
+    pub decl_multi_return_ref: HashMap<LuaDeclId, Vec<DeclMultiReturnRefAt>>,
     pub start: FlowId,
     pub unreachable: FlowId,
     pub loop_label: FlowId,
@@ -36,6 +37,7 @@ impl<'a> FlowBinder<'a> {
             flow_nodes: Vec::new(),
             multiple_antecedents: Vec::new(),
             decl_bind_expr_ref: HashMap::new(),
+            decl_multi_return_ref: HashMap::new(),
             labels: HashMap::new(),
             start: FlowId::default(),
             unreachable: FlowId::default(),
@@ -189,6 +191,7 @@ impl<'a> FlowBinder<'a> {
     pub fn finish(self) -> FlowTree {
         FlowTree::new(
             self.decl_bind_expr_ref,
+            self.decl_multi_return_ref,
             self.flow_nodes,
             self.multiple_antecedents,
             // self.labels,

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

@@ -31,15 +31,7 @@ pub fn analyze_chunk_return(analyzer: &mut LuaAnalyzer, chunk: LuaChunk) -> Opti
                 .db
                 .get_module_index_mut()
                 .get_module_mut(analyzer.file_id)?;
-            match expr_type {
-                LuaType::Variadic(multi) => {
-                    let ty = multi.get_type(0)?;
-                    module_info.export_type = Some(ty.clone());
-                }
-                _ => {
-                    module_info.export_type = Some(expr_type);
-                }
-            }
+            module_info.export_type = Some(expr_type.get_result_slot_type(0).unwrap_or(expr_type));
             module_info.semantic_id = semantic_id;
             break;
         }

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

@@ -47,10 +47,8 @@ pub fn analyze_local_stat(analyzer: &mut LuaAnalyzer, local_stat: LuaLocalStat)
         };
 
         match analyzer.infer_expr(&expr) {
-            Ok(mut expr_type) => {
-                if let LuaType::Variadic(multi) = expr_type {
-                    expr_type = multi.get_type(0)?.clone();
-                }
+            Ok(expr_type) => {
+                let expr_type = expr_type.get_result_slot_type(0).unwrap_or(expr_type);
                 let decl_id = LuaDeclId::new(analyzer.file_id, position);
                 // 当`call`参数包含表时, 表可能未被分析, 需要延迟
                 if let LuaType::Instance(instance) = &expr_type
@@ -106,12 +104,12 @@ pub fn analyze_local_stat(analyzer: &mut LuaAnalyzer, local_stat: LuaLocalStat)
         if let Some(last_expr) = last_expr {
             match analyzer.infer_expr(last_expr) {
                 Ok(last_expr_type) => {
-                    if let LuaType::Variadic(variadic) = last_expr_type {
+                    if last_expr_type.contain_multi_return() {
                         for i in expr_count..name_count {
                             let name = name_list.get(i)?;
                             let position = name.get_position();
                             let decl_id = LuaDeclId::new(analyzer.file_id, position);
-                            let ret_type = variadic.get_type(i - expr_count + 1);
+                            let ret_type = last_expr_type.get_result_slot_type(i - expr_count + 1);
                             if let Some(ret_type) = ret_type {
                                 bind_type(
                                     analyzer.db,
@@ -311,10 +309,7 @@ pub fn analyze_assign_stat(analyzer: &mut LuaAnalyzer, assign_stat: LuaAssignSta
         }
 
         let expr_type = match analyzer.infer_expr(expr) {
-            Ok(expr_type) => match expr_type {
-                LuaType::Variadic(multi) => multi.get_type(0)?.clone(),
-                _ => expr_type,
-            },
+            Ok(expr_type) => expr_type.get_result_slot_type(0).unwrap_or(expr_type),
             Err(InferFailReason::None) => LuaType::Unknown,
             Err(reason) => {
                 match type_owner {
@@ -367,7 +362,7 @@ pub fn analyze_assign_stat(analyzer: &mut LuaAnalyzer, assign_stat: LuaAssignSta
     {
         match analyzer.infer_expr(last_expr) {
             Ok(last_expr_type) => {
-                if last_expr_type.is_multi_return() {
+                if last_expr_type.contain_multi_return() {
                     for i in expr_count..var_count {
                         let var = var_list.get(i)?;
                         let type_owner = get_var_owner(analyzer, var.clone());
@@ -408,10 +403,7 @@ fn assign_merge_type_owner_and_expr_type(
     expr_type: &LuaType,
     idx: usize,
 ) -> Option<()> {
-    let mut expr_type = expr_type.clone();
-    if let LuaType::Variadic(multi) = expr_type {
-        expr_type = multi.get_type(idx).unwrap_or(&LuaType::Nil).clone();
-    }
+    let expr_type = expr_type.get_result_slot_type(idx).unwrap_or(LuaType::Nil);
 
     bind_type(analyzer.db, type_owner, LuaTypeCache::InferType(expr_type));