fix(lsp): grow the ts_lsp walk stacks with cbm_walk_stack_reserve

Follow-up to the review of #2091, on top of the contributor's change.

The three ts_lsp.c walks (infer_return_type_from_body,
rebuild_signatures_from_ast, convert_signature_type_params) now keep
their first 256 slots on the C stack, as the old fixed arrays did, and
grow on the heap through cbm_walk_stack_reserve / cbm_walk_stack_release
(#2375) instead of cutting TSNodeStack storage from ctx->arena:

- Scratch lifetime: ctx->arena in the per-file pass is the result
  arena, kept until the result cache is freed, and the return walk runs
  once per untyped top-level function. The heap chunk is now released
  when each walk ends.
- No silent drop: when the stack cannot grow, cbm_walk_stack_reserve
  logs extract.walk_stack_alloc_failed and the walk stops. The return
  walk then answers "unknown"; the signature walks leave the
  declarations not yet reached as extract_defs registered them.
- Traversal order: children are filled in cursor order (still one O(N)
  pass) and popped last-child first, exactly as main's walks did, so of
  several returns the later one is still tried first.
  tslsp_return_walk_order_prefers_later_return pins it; it fails under
  a first-child-first walk.

The TSNodeStack arena initializer in extract_node_stack.h is no longer
needed and is dropped, so that header is unchanged from main. The two
stack-cap tests now fail instead of passing when their buffer
allocation fails.

Co-authored-by: Joshua Richter <jrichter5781@gmail.com>
Signed-off-by: Martin Vogel <martin.vogel.tech@gmail.com>

internal/cbm/extract_node_stack.h

@@ -34,18 +34,6 @@ 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
@@ -53,7 +41,11 @@ static inline void ts_nstack_init_arena(TSNodeStack *s, CBMArena *arena, int ini
  * 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) {
-    ts_nstack_init_arena(s, ctx->scratch ? ctx->scratch : ctx->arena, 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;
 }
 
 /* Push a node onto the stack, growing 2x if needed. */

internal/cbm/lsp/ts_lsp.c

@@ -28,7 +28,7 @@
 
 #include "ts_lsp.h"
 #include "../../../src/foundation/platform.h"
-#include "extract_node_stack.h"
+#include "../helpers.h"
 #include <stdatomic.h>
 #include <stdio.h>
 #include <stdlib.h>
@@ -4685,21 +4685,61 @@ static void register_file_defs(CBMArena *arena, CBMTypeRegistry *reg, CBMFileRes
     }
 }
 
-// Find the first `return EXPR` inside a body and evaluate EXPR's type. Returns
-// UNKNOWN if the body has no return or every return is unresolvable.
+// AST walk stacks for the passes below. The first chunk lives on the C stack, the
+// same 256 slots the old fixed arrays used; a wider pending set grows on the heap
+// through cbm_walk_stack_reserve and is released when the walk ends, so nothing
+// is cut from the per-file result arena. A fixed cap used to stop pushing without
+// a word, which skipped every top-level declaration past the 256th and left the
+// return walk short of returns it had never reached (#2388).
+enum { TS_WALK_STACK_INLINE = 256 };
+
+// Push every child of n in cursor order (one O(N) pass). The walk pops the stack,
+// so children are visited last-child first: the order these walks have always
+// used. False means the stack could not grow; cbm_walk_stack_reserve has logged it.
+static bool ts_walk_push_children(TSNode **stack, int *cap, int *top, TSNode n,
+                                  const TSNode *inline_buf, const char *walker) {
+    uint32_t count = ts_node_child_count(n);
+    if (count == 0)
+        return true;
+    if (!cbm_walk_stack_reserve((void **)stack, cap, *top + (int)count, sizeof(TSNode), inline_buf,
+                                walker))
+        return false;
+    TSTreeCursor cursor = ts_tree_cursor_new(n);
+    if (ts_tree_cursor_goto_first_child(&cursor)) {
+        uint32_t pushed = 0;
+        do {
+            (*stack)[(*top)++] = ts_tree_cursor_current_node(&cursor);
+            pushed++;
+        } while (pushed < count && ts_tree_cursor_goto_next_sibling(&cursor));
+    }
+    ts_tree_cursor_delete(&cursor);
+    return true;
+}
+
+static bool ts_walk_push(TSNode **stack, int *cap, int *top, TSNode n, const TSNode *inline_buf,
+                         const char *walker) {
+    if (!cbm_walk_stack_reserve((void **)stack, cap, *top + 1, sizeof(TSNode), inline_buf, walker))
+        return false;
+    (*stack)[(*top)++] = n;
+    return true;
+}
+
+// Find a `return EXPR` inside a body and evaluate EXPR's type. Children are visited
+// last-child first, so of several returns the one latest in the body is tried
+// first; the first one that resolves wins. Returns UNKNOWN if the body has no
+// return, every return is unresolvable, or the walk stack could not grow.
 static const CBMType *infer_return_type_from_body(TSLSPContext *ctx, TSNode body) {
     if (ts_node_is_null(body))
         return cbm_type_unknown();
 
-    // 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);
+    TSNode inline_stack[TS_WALK_STACK_INLINE];
+    TSNode *stack = inline_stack;
+    int cap = TS_WALK_STACK_INLINE;
+    int top = 0;
+    const CBMType *found = cbm_type_unknown();
+    stack[top++] = body;
+    while (top > 0) {
+        TSNode n = stack[--top];
         if (ts_node_is_null(n))
             continue;
         const char *k = ts_node_type(n);
@@ -4712,8 +4752,10 @@ static const CBMType *infer_return_type_from_body(TSLSPContext *ctx, TSNode body
             if (ts_node_is_null(expr))
                 continue;
             const CBMType *t = ts_eval_expr_type(ctx, expr);
-            if (t && !cbm_type_is_unknown(t))
-                return t;
+            if (t && !cbm_type_is_unknown(t)) {
+                found = t;
+                break;
+            }
             continue;
         }
         // Don't recurse into nested function bodies — their returns belong to them.
@@ -4721,9 +4763,11 @@ static const CBMType *infer_return_type_from_body(TSLSPContext *ctx, TSNode body
             strcmp(k, "arrow_function") == 0 || strcmp(k, "method_definition") == 0)
             continue;
 
-        ts_nstack_push_children(&stack, n);
+        if (!ts_walk_push_children(&stack, &cap, &top, n, inline_stack, "ts_lsp_return_walk"))
+            break; // "unknown" is the safe answer for a walk that could not finish
     }
-    return cbm_type_unknown();
+    cbm_walk_stack_release(stack, inline_stack);
+    return found;
 }
 
 // Rebuild a function's CBMRegisteredFunc signature from the AST directly. extract_defs
@@ -4737,29 +4781,33 @@ static void rebuild_signatures_from_ast(TSLSPContext *ctx, TSNode root, CBMTypeR
     if (ts_node_is_null(root) || !reg || !ctx->module_qn)
         return;
 
-    /* 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);
+    // Seeded with every top-level child, so a file with more than 256 of them
+    // needs the growing stack (see TS_WALK_STACK_INLINE). If it cannot grow, the
+    // walk stops and the declarations not yet reached keep extract_defs' signature.
+    TSNode inline_stack[TS_WALK_STACK_INLINE];
+    TSNode *stack = inline_stack;
+    int cap = TS_WALK_STACK_INLINE;
+    int top = 0;
+    bool walking =
+        ts_walk_push_children(&stack, &cap, &top, root, inline_stack, "ts_lsp_rebuild_signatures");
 
-    while (stack.count > 0) {
-        TSNode n = ts_nstack_pop(&stack);
+    while (walking && top > 0) {
+        TSNode n = stack[--top];
         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) {
-            ts_nstack_push_children(&stack, n);
+            walking = ts_walk_push_children(&stack, &cap, &top, n, inline_stack,
+                                            "ts_lsp_rebuild_signatures");
             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))
-                ts_nstack_push(&stack, body);
+                walking = ts_walk_push(&stack, &cap, &top, body, inline_stack,
+                                       "ts_lsp_rebuild_signatures");
             continue;
         }
 
@@ -4889,6 +4937,7 @@ static void rebuild_signatures_from_ast(TSLSPContext *ctx, TSNode root, CBMTypeR
             break;
         }
     }
+    cbm_walk_stack_release(stack, inline_stack);
 }
 
 // Convert NAMED references to TYPE_PARAM for any function whose declaration has a
@@ -4899,29 +4948,33 @@ static void convert_signature_type_params(TSLSPContext *ctx, TSNode root, CBMTyp
         return;
 
     // Walk: function_declaration, class_declaration { method_definition }, plus exported.
-    /* 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);
+    // Seeded with every top-level child, so a file with more than 256 of them needs
+    // the growing stack (see TS_WALK_STACK_INLINE). If it cannot grow, the walk stops
+    // and the generics not yet reached keep their NAMED parameters.
+    TSNode inline_stack[TS_WALK_STACK_INLINE];
+    TSNode *stack = inline_stack;
+    int cap = TS_WALK_STACK_INLINE;
+    int top = 0;
+    bool walking =
+        ts_walk_push_children(&stack, &cap, &top, root, inline_stack, "ts_lsp_type_params");
 
-    while (stack.count > 0) {
-        TSNode n = ts_nstack_pop(&stack);
+    while (walking && top > 0) {
+        TSNode n = stack[--top];
         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) {
-            ts_nstack_push_children(&stack, n);
+            walking =
+                ts_walk_push_children(&stack, &cap, &top, n, inline_stack, "ts_lsp_type_params");
             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))
-                ts_nstack_push(&stack, body);
+                walking =
+                    ts_walk_push(&stack, &cap, &top, body, inline_stack, "ts_lsp_type_params");
             continue;
         }
 
@@ -5020,6 +5073,7 @@ static void convert_signature_type_params(TSLSPContext *ctx, TSNode root, CBMTyp
             break;
         }
     }
+    cbm_walk_stack_release(stack, inline_stack);
 }
 
 // Parse a JSDoc comment block for `@param {T} name` and `@returns {T}` tags. For each

tests/test_ts_lsp.c

@@ -310,6 +310,21 @@ TEST(tslsp_return_inferred_local) {
     PASS();
 }
 
+/* infer_return_type_from_body visits children last-child first, so of two returns
+ * the one later in the body is tried first. Pinned so that a change to the walk
+ * order (#2091) is a visible decision rather than a side effect. */
+TEST(tslsp_return_walk_order_prefers_later_return) {
+    CBMFileResult *r =
+        extract_ts("class A { a(): void {} }\n"
+                   "class B { b(): void {} }\n"
+                   "function pick(f) { if (f) { return new A(); } return new B(); }\n"
+                   "function go() { const x = pick(true); x.b(); }\n");
+    ASSERT_NOT_NULL(r);
+    ASSERT_GTE(require_resolved(r, ".go", "b"), 0);
+    cbm_free_result(r);
+    PASS();
+}
+
 TEST(tslsp_return_async_promise) {
     CBMFileResult *r = extract_ts("class User { hello(): void {} }\n"
                                   "async function load(): Promise<User> { return new User(); }\n"
@@ -3321,8 +3336,7 @@ TEST(tslsp_stress_declarations_past_stack_cap) {
      * the first extraction pass had guessed. */
     enum { BUF_CAP = 64 * 1024 };
     char *buf = (char *)malloc(BUF_CAP);
-    if (!buf)
-        PASS();
+    ASSERT_NOT_NULL(buf);
     char *p = buf;
     char *end = buf + BUF_CAP;
     for (int i = 0; i < 300; i++) {
@@ -3350,8 +3364,7 @@ TEST(tslsp_stress_return_past_stack_cap) {
      * and answer "unknown return type". */
     enum { BUF_CAP = 64 * 1024 };
     char *buf = (char *)malloc(BUF_CAP);
-    if (!buf)
-        PASS();
+    ASSERT_NOT_NULL(buf);
     char *p = buf;
     char *end = buf + BUF_CAP;
     p += snprintf(p, (size_t)(end - p), "class Conn { ping(): void {} }\n");
@@ -4799,6 +4812,7 @@ SUITE(ts_lsp) {
     RUN_TEST(tslsp_return_simple);
     RUN_TEST(tslsp_return_chain);
     RUN_TEST(tslsp_return_inferred_local);
+    RUN_TEST(tslsp_return_walk_order_prefers_later_return);
     RUN_TEST(tslsp_return_async_promise);
     RUN_TEST(tslsp_return_method_no_rttype);