Merge pull request #2091 from CaptainMittens/fix/ts-lsp-walk-stack-drops-nodes

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

internal/cbm/lsp/ts_lsp.c

@@ -20,7 +20,7 @@
  * cbm_batch_ts_lsp_cross.
  *
  * Conventions (mirrors c_lsp / go_lsp):
- *   - All allocations go through ctx->arena; no malloc / free.
+ *   - Persistent results use ctx->arena; temporary walk/index buffers use tracked allocations.
  *   - ts_emit_resolved_call requires a non-NULL enclosing_func_qn — calls outside any
  *     function (module top-level) are not emitted.
  *   - eval_depth caps at TS_LSP_MAX_EVAL_DEPTH to defend against pathological recursion.
@@ -28,6 +28,10 @@
 
 #include "ts_lsp.h"
 #include "../../../src/foundation/platform.h"
+#include "../helpers.h"
+#include "../../../src/foundation/log.h"
+#include "../../../src/foundation/mem_core.h"
+#include <limits.h>
 #include <stdatomic.h>
 #include <stdio.h>
 #include <stdlib.h>
@@ -60,6 +64,7 @@ void cbm_ts_full_registry_builds_reset(void) {
 static _Thread_local long g_ts_type_budget = -1;
 static _Thread_local bool g_ts_type_budget_warned;
 #ifdef CBM_ENABLE_TEST_SEAMS
+static _Thread_local unsigned long g_ts_signature_lookup_work;
 static _Thread_local int g_ts_max_member_lookup_steps;
 static _Thread_local int g_ts_max_method_lookup_steps;
 #endif
@@ -174,6 +179,9 @@ long cbm_ts_lsp_test_budget_remaining(void) {
 bool cbm_ts_lsp_test_budget_warned(void) {
     return g_ts_type_budget_warned;
 }
+unsigned long cbm_ts_lsp_test_signature_lookup_work(void) {
+    return g_ts_signature_lookup_work;
+}
 int cbm_ts_lsp_test_max_member_lookup_steps(void) {
     return g_ts_max_member_lookup_steps;
 }
@@ -3956,6 +3964,7 @@ void ts_lsp_init(TSLSPContext *ctx, CBMArena *arena, const char *source, int sou
     if (!ctx)
         return;
 #ifdef CBM_ENABLE_TEST_SEAMS
+    g_ts_signature_lookup_work = 0;
     g_ts_max_member_lookup_steps = 0;
     g_ts_max_method_lookup_steps = 0;
 #endif
@@ -4805,16 +4814,77 @@ 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 };
+
+static int ts_walk_pending_count(int top, uint32_t count, const char *walker) {
+    if (top < 0 || count > (uint32_t)(INT_MAX - top)) {
+        char pending[24];
+        snprintf(pending, sizeof(pending), "%llu", (unsigned long long)count + (unsigned)top);
+        cbm_log_error("extract.walk_stack_alloc_failed", "walker", walker, "pending", pending);
+        return -1;
+    }
+    return top + (int)count;
+}
+
+#ifdef CBM_ENABLE_TEST_SEAMS
+int cbm_ts_lsp_test_walk_pending_count(int top, uint32_t count) {
+    return ts_walk_pending_count(top, count, "ts_lsp_test");
+}
+#endif
+
+// 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;
+    int need = ts_walk_pending_count(*top, count, walker);
+    if (need < 0 ||
+        !cbm_walk_stack_reserve((void **)stack, cap, need, 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) {
+    int need = ts_walk_pending_count(*top, 1, walker);
+    if (need < 0 ||
+        !cbm_walk_stack_reserve((void **)stack, cap, need, 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 using a small fixed-size stack to avoid C-stack blowup on big bodies.
-    enum { STACK_CAP = 256 };
-    TSNode stack[STACK_CAP];
+    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];
@@ -4830,8 +4900,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.
@@ -4839,14 +4911,96 @@ 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);
-        }
+        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;
 }
 
+// These two AST passes patch the FIRST matching local registration, even when
+// the general registry index prefers the last duplicate. Do not follow fallback
+// or copy its entries into the overlay. The snapshot lives only across these
+// passes, which mutate signatures but never append or rename registered funcs.
+typedef struct {
+    CBMRegisteredFunc **slots;
+    size_t capacity;
+} TSSignatureIndex;
+
+static void ts_signature_work(void) {
+#ifdef CBM_ENABLE_TEST_SEAMS
+    g_ts_signature_lookup_work++;
+#endif
+}
+
+static size_t ts_signature_slot(const TSSignatureIndex *index, const char *qn) {
+    uint64_t hash = 0xcbf29ce484222325ULL;
+    for (const unsigned char *p = (const unsigned char *)qn; *p; p++) {
+        hash ^= *p;
+        hash *= 0x100000001b3ULL;
+    }
+    size_t slot = (size_t)hash & (index->capacity - 1);
+    for (;;) {
+        ts_signature_work();
+        if (!index->slots[slot] || strcmp(index->slots[slot]->qualified_name, qn) == 0)
+            return slot;
+        slot = (slot + 1) & (index->capacity - 1);
+    }
+}
+
+static TSSignatureIndex ts_signature_index(CBMTypeRegistry *reg) {
+    TSSignatureIndex index = {0};
+    if (reg->func_count <= 0 || (size_t)reg->func_count > SIZE_MAX / 2)
+        return index;
+    size_t capacity = 16;
+    while (capacity < (size_t)reg->func_count * 2) {
+        if (capacity > SIZE_MAX / 2)
+            return index;
+        capacity *= 2;
+    }
+    if (capacity > SIZE_MAX / sizeof(*index.slots))
+        return index;
+    index.slots = cbm_alloc(CBM_MEM_CLASS_EXTRACT, capacity * sizeof(*index.slots));
+    if (!index.slots)
+        return index; // Preserve the original lookup on allocation failure.
+    index.capacity = capacity;
+    memset(index.slots, 0, capacity * sizeof(*index.slots));
+    for (int i = 0; i < reg->func_count; i++) {
+        ts_signature_work();
+        CBMRegisteredFunc *func = &reg->funcs[i];
+        if (!func->qualified_name)
+            continue;
+        size_t slot = ts_signature_slot(&index, func->qualified_name);
+        if (!index.slots[slot])
+            index.slots[slot] = func;
+    }
+    return index;
+}
+
+static CBMRegisteredFunc *ts_signature_func(CBMTypeRegistry *reg, const TSSignatureIndex *index,
+                                            const char *qn) {
+    if (!qn)
+        return NULL;
+    if (index->slots)
+        return index->slots[ts_signature_slot(index, qn)];
+    for (int i = 0; i < reg->func_count; i++) {
+        ts_signature_work();
+        if (reg->funcs[i].qualified_name && strcmp(reg->funcs[i].qualified_name, qn) == 0)
+            return &reg->funcs[i];
+    }
+    return NULL;
+}
+
+#ifdef CBM_ENABLE_TEST_SEAMS
+int cbm_ts_lsp_test_signature_func_index(CBMTypeRegistry *reg, const char *qn) {
+    TSSignatureIndex index = ts_signature_index(reg);
+    CBMRegisteredFunc *func = ts_signature_func(reg, &index, qn);
+    int result = func ? (int)(func - reg->funcs) : -1;
+    cbm_free(CBM_MEM_CLASS_EXTRACT, index.slots);
+    return result;
+}
+#endif
+
 // Rebuild a function's CBMRegisteredFunc signature from the AST directly. extract_defs
 // only handles type_identifier / generic_type / predefined_type for params, so function
 // types (`(x: Foo) => void`), object types, tuple types etc. are missing from
@@ -4854,18 +5008,22 @@ static const CBMType *infer_return_type_from_body(TSLSPContext *ctx, TSNode body
 //
 // Called after `register_file_defs` so the func entries already exist; we only mutate
 // in place.
-static void rebuild_signatures_from_ast(TSLSPContext *ctx, TSNode root, CBMTypeRegistry *reg) {
+static void rebuild_signatures_from_ast(TSLSPContext *ctx, TSNode root, CBMTypeRegistry *reg,
+                                        const TSSignatureIndex *index) {
     if (ts_node_is_null(root) || !reg || !ctx->module_qn)
         return;
 
-    enum { STACK_CAP = 256 };
-    TSNode stack[STACK_CAP];
+    // 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;
-    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);
+    bool walking =
+        ts_walk_push_children(&stack, &cap, &top, root, inline_stack, "ts_lsp_rebuild_signatures");
 
-    while (top > 0) {
+    while (walking && top > 0) {
         TSNode n = stack[--top];
         if (ts_node_is_null(n))
             continue;
@@ -4873,16 +5031,15 @@ static void rebuild_signatures_from_ast(TSLSPContext *ctx, TSNode root, CBMTypeR
 
         // 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);
-            }
+            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) && top < STACK_CAP)
-                stack[top++] = body;
+            if (!ts_node_is_null(body))
+                walking = ts_walk_push(&stack, &cap, &top, body, inline_stack,
+                                       "ts_lsp_rebuild_signatures");
             continue;
         }
 
@@ -4993,43 +5150,42 @@ static void rebuild_signatures_from_ast(TSLSPContext *ctx, TSNode root, CBMTypeR
         // Build new signature and patch in registry.
         const CBMType *rets[2] = {ret, NULL};
         const CBMType *new_sig = cbm_type_func(ctx->arena, p_names, p_types, rets);
-        for (int fi = 0; fi < reg->func_count; fi++) {
-            if (!reg->funcs[fi].qualified_name)
-                continue;
-            if (strcmp(reg->funcs[fi].qualified_name, fqn) != 0)
-                continue;
+        CBMRegisteredFunc *func = ts_signature_func(reg, index, fqn);
+        if (func) {
             // Preserve the ret type if AST gave UNKNOWN but the existing reg sig had one.
-            if (cbm_type_is_unknown(ret) && reg->funcs[fi].signature &&
-                reg->funcs[fi].signature->kind == CBM_TYPE_FUNC &&
-                reg->funcs[fi].signature->data.func.return_types &&
-                reg->funcs[fi].signature->data.func.return_types[0] &&
-                !cbm_type_is_unknown(reg->funcs[fi].signature->data.func.return_types[0])) {
-                const CBMType *preserved[2] = {reg->funcs[fi].signature->data.func.return_types[0],
-                                               NULL};
+            if (cbm_type_is_unknown(ret) && func->signature &&
+                func->signature->kind == CBM_TYPE_FUNC && func->signature->data.func.return_types &&
+                func->signature->data.func.return_types[0] &&
+                !cbm_type_is_unknown(func->signature->data.func.return_types[0])) {
+                const CBMType *preserved[2] = {func->signature->data.func.return_types[0], NULL};
                 new_sig = cbm_type_func(ctx->arena, p_names, p_types, preserved);
             }
-            reg->funcs[fi].signature = new_sig;
-            break;
+            func->signature = new_sig;
         }
     }
+    cbm_walk_stack_release(stack, inline_stack);
 }
 
 // Convert NAMED references to TYPE_PARAM for any function whose declaration has a
 // `type_parameters` block. Without this, generic functions like `function id<T>(x: T): T`
 // register T as NAMED("module.T") instead of TYPE_PARAM("T"), defeating substitution.
-static void convert_signature_type_params(TSLSPContext *ctx, TSNode root, CBMTypeRegistry *reg) {
+static void convert_signature_type_params(TSLSPContext *ctx, TSNode root, CBMTypeRegistry *reg,
+                                          const TSSignatureIndex *index) {
     if (ts_node_is_null(root) || !reg || !ctx->module_qn)
         return;
 
     // Walk: function_declaration, class_declaration { method_definition }, plus exported.
-    enum { STACK_CAP = 256 };
-    TSNode stack[STACK_CAP];
+    // 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;
-    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);
+    bool walking =
+        ts_walk_push_children(&stack, &cap, &top, root, inline_stack, "ts_lsp_type_params");
 
-    while (top > 0) {
+    while (walking && top > 0) {
         TSNode n = stack[--top];
         if (ts_node_is_null(n))
             continue;
@@ -5037,16 +5193,15 @@ static void convert_signature_type_params(TSLSPContext *ctx, TSNode root, CBMTyp
 
         // 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);
-            }
+            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) && top < STACK_CAP)
-                stack[top++] = body;
+            if (!ts_node_is_null(body))
+                walking =
+                    ts_walk_push(&stack, &cap, &top, body, inline_stack, "ts_lsp_type_params");
             continue;
         }
 
@@ -5122,17 +5277,14 @@ static void convert_signature_type_params(TSLSPContext *ctx, TSNode root, CBMTyp
         }
 
         // Patch the registered func's signature: rewrite NAMED → TYPE_PARAM.
-        for (int fi = 0; fi < reg->func_count; fi++) {
-            if (!reg->funcs[fi].qualified_name)
-                continue;
-            if (strcmp(reg->funcs[fi].qualified_name, fqn) != 0)
-                continue;
-            const CBMType *old_sig = reg->funcs[fi].signature;
+        CBMRegisteredFunc *func = ts_signature_func(reg, index, fqn);
+        if (func) {
+            const CBMType *old_sig = func->signature;
             if (!old_sig)
-                break;
+                continue;
             const CBMType *new_sig = cbm_type_substitute(ctx->arena, old_sig, names, args);
             if (new_sig)
-                reg->funcs[fi].signature = new_sig;
+                func->signature = new_sig;
             // Also store the type_param_names for later overload-by-types use.
             const char **stored = (const char **)cbm_arena_alloc(
                 ctx->arena, (size_t)(tpc + 1) * sizeof(const char *));
@@ -5140,11 +5292,20 @@ static void convert_signature_type_params(TSLSPContext *ctx, TSNode root, CBMTyp
                 for (int ti = 0; ti < tpc; ti++)
                     stored[ti] = names[ti];
                 stored[tpc] = NULL;
-                reg->funcs[fi].type_param_names = stored;
+                func->type_param_names = stored;
             }
-            break;
         }
     }
+    cbm_walk_stack_release(stack, inline_stack);
+}
+
+// Build once after shape registration, then release before later passes may
+// grow the registry. This keeps per-file lookup storage out of the result arena.
+static void refine_ast_signatures(TSLSPContext *ctx, TSNode root, CBMTypeRegistry *reg) {
+    TSSignatureIndex index = ts_signature_index(reg);
+    rebuild_signatures_from_ast(ctx, root, reg, &index);
+    convert_signature_type_params(ctx, root, reg, &index);
+    cbm_free(CBM_MEM_CLASS_EXTRACT, index.slots);
 }
 
 // Parse a JSDoc comment block for `@param {T} name` and `@returns {T}` tags. For each
@@ -5789,8 +5950,7 @@ void cbm_run_ts_lsp(CBMArena *arena, CBMFileResult *result, const char *source,
     }
 
     ast_sweep_shapes(&ctx, root, &reg);
-    rebuild_signatures_from_ast(&ctx, root, &reg);
-    convert_signature_type_params(&ctx, root, &reg);
+    refine_ast_signatures(&ctx, root, &reg);
     apply_jsdoc_signatures(&ctx, root, &reg);
     infer_implicit_returns(&ctx, root, &reg);
 
@@ -6043,8 +6203,7 @@ void cbm_run_ts_lsp_cross_with_registry(CBMArena *arena, const char *source, int
     }
 
     ast_sweep_shapes(&ctx, root, &overlay);
-    rebuild_signatures_from_ast(&ctx, root, &overlay);
-    convert_signature_type_params(&ctx, root, &overlay);
+    refine_ast_signatures(&ctx, root, &overlay);
     apply_jsdoc_signatures(&ctx, root, &overlay);
     infer_implicit_returns(&ctx, root, &overlay);
 
@@ -6259,8 +6418,7 @@ void cbm_run_ts_lsp_cross(CBMArena *arena, const char *source, int source_len,
     }
 
     ast_sweep_shapes(&ctx, root, &reg);
-    rebuild_signatures_from_ast(&ctx, root, &reg);
-    convert_signature_type_params(&ctx, root, &reg);
+    refine_ast_signatures(&ctx, root, &reg);
     apply_jsdoc_signatures(&ctx, root, &reg);
     infer_implicit_returns(&ctx, root, &reg);
 

internal/cbm/lsp/ts_lsp.h

@@ -72,6 +72,11 @@ typedef struct {
 // for the calling thread, valid after a cbm_run_ts_lsp on the same thread.
 long cbm_ts_lsp_test_budget_remaining(void);
 bool cbm_ts_lsp_test_budget_warned(void);
+// Registry entries/probes visited by the two AST signature passes in the last run.
+unsigned long cbm_ts_lsp_test_signature_lookup_work(void);
+// Inspect signature-index selection and checked walker addition without huge ASTs.
+int cbm_ts_lsp_test_signature_func_index(CBMTypeRegistry *reg, const char *qn);
+int cbm_ts_lsp_test_walk_pending_count(int top, uint32_t count);
 // Maximum recursive calls in one member/method lookup during the last TS LSP run,
 // including attempts rejected by the depth or visited-type guard.
 int cbm_ts_lsp_test_max_member_lookup_steps(void);

tests/test_ts_lsp.c

@@ -26,6 +26,7 @@
 #include "../src/foundation/compat.h"
 #include "../src/pipeline/lsp_resolve.h"
 #include "lsp/ts_lsp.h"
+#include <limits.h>
 
 /* ── Helpers ───────────────────────────────────────────────────────────────── */
 
@@ -310,6 +311,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"
@@ -3433,6 +3449,122 @@ TEST(tslsp_stress_many_classes) {
     PASS();
 }
 
+TEST(tslsp_signature_walk_keeps_first_local_registration) {
+    CBMRegisteredFunc base_funcs[] = {{.qualified_name = "test.base.only"}};
+    CBMTypeRegistry base = {.funcs = base_funcs, .func_count = 1, .read_only = true};
+    CBMRegisteredFunc funcs[] = {
+        {.qualified_name = NULL},
+        {.qualified_name = "test.main.identity"},
+        {.qualified_name = "test.main.identity"},
+        {.qualified_name = "test.main.Box.identity"},
+    };
+    CBMTypeRegistry overlay = {.funcs = funcs, .func_count = 4, .fallback = &base};
+    ASSERT_EQ(cbm_ts_lsp_test_signature_func_index(&overlay, "test.main.identity"), 1);
+    ASSERT_EQ(cbm_ts_lsp_test_signature_func_index(&overlay, "test.main.Box.identity"), 3);
+    ASSERT_EQ(cbm_ts_lsp_test_signature_func_index(&overlay, "test.base.only"), -1);
+    ASSERT_EQ(cbm_ts_lsp_test_signature_func_index(&overlay, "test.missing"), -1);
+    ASSERT_EQ(overlay.func_count, 4);
+    ASSERT_EQ(base.func_count, 1);
+    PASS();
+}
+
+TEST(tslsp_walk_pending_count_rejects_overflow) {
+    ASSERT_EQ(cbm_ts_lsp_test_walk_pending_count(255, 300), 555);
+    ASSERT_EQ(cbm_ts_lsp_test_walk_pending_count(INT_MAX - 1, 1), INT_MAX);
+    ASSERT_EQ(cbm_ts_lsp_test_walk_pending_count(INT_MAX, 1), -1);
+    ASSERT_EQ(cbm_ts_lsp_test_walk_pending_count(1, (uint32_t)INT_MAX), -1);
+    ASSERT_EQ(cbm_ts_lsp_test_walk_pending_count(0, UINT32_MAX), -1);
+    PASS();
+}
+
+/* Count actual registry work, independent of machine speed. Both inputs exceed
+ * the old walk cap, and every declaration exercises rebuild and generic rewrite. */
+TEST(tslsp_signature_walk_lookup_scales_linearly) {
+    unsigned long work[2] = {0};
+    for (int scale = 0; scale < 2; scale++) {
+        const int count = 300 << scale;
+        const size_t capacity = (size_t)count * 160 + 512;
+        char *source = (char *)malloc(capacity);
+        ASSERT_NOT_NULL(source);
+        size_t used = 0;
+        for (int i = 0; i < count; i++) {
+            used += (size_t)snprintf(source + used, capacity - used,
+                                     "export function identity%d<T>(x: T): T { return x; }\n", i);
+        }
+        snprintf(source + used, capacity - used,
+                 "class Box { use(): void {} }\n"
+                 "function go() { const b = identity%d(new Box()); b.use(); }\n",
+                 count - 1);
+        CBMFileResult *r = extract_ts(source);
+        free(source);
+        ASSERT_NOT_NULL(r);
+        work[scale] = cbm_ts_lsp_test_signature_lookup_work();
+        ASSERT_GTE(require_resolved(r, ".go", "use"), 0);
+        cbm_free_result(r);
+    }
+    printf("  signature lookup work: 300=%lu 600=%lu\n", work[0], work[1]);
+    ASSERT_GT(work[0], 0);
+    /* Allow 25% above ideal doubling for hash-bucket distribution. */
+    ASSERT_LT(work[1], work[0] * 5 / 2);
+    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);
+    ASSERT_NOT_NULL(buf);
+    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);
+    ASSERT_NOT_NULL(buf);
+    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];
@@ -4860,6 +4992,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);
 
@@ -5189,6 +5322,11 @@ SUITE(ts_lsp) {
 
     /* Stress tests */
     RUN_TEST(tslsp_stress_many_classes);
+    RUN_TEST(tslsp_signature_walk_keeps_first_local_registration);
+    RUN_TEST(tslsp_walk_pending_count_rejects_overflow);
+    RUN_TEST(tslsp_signature_walk_lookup_scales_linearly);
+    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);