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

[lewis6991]

This PR adds correlated ---@return_overload rows for multi-return APIs.

Today, plain ---@return docs flatten each return slot independently, which loses the relationship between values
like ok and result. ---@return_overload lets each annotation row describe one possible return tuple, so the
analyzer can preserve that correlation through inference and control-flow narrowing.

---@return_overload true, T
---@return_overload false, string
local ok, result = pcall(foo)

if not ok then
  error(result)
end

-- `result` narrows to `T` here

What Changed

• Parse and document ---@return_overload.
• Store return overload rows separately from plain return docs and surface them in tooling output.
• Use overload rows during return inference, including generic and variadic tails.
• Pad missing trailing overload slots with nil so short rows still produce correct tuple shapes.
• Narrow correlated return slots on truthy/falsy and literal-equality checks while staying sound across reassignments,
  mixed RHS calls, and branch joins.
• Annotate stdlib pcall with correlated return rows so status/result flows narrow correctly.
• Update incomplete-signature diagnostics to understand return overload rows.

Coverage Added

• Parser/doc tests for the new tag.
• Analyzer tests for generic and variadic return rows.
• Flow tests for correlated narrowing, reassignments, joins, and equality guards.
• pcall regression coverage.
• Hover, semantic-token, and diagnostic coverage for the new annotation.

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the type inference system by introducing support for ---@return_overload doc tags. This new feature allows the type checker to understand and narrow down the types of multiple return values from a function based on specific conditions, such as the success or failure status returned by a pcall operation. The changes span across parsing, semantic analysis, flow control, and standard library definitions, providing a more robust and accurate type analysis experience for complex Lua return patterns.

Highlights

• New Doc Tag Support: Introduced the ---@return_overload doc tag to specify discriminated return types for functions, allowing for more precise type inference based on return conditions.
• Correlated Multi-Return Narrowing: Implemented flow narrowing logic that correlates multiple return values from a single function call. This enables the type of one return value to be narrowed based on the condition of another return value (e.g., ok, result = pcall(); if ok then ...).
• Standard Library Integration: Updated the standard library definition for pcall to utilize the new ---@return_overload tag, enhancing its type inference capabilities.
• Parsing and AST Updates: Added new token and syntax kinds for return_overload and integrated its parsing into the AST, ensuring the doc tag is correctly recognized and structured.
• Analyzer and Flow Infrastructure: Modified the analyzer to handle LuaDocTag::ReturnOverload, store overload rows on LuaSignature, and compute return types by slot-wise union. The flow infrastructure now tracks decl -> (call_expr, return_index) mappings to enable correlated narrowing.
• Enhanced Narrowing Logic: Developed an overload condition model (truthy/falsy/eq/neq) and improved binary condition flow to support narrowing with swapped equality operand order (e.g., literal == variable).
• Semantic Token Highlighting: Added highlighting for TkTagReturnOverload to improve readability and consistency with other doc tags in the editor.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• crates/emmylua_code_analysis/resources/std/global.lua
  • Updated pcall function signature to use ---@return_overload for discriminated return types.
• crates/emmylua_code_analysis/src/compilation/analyzer/doc/tags.rs
  • Integrated analyze_return_overload into the main tag analysis function.
• crates/emmylua_code_analysis/src/compilation/analyzer/doc/type_ref_tags.rs
  • Added analyze_return_overload function to process the new doc tag and store its type information.
• crates/emmylua_code_analysis/src/compilation/analyzer/flow/bind_analyze/stats.rs
  • Modified bind_local_stat and bind_assign_stat to track multi-return references.
  • Introduced bind_multi_return_refs to link declaration IDs to call-return slots for flow narrowing.
• crates/emmylua_code_analysis/src/compilation/analyzer/flow/binder.rs
  • Added decl_multi_return_ref to FlowBinder and FlowTree to store mappings for multi-return values.
• crates/emmylua_code_analysis/src/compilation/test/flow.rs
  • Added new test cases for return_overload narrowing, covering boolean guards, swapped equality operands, reassignment, and pcall integration.
• crates/emmylua_code_analysis/src/db_index/flow/flow_tree.rs
  • Introduced DeclMultiReturnRef struct to represent call expression and return index.
  • Added decl_multi_return_ref field and accessor to FlowTree.
• crates/emmylua_code_analysis/src/db_index/flow/mod.rs
  • Exported DeclMultiReturnRef for external use.
• crates/emmylua_code_analysis/src/db_index/signature/signature.rs
  • Added return_overloads field to LuaSignature to store discriminated return type rows.
  • Implemented get_return_type_by_overloads, union_overload_row_slot, and get_overload_row_slot for computing return types from overload rows.
• crates/emmylua_code_analysis/src/semantic/infer/narrow/condition_flow/binary_flow.rs
  • Modified try_get_at_eq_or_neq_expr to support narrowing with swapped equality operands.
  • Integrated narrow_var_from_return_overload_condition for correlated return value narrowing.
  • Removed the const_type_eq helper function.
• crates/emmylua_code_analysis/src/semantic/infer/narrow/condition_flow/mod.rs
  • Defined ReturnOverloadCondition enum for various narrowing conditions.
  • Implemented narrow_var_from_return_overload_condition to apply narrowing based on correlated return overloads.
  • Added helper functions get_or_instantiate_return_overload_rows, instantiate_return_overload_rows, is_return_overload_row_match, always_literal_equal, always_literal_not_equal, and scalar_literal_eq for overload row processing and literal comparison.
• crates/emmylua_ls/src/handlers/semantic_token/build_semantic_tokens.rs
  • Added TkTagReturnOverload to the list of tokens to be highlighted as a doc tag.
• crates/emmylua_parser/src/grammar/doc/tag.rs
  • Added parsing logic for the new ---@return_overload doc tag.
• crates/emmylua_parser/src/grammar/doc/test.rs
  • Added a test case to verify correct parsing of the ---@return_overload tag.
• crates/emmylua_parser/src/kind/lua_syntax_kind.rs
  • Added DocTagReturnOverload to the LuaSyntaxKind enum.
• crates/emmylua_parser/src/kind/lua_token_kind.rs
  • Added TkTagReturnOverload to the LuaTokenKind enum.
• crates/emmylua_parser/src/lexer/lua_doc_lexer.rs
  • Mapped the string "return_overload" to the new TkTagReturnOverload token kind.
• crates/emmylua_parser/src/syntax/node/doc/tag.rs
  • Added ReturnOverload variant to LuaDocTag enum.
  • Defined LuaDocTagReturnOverload struct for the new doc tag.
• crates/emmylua_parser/src/syntax/node/mod.rs
  • Updated LuaAst enum and its associated implementations to include LuaDocTagReturnOverload.

Activity

• No human activity has occurred on this pull request yet.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces comprehensive support for @return_overload tags, enabling advanced flow narrowing for correlated multi-return values in condition flow. The changes span across parsing, AST representation, semantic analysis, and type inference, including updates to the standard library definitions and semantic token highlighting. The implementation correctly handles generic type instantiation and variadic returns within overload rows, and includes robust test cases for various narrowing scenarios, including re-assignment and swapped equality operands. Overall, this is a well-implemented and significant feature enhancement.

[lewis6991]

@CppCXY what do you think about this?

Admittedly I got AI to implement all of this, but I guided it best I could and reviewed all the changes.

I'm not sure how good the overall strategy is here, but it seems to align with how other stuff is analysed. There's some slight code duplication with literal comparison narrowing, but that can be improved later.

[CppCXY]

How do return overloads match with generics, and what would pcall(pcall, f) return?

[lewis6991]

Good catch, It didn't handle it well, though either did main. I've made some changes to fix this.

I ended up running a "review, fix, minimize" loop in codex which ended up finding lots of edge cases to fix. This has made this PR much bigger, but ~40% is tests.

On main:

local ok, status, payload = pcall(pcall, produce)
-- ok - boolean
-- status - boolean|string
-- payload - R1|string

This PR (with and without @return_overload):

local ok, status, payload = pcall(pcall, produce)
-- ok - true|false
-- status - string|true|false
-- payload - string|integer

[xuhuanzy]

You might consider using conditional generics to handle this instead of introducing new syntax.

[CppCXY]

You might consider using conditional generics to handle this instead of introducing new syntax.

I think we should stick to simple rules. The current complex generics are immature and need another round of refactoring and rethinking. We can adopt @return_overload, but I don't want to have to write @return again after using @return_overload.

[xuhuanzy]

After analysis, for conditional generics to work as intended, built-in generics would have to be introduced. In that case, there’s essentially no fundamental difference from @return_overload, so using @return_overload is a viable approach.

[CppCXY]

All approaches based on @return have a problem: because multiple @return tags are allowed, the semantics become completely ambiguous when more than one @return appears. I suggest keeping return_overload unchanged.