Fix exponential flow tracing in files with loops

crates/emmylua_code_analysis/src/semantic_db/flow/flow_tree.rs

@@ -42,6 +42,11 @@ pub struct FlowTree {
     /// backward edges whose source positions are not monotonic, so the member
     /// assignment-position fast path is conservatively disabled.
     has_complex_control: bool,
+    /// Per node: whether a loop / label / break / continue node can reach it by
+    /// following antecedent edges (computed by forward propagation from those
+    /// nodes). Unmarked nodes can memoize trace states even when another function
+    /// in the same file has loops.
+    loop_reachable: Vec<bool>,
 }
 
 impl FlowTree {
@@ -71,6 +76,11 @@ impl FlowTree {
                     | FlowNodeKind::Continue
             )
         });
+        let loop_reachable = if has_complex_control {
+            Self::compute_loop_reachable(&flow_nodes, &multiple_antecedents)
+        } else {
+            vec![false; flow_nodes.len()]
+        };
         for node in &flow_nodes {
             let flow_id = node.id;
             if matches!(node.antecedent, Some(FlowAntecedent::Multiple(_))) {
@@ -155,7 +165,68 @@ impl FlowTree {
             branch_nodes,
             unpositioned_branch,
             has_complex_control,
+            loop_reachable,
+        }
+    }
+
+    /// Marks every node that a loop / label / break / continue node can reach by
+    /// following antecedent edges. Computed by propagating forward from those nodes
+    /// along successor edges, so each node is visited once.
+    fn compute_loop_reachable(
+        flow_nodes: &[FlowNode],
+        multiple_antecedents: &[Vec<FlowId>],
+    ) -> Vec<bool> {
+        let mut reachable = vec![false; flow_nodes.len()];
+        let mut successors: Vec<Vec<u32>> = vec![Vec::new(); flow_nodes.len()];
+        let mut queue = VecDeque::new();
+        for node in flow_nodes {
+            let id = node.id.0;
+            let antecedents: &[FlowId] = match &node.antecedent {
+                Some(FlowAntecedent::Single(prev)) => std::slice::from_ref(prev),
+                Some(FlowAntecedent::Multiple(multi_id)) => multiple_antecedents
+                    .get(*multi_id as usize)
+                    .map(|v| v.as_slice())
+                    .unwrap_or_default(),
+                None => &[],
+            };
+            for prev in antecedents {
+                debug_assert!(
+                    (prev.0 as usize) < flow_nodes.len(),
+                    "flow antecedent {} is out of range",
+                    prev.0
+                );
+                if let Some(list) = successors.get_mut(prev.0 as usize) {
+                    list.push(id);
+                }
+            }
+            if matches!(
+                node.kind,
+                FlowNodeKind::LoopLabel
+                    | FlowNodeKind::NamedLabel(_)
+                    | FlowNodeKind::ForIStat(_)
+                    | FlowNodeKind::Break
+                    | FlowNodeKind::Continue
+            ) && let Some(slot) = reachable.get_mut(id as usize)
+                && !*slot
+            {
+                *slot = true;
+                queue.push_back(id);
+            }
+        }
+        while let Some(id) = queue.pop_front() {
+            let Some(next_ids) = successors.get(id as usize) else {
+                continue;
+            };
+            for &next in next_ids {
+                if let Some(slot) = reachable.get_mut(next as usize)
+                    && !*slot
+                {
+                    *slot = true;
+                    queue.push_back(next);
+                }
+            }
         }
+        reachable
     }
 
     pub fn get_flow_id(&self, syntax_id: LuaSyntaxId) -> Option<FlowId> {
@@ -271,6 +342,16 @@ impl FlowTree {
         self.has_complex_control
     }
 
+    /// Whether a loop / label / break / continue node can reach `flow_id` by
+    /// following antecedent edges.
+    pub fn is_loop_reachable(&self, flow_id: FlowId) -> bool {
+        // An unknown node stays on the conservative, uncached path.
+        self.loop_reachable
+            .get(flow_id.0 as usize)
+            .copied()
+            .unwrap_or(true)
+    }
+
     pub fn get_multi_antecedents(&self, id: u32) -> Option<&[FlowId]> {
         self.multiple_antecedents
             .get(id as usize)

crates/emmylua_code_analysis/src/semantic_db/phase_tests/flow_loop_memo_tests.rs

@@ -0,0 +1,105 @@
+//! Regression guards for flow backtracking in files that contain a loop.
+//!
+//! One loop anywhere in a file disabled trace-state memoization for every
+//! function in it, so a chain of `n` guards in a loop-free function walked
+//! ~2^n routes (`walking.lua` of a real game client project never finished).
+
+use emmylua_parser::{LuaAstNode, LuaNameExpr};
+use rowan::TextSize;
+
+use crate::semantic_model::flow::flow_metrics;
+use crate::{FileId, VirtualWorkspace};
+
+/// Every loop statement the binder knows; each one used to turn the memo off for
+/// the whole file.
+const LOOP_KINDS: &[(&str, &str)] = &[
+    ("for-in", "for _ in pairs(t) do end"),
+    ("numeric for", "for i = 1, #t do end"),
+    ("while", "while t do t = nil end"),
+    ("repeat", "repeat t = nil until not t"),
+];
+
+/// `n` sequential guards in `walk`, plus an unrelated `loop_stat` in `other`.
+fn guard_chain_with_loop_elsewhere(n: usize, loop_stat: &str) -> String {
+    let mut content = String::from("lastDir = nil\nwalkLock = 0\n\nfunction walk(dir)\n");
+    for i in 0..n {
+        content.push_str(&format!(
+            "  if walkLock > {i} and lastDir == dir then\n    walkLock = walkLock + 1\n  end\n"
+        ));
+    }
+    content.push_str(&format!(
+        "  lastDir = dir\nend\n\nfunction other(t)\n  {loop_stat}\nend\n"
+    ));
+    content
+}
+
+const BACKWARD_GOTO: &str = "
+function retry(t)
+  ::again::
+  walkLock = 0
+  if t then goto again end
+end
+";
+
+/// Read the flow type of every name, the way checkers do for each reference.
+fn read_every_name(ws: &VirtualWorkspace, file_id: FileId) {
+    let model = ws.analysis.semantic_model(file_id);
+    for name in model.chunk().expect("chunk").descendants::<LuaNameExpr>() {
+        let _ = model.type_of_expr_at(name.get_syntax_id(), name.get_position());
+    }
+}
+
+/// Only nodes with a real cycle behind them (a backward `goto`) lose the memo; a loop
+/// or label elsewhere in the file no longer affects `walk`.
+#[test]
+fn flow_loop_reachability_is_tracked_per_node() {
+    let mut ws = VirtualWorkspace::new();
+    let source = guard_chain_with_loop_elsewhere(2, LOOP_KINDS[0].1) + BACKWARD_GOTO;
+    let file_id = ws.def_file("walking.lua", &source);
+    let model = ws.analysis.semantic_model(file_id);
+    let tree = model.flow_tree().expect("flow tree");
+    let flow_at = |needle: &str| {
+        let offset = source.find(needle).expect("needle in source") as u32;
+        tree.get_flow_id_at(TextSize::from(offset))
+            .unwrap_or_else(|| panic!("no flow node at `{needle}` (offset {offset})"))
+    };
+
+    assert!(tree.has_complex_control());
+    assert!(
+        !tree.is_loop_reachable(flow_at("lastDir = dir")),
+        "`walk` has no loop behind it, so its trace states stay memoizable"
+    );
+    assert!(
+        tree.is_loop_reachable(flow_at("walkLock = 0\n  if t")),
+        "a statement behind the backward `goto` must keep the uncached path"
+    );
+}
+
+/// The number of uncached walks must stay linear in the number of guards for
+/// every loop kind in another function: `16n + 3` (`for`) or `16n + 4` (`while`,
+/// `repeat`) with the per-node gate. With the per-file gate it was already 631 at
+/// n=4 and `[(8, 52455), (16, 2065146), (32, 6638442)]` after that, so the check
+/// runs per size and fails on the first one instead of walking the larger chains.
+#[test]
+fn flow_loop_in_other_function_keeps_guard_chain_linear() {
+    for (kind, loop_stat) in LOOP_KINDS {
+        let mut measurements = Vec::new();
+        for n in [4usize, 8, 16, 32] {
+            let mut ws = VirtualWorkspace::new();
+            let source = guard_chain_with_loop_elsewhere(n, loop_stat);
+            let file_id = ws.def_file("walking.lua", &source);
+            flow_metrics::reset();
+            read_every_name(&ws, file_id);
+            let entries = flow_metrics::trace_decl_entries();
+            measurements.push((n, entries));
+
+            let bound = 32 * (n as u64) + 64;
+            assert!(
+                entries <= bound,
+                "the declaration walk ran {entries} times for {n} guards (bound {bound}); \
+                 a `{kind}` loop in another function disabled the trace-state memo. \
+                 So far: {measurements:?}"
+            );
+        }
+    }
+}

crates/emmylua_code_analysis/src/semantic_db/phase_tests/mod.rs

@@ -1,4 +1,5 @@
 mod fix_tests;
+mod flow_loop_memo_tests;
 mod module_roots_tests;
 mod p0_tests;
 mod p10_tests;

crates/emmylua_code_analysis/src/semantic_model/flow.rs

@@ -1502,9 +1502,10 @@ fn trace_decl(
     // a state's value depends on the node and the options alone — never on the
     // route. That makes the state cacheable, which matters because a nested
     // conditional chain otherwise re-derives the same states along exponentially
-    // many routes. Graphs with loops keep the uncached path (see the type note on
-    // `cache::TraceStateKey`).
-    if tree.has_complex_control() {
+    // many routes. Only nodes that can reach a loop backwards keep the uncached
+    // path (see the type note on `cache::TraceStateKey`), so a loop in one function
+    // does not disable memoization for the rest of the file.
+    if tree.is_loop_reachable(flow_id) {
         return trace_decl_uncached(model, decl, tree, flow_id, options, mode, visited, path);
     }
     let key = cache::TraceStateKey {