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

Add `---@return_overload` support end-to-end and use it to narrow correlated
multi-return values in condition flow.

- Parse/AST:
  - add `return_overload` token/syntax kinds
  - parse `---@return_overload <type>(, <type>)*`
  - add doc AST node and lexer/tag wiring
- Analyzer/index:
  - handle `LuaDocTag::ReturnOverload`
  - store overload rows on `LuaSignature`
  - compute return type from overload rows (slot-wise union, variadic-aware)
- Flow infra:
  - track `decl -> (call_expr, return_index)` mappings in binder/flow tree
  - clear stale mappings on reassignment
- Narrowing:
  - add overload condition model (truthy/falsy/eq/neq)
  - narrow target vars using discriminant vars from the same call
  - support swapped equality operand order in binary condition flow
- Stdlib/semantic tokens:
  - annotate `pcall` with `@return_overload` rows
  - highlight `TkTagReturnOverload` as a doc tag

Author: lewis6991

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, DiagnosticCode, FlowId, FlowNodeKind, LuaClosureId,
+    LuaDeclId,
     compilation::analyzer::flow::{
         bind_analyze::{
             bind_block, bind_each_child, bind_node,
@@ -33,11 +34,19 @@ 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<_>>();
+
+    // Track `local a, b = call()` slot ownership for later flow narrowing.
+    bind_multi_return_refs(binder, &decl_ids, &values);
+
     let local_flow_id = binder.create_decl(local_stat.get_position());
     binder.add_antecedent(local_flow_id, current);
     local_flow_id
@@ -88,13 +97,80 @@ 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<_>>();
+
+    // Track `a, b = call()` slot ownership for later flow narrowing.
+    bind_multi_return_refs(binder, &decl_ids, &values);
+
     let assignment_kind = FlowNodeKind::Assignment(assign_stat.to_ptr());
     let flow_id = binder.create_node(assignment_kind);
     binder.add_antecedent(flow_id, current);
 
     flow_id
 }
 
+/// Binds declaration IDs to call-return slots in assignment/local statements.
+///
+/// This lets condition flow recover which variable is the discriminant/result from the same call.
+fn bind_multi_return_refs(
+    binder: &mut FlowBinder,
+    decl_ids: &[Option<LuaDeclId>],
+    values: &[LuaExpr],
+) {
+    // Rebinding invalidates previous call-slot links.
+    for decl_id in decl_ids.iter().flatten() {
+        binder.decl_multi_return_ref.remove(&decl_id);
+    }
+
+    if values.is_empty() {
+        return;
+    }
+
+    let min_len = decl_ids.len().min(values.len());
+    for i in 0..min_len {
+        let Some(decl_id) = decl_ids[i] else {
+            continue;
+        };
+        let LuaExpr::CallExpr(call_expr) = &values[i] else {
+            continue;
+        };
+        binder.decl_multi_return_ref.insert(
+            decl_id,
+            DeclMultiReturnRef {
+                // Direct pair uses the first return slot.
+                call_expr: call_expr.to_ptr(),
+                return_index: 0,
+            },
+        );
+    }
+
+    if decl_ids.len() > values.len()
+        && let Some(LuaExpr::CallExpr(call_expr)) = values.last()
+    {
+        for i in values.len()..decl_ids.len() {
+            let Some(decl_id) = decl_ids[i] else {
+                continue;
+            };
+            binder.decl_multi_return_ref.insert(
+                decl_id,
+                DeclMultiReturnRef {
+                    call_expr: call_expr.to_ptr(),
+                    // Remaining LHS names consume expanded trailing return slots.
+                    return_index: i - values.len() + 1,
+                },
+            );
+        }
+    }
+}
+
 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, DeclMultiReturnRef, 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, DeclMultiReturnRef>,
     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/test/flow.rs

@@ -1710,6 +1710,169 @@ _2 = a[1]
         assert_eq!(e, e_expected);
     }
 
+    #[test]
+    fn test_return_overload_narrow_after_not() {
+        // Boolean guard on discriminant should narrow correlated result slot.
+        let mut ws = VirtualWorkspace::new();
+
+        ws.def(
+            r#"
+            ---@generic T, E
+            ---@param ok boolean
+            ---@param success T
+            ---@param failure E
+            ---@return boolean
+            ---@return T|E
+            ---@return_overload true, T
+            ---@return_overload false, E
+            local function pick(ok, success, failure)
+                if ok then
+                    return true, success
+                end
+                return false, failure
+            end
+
+            local cond ---@type boolean
+            local ok, result = pick(cond, 1, "error")
+
+            if not ok then
+                error(result)
+            end
+
+            a = result
+            "#,
+        );
+
+        let a = ws.expr_ty("a");
+        let expected = ws.ty("integer");
+        assert_eq!(a, expected);
+    }
+
+    #[test]
+    fn test_return_overload_narrow_with_swapped_operand_eq() {
+        // Equality narrowing should work when literal is on the left side.
+        let mut ws = VirtualWorkspace::new();
+
+        ws.def(
+            r#"
+            ---@generic T, E
+            ---@param ok boolean
+            ---@param success T
+            ---@param failure E
+            ---@return "ok"|"err"
+            ---@return T|E
+            ---@return_overload "ok", T
+            ---@return_overload "err", E
+            local function pick(ok, success, failure)
+                if ok then
+                    return "ok", success
+                end
+                return "err", failure
+            end
+
+            local cond ---@type boolean
+            local tag, result = pick(cond, 1, "error")
+
+            if "err" == tag then
+                error(result)
+            end
+
+            d = result
+            "#,
+        );
+
+        let d = ws.expr_ty("d");
+        let expected = ws.ty("integer");
+        assert_eq!(d, expected);
+    }
+
+    #[test]
+    fn test_swapped_literal_eq_narrow_without_return_overload() {
+        // Baseline: swapped literal equality still narrows regular unions.
+        let mut ws = VirtualWorkspace::new();
+
+        assert!(!ws.check_code_for(
+            DiagnosticCode::ReturnTypeMismatch,
+            r#"
+            ---@return "x"
+            local function test()
+                local a ---@type "x"|nil
+                if "x" == a then
+                    return a
+                end
+                return "x"
+            end
+            "#,
+        ));
+    }
+
+    #[test]
+    fn test_return_overload_reassign_clears_multi_return_mapping() {
+        // Reassignment should break call-slot correlation, preventing stale narrowing.
+        let mut ws = VirtualWorkspace::new();
+
+        ws.def(
+            r#"
+            ---@generic T, E
+            ---@param ok boolean
+            ---@param success T
+            ---@param failure E
+            ---@return boolean
+            ---@return T|E
+            ---@return_overload true, T
+            ---@return_overload false, E
+            local function pick(ok, success, failure)
+                if ok then
+                    return true, success
+                end
+                return false, failure
+            end
+
+            local cond ---@type boolean
+            local random ---@type boolean
+            local ok, result = pick(cond, 1, "error")
+            result = random and 1 or "override"
+
+            if not ok then
+                error(result)
+            end
+
+            f = result
+            "#,
+        );
+
+        let f = ws.expr_ty("f");
+        let expected = ws.ty("integer|string");
+        assert_eq!(f, expected);
+    }
+
+    #[test]
+    fn test_pcall_return_overload_narrow_after_error_guard() {
+        // Stdlib `pcall` overload rows should narrow result after error branch exits.
+        let mut ws = VirtualWorkspace::new_with_init_std_lib();
+
+        ws.def(
+            r#"
+            ---@return integer
+            local function foo()
+                return 2
+            end
+
+            local ok, result = pcall(foo)
+
+            if not ok then
+                error(result)
+            end
+
+            a = result
+            "#,
+        );
+
+        let a = ws.expr_ty("a");
+        let expected = ws.ty("integer");
+        assert_eq!(a, expected);
+    }
+
     #[test]
     fn test_issue_868() {
         let mut ws = VirtualWorkspace::new();