From e8ccc330576ce3fdf826d2dbb0a0a515f0d11473 Mon Sep 17 00:00:00 2001 From: Martin Vogel Date: Tue, 20 Oct 2026 19:27:42 +0200 Subject: [PATCH] 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 Signed-off-by: Martin Vogel --- internal/cbm/extract_node_stack.h | 18 ++--- internal/cbm/lsp/ts_lsp.c | 130 +++++++++++++++++++++--------- tests/test_ts_lsp.c | 22 ++++- 3 files changed, 115 insertions(+), 55 deletions(-) diff --git a/internal/cbm/extract_node_stack.h b/internal/cbm/extract_node_stack.h index 3c40838cd..30a27968b 100644 --- a/internal/cbm/extract_node_stack.h +++ b/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. */ diff --git a/internal/cbm/lsp/ts_lsp.c b/internal/cbm/lsp/ts_lsp.c index 228b32736..cf827e7c8 100644 --- a/internal/cbm/lsp/ts_lsp.c +++ b/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 #include #include @@ -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 diff --git a/tests/test_ts_lsp.c b/tests/test_ts_lsp.c index a6261e967..bcbb89d21 100644 --- a/tests/test_ts_lsp.c +++ b/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 { 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);