fix: grow the ts_lsp AST walk stacks so wide files keep every declaration

internal/cbm/extract_node_stack.h

@@ -34,18 +34,26 @@ typedef struct {
     CBMArena *scratch;
 } TSNodeStack;
 
+/* Initialize a stack against an arena directly. The LSP passes hold a
+ * TSLSPContext, not a CBMExtractCtx, and every allocation in those files
+ * already comes from that one arena — see collect_children in ts_lsp.c. An
+ * extraction caller should still use ts_nstack_init below, whose context
+ * argument is what makes handing over the wrong arena a type error. */
+static inline void ts_nstack_init_arena(TSNodeStack *s, CBMArena *arena, int initial_cap) {
+    s->scratch = arena;
+    s->items = (TSNode *)cbm_arena_alloc(arena, (size_t)initial_cap * sizeof(TSNode));
+    s->count = 0;
+    s->cap = s->items ? initial_cap : 0;
+}
+
 /* Initialize a stack with the given initial capacity, allocated from the
  * context's traversal scratch. Taking the context rather than an arena is
  * deliberate: it makes handing over ctx->arena, or a local alias of it, a type
  * error rather than a retention bug nobody notices. A context built without a
  * scratch falls back to ctx->arena, which is the behaviour that shipped before
  * #1997, so no caller ever gets a NULL arena and silently loses nodes. */
 static inline void ts_nstack_init(TSNodeStack *s, const CBMExtractCtx *ctx, int initial_cap) {
-    CBMArena *arena = ctx->scratch ? ctx->scratch : ctx->arena;
-    s->scratch = arena;
-    s->items = (TSNode *)cbm_arena_alloc(arena, (size_t)initial_cap * sizeof(TSNode));
-    s->count = 0;
-    s->cap = s->items ? initial_cap : 0;
+    ts_nstack_init_arena(s, ctx->scratch ? ctx->scratch : ctx->arena, initial_cap);
 }
 
 /* Push a node onto the stack, growing 2x if needed. */

internal/cbm/lsp/ts_lsp.c

@@ -27,6 +27,7 @@
  */
 
 #include "ts_lsp.h"
+#include "extract_node_stack.h"
 #include <stdatomic.h>
 #include <stdio.h>
 #include <stdlib.h>
@@ -4634,13 +4635,15 @@ static const CBMType *infer_return_type_from_body(TSLSPContext *ctx, TSNode body
     if (ts_node_is_null(body))
         return cbm_type_unknown();
 
-    // Iterative DFS using a small fixed-size stack to avoid C-stack blowup on big bodies.
-    enum { STACK_CAP = 256 };
-    TSNode stack[STACK_CAP];
-    int top = 0;
-    stack[top++] = body;
-    while (top > 0) {
-        TSNode n = stack[--top];
+    // Iterative DFS to avoid C-stack blowup on big bodies. The stack grows, so a
+    // body wider than any fixed cap is still walked in full: a fixed 256 slots
+    // used to stop pushing without a word, and this function then answered
+    // "unknown return type" for a body whose return it had never reached.
+    TSNodeStack stack;
+    ts_nstack_init_arena(&stack, ctx->arena, 256);
+    ts_nstack_push(&stack, body);
+    while (stack.count > 0) {
+        TSNode n = ts_nstack_pop(&stack);
         if (ts_node_is_null(n))
             continue;
         const char *k = ts_node_type(n);
@@ -4662,10 +4665,7 @@ static const CBMType *infer_return_type_from_body(TSLSPContext *ctx, TSNode body
             strcmp(k, "arrow_function") == 0 || strcmp(k, "method_definition") == 0)
             continue;
 
-        uint32_t cnt = ts_node_child_count(n);
-        for (uint32_t i = 0; i < cnt && top < STACK_CAP; i++) {
-            stack[top++] = ts_node_child(n, i);
-        }
+        ts_nstack_push_children(&stack, n);
     }
     return cbm_type_unknown();
 }
@@ -4681,31 +4681,29 @@ static void rebuild_signatures_from_ast(TSLSPContext *ctx, TSNode root, CBMTypeR
     if (ts_node_is_null(root) || !reg || !ctx->module_qn)
         return;
 
-    enum { STACK_CAP = 256 };
-    TSNode stack[STACK_CAP];
-    int top = 0;
-    uint32_t nc = ts_node_child_count(root);
-    for (uint32_t i = 0; i < nc && top < STACK_CAP; i++)
-        stack[top++] = ts_node_child(root, i);
+    /* The stack grows. A fixed 256 slots used to be filled by the seeding loop
+     * alone on a file with that many top-level children, and every declaration
+     * past it was then skipped without a word — it kept whatever signature
+     * extract_defs had guessed. */
+    TSNodeStack stack;
+    ts_nstack_init_arena(&stack, ctx->arena, 256);
+    ts_nstack_push_children(&stack, root);
 
-    while (top > 0) {
-        TSNode n = stack[--top];
+    while (stack.count > 0) {
+        TSNode n = ts_nstack_pop(&stack);
         if (ts_node_is_null(n))
             continue;
         const char *k = ts_node_type(n);
 
         // Recurse into export_statement and class bodies.
         if (strcmp(k, "export_statement") == 0 || strcmp(k, "class_body") == 0) {
-            uint32_t cnt = ts_node_child_count(n);
-            for (uint32_t i = 0; i < cnt && top < STACK_CAP; i++) {
-                stack[top++] = ts_node_child(n, i);
-            }
+            ts_nstack_push_children(&stack, n);
             continue;
         }
         if (strcmp(k, "class_declaration") == 0) {
             TSNode body = ts_node_child_by_field_name(n, "body", TS_LSP_FIELD_LEN("body"));
-            if (!ts_node_is_null(body) && top < STACK_CAP)
-                stack[top++] = body;
+            if (!ts_node_is_null(body))
+                ts_nstack_push(&stack, body);
             continue;
         }
 
@@ -4845,31 +4843,29 @@ static void convert_signature_type_params(TSLSPContext *ctx, TSNode root, CBMTyp
         return;
 
     // Walk: function_declaration, class_declaration { method_definition }, plus exported.
-    enum { STACK_CAP = 256 };
-    TSNode stack[STACK_CAP];
-    int top = 0;
-    uint32_t nc = ts_node_child_count(root);
-    for (uint32_t i = 0; i < nc && top < STACK_CAP; i++)
-        stack[top++] = ts_node_child(root, i);
-
-    while (top > 0) {
-        TSNode n = stack[--top];
+    /* The stack grows. A fixed 256 slots used to be filled by the seeding loop
+     * alone on a file with that many top-level children, and every declaration
+     * past it was then skipped without a word — it kept whatever signature
+     * extract_defs had guessed. */
+    TSNodeStack stack;
+    ts_nstack_init_arena(&stack, ctx->arena, 256);
+    ts_nstack_push_children(&stack, root);
+
+    while (stack.count > 0) {
+        TSNode n = ts_nstack_pop(&stack);
         if (ts_node_is_null(n))
             continue;
         const char *k = ts_node_type(n);
 
         // Recurse into export_statement and class_body.
         if (strcmp(k, "export_statement") == 0 || strcmp(k, "class_body") == 0) {
-            uint32_t cnt = ts_node_child_count(n);
-            for (uint32_t i = 0; i < cnt && top < STACK_CAP; i++) {
-                stack[top++] = ts_node_child(n, i);
-            }
+            ts_nstack_push_children(&stack, n);
             continue;
         }
         if (strcmp(k, "class_declaration") == 0) {
             TSNode body = ts_node_child_by_field_name(n, "body", TS_LSP_FIELD_LEN("body"));
-            if (!ts_node_is_null(body) && top < STACK_CAP)
-                stack[top++] = body;
+            if (!ts_node_is_null(body))
+                ts_nstack_push(&stack, body);
             continue;
         }
 

tests/test_ts_lsp.c

@@ -2817,6 +2817,64 @@ TEST(tslsp_stress_many_classes) {
     PASS();
 }
 
+TEST(tslsp_stress_declarations_past_stack_cap) {
+    /* 300 padding functions, then the generic pair that needs the AST walk.
+     * The walk used to hold 256 nodes at most, so every top-level declaration
+     * after the 256th was dropped without a word and kept whatever signature
+     * the first extraction pass had guessed. */
+    enum { BUF_CAP = 64 * 1024 };
+    char *buf = (char *)malloc(BUF_CAP);
+    if (!buf)
+        PASS();
+    char *p = buf;
+    char *end = buf + BUF_CAP;
+    for (int i = 0; i < 300; i++) {
+        p += snprintf(p, (size_t)(end - p), "function pad%d(): void {}\n", i);
+    }
+    p += snprintf(p, (size_t)(end - p), "function identity<T>(x: T): T { return x; }\n");
+    p += snprintf(p, (size_t)(end - p), "class Box { use(): void {} }\n");
+    p += snprintf(p, (size_t)(end - p),
+                  "function go() { const b = identity(new Box()); b.use(); }\n");
+    *p = '\0';
+
+    CBMFileResult *r = extract_ts(buf);
+    free(buf);
+    ASSERT_NOT_NULL(r);
+    /* Same claim as tslsp_generic_identity_inference, 300 declarations deeper. */
+    ASSERT_GTE(require_resolved(r, ".go", "use"), 0);
+    cbm_free_result(r);
+    PASS();
+}
+
+TEST(tslsp_stress_return_past_stack_cap) {
+    /* infer_return_type_from_body reads the return type out of a body that
+     * carries no written one. The return here sits after 300 statements, more
+     * than the old fixed cap held, so the walk used to stop before reaching it
+     * and answer "unknown return type". */
+    enum { BUF_CAP = 64 * 1024 };
+    char *buf = (char *)malloc(BUF_CAP);
+    if (!buf)
+        PASS();
+    char *p = buf;
+    char *end = buf + BUF_CAP;
+    p += snprintf(p, (size_t)(end - p), "class Conn { ping(): void {} }\n");
+    p += snprintf(p, (size_t)(end - p), "function makeConn() {\n");
+    for (int i = 0; i < 300; i++) {
+        p += snprintf(p, (size_t)(end - p), "  let a%d = 0;\n", i);
+    }
+    p += snprintf(p, (size_t)(end - p), "  return new Conn();\n}\n");
+    p += snprintf(p, (size_t)(end - p), "function go() { const c = makeConn(); c.ping(); }\n");
+    *p = '\0';
+
+    CBMFileResult *r = extract_ts(buf);
+    free(buf);
+    ASSERT_NOT_NULL(r);
+    /* Same claim as tslsp_return_inferred_local, 300 statements deeper. */
+    ASSERT_GTE(require_resolved(r, ".go", "ping"), 0);
+    cbm_free_result(r);
+    PASS();
+}
+
 TEST(tslsp_stress_deep_inheritance) {
     /* Chain of 30 classes via extends, leaf method on root. */
     char buf[16 * 1024];
@@ -4565,6 +4623,8 @@ SUITE(ts_lsp) {
 
     /* Stress tests */
     RUN_TEST(tslsp_stress_many_classes);
+    RUN_TEST(tslsp_stress_declarations_past_stack_cap);
+    RUN_TEST(tslsp_stress_return_past_stack_cap);
     RUN_TEST(tslsp_stress_deep_inheritance);
     RUN_TEST(tslsp_stress_long_method_chain);
     RUN_TEST(tslsp_stress_megafile_mixed);