fix(flow): guard correlated return-overload reentry

lewis6991 · lewis6991 · commit 81ac6be37e02 · 2026-03-20T12:46:15.000Z
Correlated @return_overload narrowing could recursively re-enter itself
for the same target return slot while the first inference was still in
flight.

The failure mode is:
- get_type_at_flow(target) reaches a condition node
- condition narrowing sees that the discriminant is a sibling return
  value
- narrow_var_from_return_overload_condition tries to recover the
  target's antecedent and uncorrelated fallback types via
  get_type_at_flow(target, ...)
- that predecessor walk can hit another condition on the same
  discriminant before the first correlated query has unwound
- the same target VarRefId is correlated again, repeating the cycle
  until indexing appears hung or the stack overflows

This is not just extra work. The algorithm was missing an in-flight
recursion break for correlated condition analysis, so the same
unresolved question could be asked again before any cache entry existed.

Add a small correlated_condition_guard keyed by VarRefId in
LuaInferCache and use it only around
narrow_var_from_return_overload_condition. Re-entry now returns
Continue, letting the existing flow walk fall back to non-correlated
inference instead of recursively opening the same correlated query.

The regression test is a focused Lua block with a single pick() call and
many repeated if-not-ok guards. Without this fix that block overflows
the test stack. With the guard it builds a semantic model and the rest
of the return-overload suite still passes.
diff --git a/crates/emmylua_code_analysis/src/compilation/test/return_overload_flow_test.rs b/crates/emmylua_code_analysis/src/compilation/test/return_overload_flow_test.rs
@@ -2,6 +2,8 @@
 mod test {
     use crate::{DiagnosticCode, VirtualWorkspace};
 
+    const STACKED_CORRELATED_GUARDS: usize = 180;
+
     #[test]
     fn test_return_overload_narrow_after_not() {
         let mut ws = VirtualWorkspace::new();
@@ -541,4 +543,46 @@ mod test {
         assert!(after_guard.contains("integer"));
         assert!(!after_guard.contains("string"));
     }
+
+    #[test]
+    fn test_return_overload_stacked_same_discriminant_guards_build_semantic_model() {
+        let mut ws = VirtualWorkspace::new();
+        let repeated_guards =
+            "if not ok then error(result) end\n".repeat(STACKED_CORRELATED_GUARDS);
+        let block = format!(
+            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")
+
+            {repeated_guards}
+            narrowed = result
+            "#,
+        );
+
+        let file_id = ws.def(&block);
+
+        assert!(
+            ws.analysis
+                .compilation
+                .get_semantic_model(file_id)
+                .is_some(),
+            "expected semantic model for stacked correlated-guard repro"
+        );
+        assert_eq!(ws.expr_ty("narrowed"), ws.ty("integer"));
+    }
 }
diff --git a/crates/emmylua_code_analysis/src/semantic/cache/mod.rs b/crates/emmylua_code_analysis/src/semantic/cache/mod.rs
@@ -23,6 +23,7 @@ pub struct LuaInferCache {
     pub call_cache:
         HashMap<(LuaSyntaxId, Option<usize>, LuaType), CacheEntry<Arc<LuaFunctionType>>>,
     pub flow_node_cache: HashMap<(VarRefId, FlowId), CacheEntry<LuaType>>,
+    pub correlated_condition_guard: HashSet<VarRefId>,
     pub index_ref_origin_type_cache: HashMap<VarRefId, CacheEntry<LuaType>>,
     pub expr_var_ref_id_cache: HashMap<LuaSyntaxId, VarRefId>,
     pub narrow_by_literal_stop_position_cache: HashSet<LuaSyntaxId>,
@@ -36,6 +37,7 @@ impl LuaInferCache {
             expr_cache: HashMap::new(),
             call_cache: HashMap::new(),
             flow_node_cache: HashMap::new(),
+            correlated_condition_guard: HashSet::new(),
             index_ref_origin_type_cache: HashMap::new(),
             expr_var_ref_id_cache: HashMap::new(),
             narrow_by_literal_stop_position_cache: HashSet::new(),
@@ -58,6 +60,7 @@ impl LuaInferCache {
         self.expr_cache.clear();
         self.call_cache.clear();
         self.flow_node_cache.clear();
+        self.correlated_condition_guard.clear();
         self.index_ref_origin_type_cache.clear();
         self.expr_var_ref_id_cache.clear();
     }
diff --git a/crates/emmylua_code_analysis/src/semantic/infer/narrow/condition_flow/correlated_flow.rs b/crates/emmylua_code_analysis/src/semantic/infer/narrow/condition_flow/correlated_flow.rs
@@ -1,4 +1,7 @@
-use std::collections::HashSet;
+use std::{
+    collections::HashSet,
+    panic::{AssertUnwindSafe, catch_unwind, resume_unwind},
+};
 
 use emmylua_parser::{LuaAstPtr, LuaCallExpr, LuaChunk};
 
@@ -23,76 +26,90 @@ pub(in crate::semantic::infer::narrow::condition_flow) fn narrow_var_from_return
     condition_position: rowan::TextSize,
     narrowed_discriminant_type: &LuaType,
 ) -> Result<ResultTypeOrContinue, InferFailReason> {
-    let Some(target_decl_id) = var_ref_id.get_decl_id_ref() else {
-        return Ok(ResultTypeOrContinue::Continue);
-    };
-    if !tree.has_decl_multi_return_refs(&discriminant_decl_id)
-        || !tree.has_decl_multi_return_refs(&target_decl_id)
-    {
+    let guard_key = var_ref_id.clone();
+    if !cache.correlated_condition_guard.insert(guard_key.clone()) {
         return Ok(ResultTypeOrContinue::Continue);
     }
 
-    let antecedent_flow_id = get_single_antecedent(tree, flow_node)?;
-    let search_root_flow_ids = tree.get_decl_multi_return_search_roots(
-        &discriminant_decl_id,
-        &target_decl_id,
-        condition_position,
-        antecedent_flow_id,
-    );
-    let mut matching_target_types = Vec::new();
-    let mut uncorrelated_target_types = Vec::new();
-    for search_root_flow_id in search_root_flow_ids {
-        let (root_matching_target_types, root_uncorrelated_target_type) =
-            collect_correlated_types_from_search_root(
-                db,
-                tree,
-                cache,
-                root,
-                var_ref_id,
-                discriminant_decl_id,
-                target_decl_id,
-                condition_position,
-                search_root_flow_id,
-                narrowed_discriminant_type,
-            )?;
-        matching_target_types.extend(root_matching_target_types);
-        if let Some(root_uncorrelated_target_type) = root_uncorrelated_target_type {
-            uncorrelated_target_types.push(root_uncorrelated_target_type);
+    let result = catch_unwind(AssertUnwindSafe(|| {
+        let Some(target_decl_id) = var_ref_id.get_decl_id_ref() else {
+            return Ok(ResultTypeOrContinue::Continue);
+        };
+        if !tree.has_decl_multi_return_refs(&discriminant_decl_id)
+            || !tree.has_decl_multi_return_refs(&target_decl_id)
+        {
+            return Ok(ResultTypeOrContinue::Continue);
         }
-    }
 
-    if matching_target_types.is_empty() {
-        return Ok(ResultTypeOrContinue::Continue);
-    }
+        let antecedent_flow_id = get_single_antecedent(tree, flow_node)?;
+        let search_root_flow_ids = tree.get_decl_multi_return_search_roots(
+            &discriminant_decl_id,
+            &target_decl_id,
+            condition_position,
+            antecedent_flow_id,
+        );
+        let mut matching_target_types = Vec::new();
+        let mut uncorrelated_target_types = Vec::new();
+        for search_root_flow_id in search_root_flow_ids {
+            let (root_matching_target_types, root_uncorrelated_target_type) =
+                collect_correlated_types_from_search_root(
+                    db,
+                    tree,
+                    cache,
+                    root,
+                    var_ref_id,
+                    discriminant_decl_id,
+                    target_decl_id,
+                    condition_position,
+                    search_root_flow_id,
+                    narrowed_discriminant_type,
+                )?;
+            matching_target_types.extend(root_matching_target_types);
+            if let Some(root_uncorrelated_target_type) = root_uncorrelated_target_type {
+                uncorrelated_target_types.push(root_uncorrelated_target_type);
+            }
+        }
 
-    let matching_target_type = LuaType::from_vec(matching_target_types);
-    let antecedent_type = get_type_at_flow(db, tree, cache, root, var_ref_id, antecedent_flow_id)?;
-    let narrowed_correlated_type =
-        TypeOps::Intersect.apply(db, &antecedent_type, &matching_target_type);
-    if narrowed_correlated_type.is_never() {
-        return Ok(ResultTypeOrContinue::Continue);
-    }
+        if matching_target_types.is_empty() {
+            return Ok(ResultTypeOrContinue::Continue);
+        }
 
-    if uncorrelated_target_types.is_empty() {
-        return Ok(if narrowed_correlated_type == antecedent_type {
-            ResultTypeOrContinue::Continue
+        let matching_target_type = LuaType::from_vec(matching_target_types);
+        let antecedent_type =
+            get_type_at_flow(db, tree, cache, root, var_ref_id, antecedent_flow_id)?;
+        let narrowed_correlated_type =
+            TypeOps::Intersect.apply(db, &antecedent_type, &matching_target_type);
+        if narrowed_correlated_type.is_never() {
+            return Ok(ResultTypeOrContinue::Continue);
+        }
+
+        if uncorrelated_target_types.is_empty() {
+            return Ok(if narrowed_correlated_type == antecedent_type {
+                ResultTypeOrContinue::Continue
+            } else {
+                ResultTypeOrContinue::Result(narrowed_correlated_type)
+            });
+        }
+
+        let uncorrelated_target_type = LuaType::from_vec(uncorrelated_target_types);
+        let merged_type = if uncorrelated_target_type.is_never() {
+            narrowed_correlated_type
         } else {
-            ResultTypeOrContinue::Result(narrowed_correlated_type)
-        });
-    }
+            LuaType::from_vec(vec![narrowed_correlated_type, uncorrelated_target_type])
+        };
 
-    let uncorrelated_target_type = LuaType::from_vec(uncorrelated_target_types);
-    let merged_type = if uncorrelated_target_type.is_never() {
-        narrowed_correlated_type
-    } else {
-        LuaType::from_vec(vec![narrowed_correlated_type, uncorrelated_target_type])
-    };
+        Ok(if merged_type == antecedent_type {
+            ResultTypeOrContinue::Continue
+        } else {
+            ResultTypeOrContinue::Result(merged_type)
+        })
+    }));
 
-    Ok(if merged_type == antecedent_type {
-        ResultTypeOrContinue::Continue
-    } else {
-        ResultTypeOrContinue::Result(merged_type)
-    })
+    cache.correlated_condition_guard.remove(&guard_key);
+    match result {
+        Ok(result) => result,
+        Err(payload) => resume_unwind(payload),
+    }
 }
 
 #[allow(clippy::too_many_arguments)]
