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/type_ref_tags.rs

@@ -245,12 +245,33 @@ pub fn analyze_param(analyzer: &mut DocAnalyzer, tag: LuaDocTagParam) -> Option<
 }
 
 pub fn analyze_return(analyzer: &mut DocAnalyzer, tag: LuaDocTagReturn) -> Option<()> {
+    let is_return_overload = tag
+        .token_by_kind(LuaTokenKind::TkTagReturnOverload)
+        .is_some();
     let description = tag
         .get_description()
         .map(|des| preprocess_description(&des.get_description_text(), None));
 
     if let Some(closure) = find_owner_closure_or_report(analyzer, &tag) {
         let signature_id = LuaSignatureId::from_closure(analyzer.file_id, &closure);
+        if is_return_overload {
+            let overload_types = tag
+                .get_types()
+                .map(|doc_type| infer_type(analyzer, doc_type))
+                .collect::<Vec<_>>();
+            if overload_types.is_empty() {
+                return Some(());
+            }
+
+            let signature = analyzer
+                .db
+                .get_signature_index_mut()
+                .get_or_create(signature_id);
+            signature.return_overloads.push(overload_types);
+            signature.resolve_return = SignatureReturnStatus::DocResolve;
+            return Some(());
+        }
+
         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());

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

@@ -5,7 +5,8 @@ use emmylua_parser::{
 };
 
 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 +34,25 @@ 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 decl_ids = local_names
+        .iter()
+        .map(|name| Some(LuaDeclId::new(binder.file_id, name.get_position())))
+        .collect::<Vec<_>>();
+
     let local_flow_id = binder.create_decl(local_stat.get_position());
     binder.add_antecedent(local_flow_id, current);
+    bind_multi_return_refs(
+        binder,
+        &decl_ids,
+        &values,
+        local_stat.get_position(),
+        local_flow_id,
+    );
     local_flow_id
 }
 
@@ -88,13 +101,72 @@ pub fn bind_assign_stat(
         }
     }
 
+    let decl_ids = vars
+        .iter()
+        .map(|var| {
+            binder
+                .db
+                .get_reference_index()
+                .get_var_reference_decl(&binder.file_id, var.get_range())
+        })
+        .collect::<Vec<_>>();
+
     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,
+        &decl_ids,
+        &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| {
+        if let LuaExpr::CallExpr(call_expr) = value {
+            Some((values.len() - 1, call_expr.to_ptr()))
+        } else {
+            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,