Fix(memorization) issue 710 py lsp memoization

internal/cbm/lsp/py_lsp.c

@@ -26,6 +26,7 @@
 // Forward decls
 static void py_resolve_calls_in(PyLSPContext *ctx, TSNode node);
 static const CBMType *py_eval_expr_type(PyLSPContext *ctx, TSNode node);
+static const CBMType *py_eval_expr_type_uncached(PyLSPContext *ctx, TSNode node);
 static void py_process_statement(PyLSPContext *ctx, TSNode node);
 static const CBMRegisteredFunc *py_lookup_attribute(PyLSPContext *ctx, const char *type_qn,
                                                     const char *member_name);
@@ -692,7 +693,8 @@ static const CBMType *py_iterable_element_type(PyLSPContext *ctx, const CBMType
     return cbm_type_unknown();
 }
 
-static const CBMType *py_eval_expr_type(PyLSPContext *ctx, TSNode node) {
+// previously named py_eval_expr_type
+static const CBMType *py_eval_expr_type_uncached(PyLSPContext *ctx, TSNode node) {
     if (!ctx || ts_node_is_null(node))
         return cbm_type_unknown();
 
@@ -1286,7 +1288,114 @@ static const CBMType *py_eval_expr_type(PyLSPContext *ctx, TSNode node) {
     return cbm_type_unknown();
 }
 
-/* ── statement processing: bind from assignments, for-loops, with-as ──── */
+/* ── py_eval_expr_type memoization (issue #710) ──────────────────────────
+ *
+ * py_eval_expr_type_uncached is a recursive-descent type evaluator. For a
+ * deeply chained call expression — `Builder().add(x).add(y)....add(z)` —
+ * resolving the outer call's type requires resolving the inner call's
+ * type, recursively, down the whole chain. On its own that is O(depth),
+ * which is fine even for 60+ levels.
+ *
+ * Cache lives on PyLSPContext (one per file), arena-allocated, so it's
+ * freed automatically when the file's arena is destroyed — no manual
+ * cleanup needed.
+ */
+
+enum { PY_TYPE_CACHE_INITIAL_CAP = 256 };
+
+/* Linear-probe lookup. Returns the cached type, or NULL if this node
+ * hasn't been evaluated yet (a NULL result is never itself cached — see
+ * py_type_cache_insert — so NULL unambiguously means "not yet computed"). */
+static const CBMType *py_type_cache_lookup(PyLSPContext *ctx, uint32_t key) {
+    if (!ctx->type_cache || ctx->type_cache_cap == 0)
+        return NULL;
+    int cap = ctx->type_cache_cap;
+    int idx = (int)(key % (uint32_t)cap);
+    for (int probed = 0; probed < cap; probed++) {
+        CBMPyTypeCacheEntry *e = &ctx->type_cache[idx];
+        if (!e->result)
+            return NULL; /* empty slot: key was never inserted */
+        if (e->node_start_byte == key)
+            return e->result;
+        idx = (idx + 1) % cap;
+    }
+    return NULL; /* table full and key not found (shouldn't happen: we grow before full) */
+}
+
+/* Grow the cache (arena allocation: old table is simply abandoned, freed
+ * when the per-file arena is destroyed, same pattern as
+ * ts_nstack_push in extract_node_stack.h). Rehashes all live entries. */
+static void py_type_cache_grow(PyLSPContext *ctx) {
+    int old_cap = ctx->type_cache_cap;
+    CBMPyTypeCacheEntry *old_entries = ctx->type_cache;
+
+    int new_cap = old_cap ? old_cap * 2 : PY_TYPE_CACHE_INITIAL_CAP;
+    CBMPyTypeCacheEntry *new_entries = (CBMPyTypeCacheEntry *)cbm_arena_alloc(
+        ctx->arena, (size_t)new_cap * sizeof(CBMPyTypeCacheEntry));
+    if (!new_entries)
+        return;
+    memset(new_entries, 0, (size_t)new_cap * sizeof(CBMPyTypeCacheEntry));
+
+    ctx->type_cache = new_entries;
+    ctx->type_cache_cap = new_cap;
+    ctx->type_cache_count = 0;
+
+    for (int i = 0; i < old_cap; i++) {
+        if (old_entries[i].result) {
+            /* Re-insert via the now-larger table. Safe to call insert here:
+             * it only grows again if count*4 >= cap*3, which can't happen
+             * mid-rehash since we just doubled capacity. */
+            int idx = (int)(old_entries[i].node_start_byte % (uint32_t)new_cap);
+            while (new_entries[idx].result) {
+                idx = (idx + 1) % new_cap;
+            }
+            new_entries[idx] = old_entries[i];
+            ctx->type_cache_count++;
+        }
+    }
+}
+
+/* Insert a result into the cache. A NULL result is never cached (so a
+ * lookup miss and an explicit "evaluated to NULL" are indistinguishable —
+ * acceptable since py_eval_expr_type_uncached treats NULL the same as
+ * cbm_type_unknown() everywhere it's checked, so re-evaluating a rare NULL
+ * case costs nothing asymptotically). */
+static void py_type_cache_insert(PyLSPContext *ctx, uint32_t key, const CBMType *result) {
+    if (!result)
+        return;
+    if (!ctx->type_cache || ctx->type_cache_count * 4 >= ctx->type_cache_cap * 3) {
+        py_type_cache_grow(ctx);
+        if (!ctx->type_cache)
+            return; /* OOM during first grow: skip caching, fall back to recompute */
+    }
+    int cap = ctx->type_cache_cap;
+    int idx = (int)(key % (uint32_t)cap);
+    while (ctx->type_cache[idx].result) {
+        if (ctx->type_cache[idx].node_start_byte == key)
+            return; /* already cached (re-entrant evaluation of the same node) */
+        idx = (idx + 1) % cap;
+    }
+    ctx->type_cache[idx].node_start_byte = key;
+    ctx->type_cache[idx].result = result;
+    ctx->type_cache_count++;
+}
+
+/* Memoizing wrapper — the function every call site in this file actually
+ * calls. py_eval_expr_type_uncached holds the real evaluation logic and is
+ * called at most once per distinct node. */
+static const CBMType *py_eval_expr_type(PyLSPContext *ctx, TSNode node) {
+    if (!ctx || ts_node_is_null(node))
+        return cbm_type_unknown();
+
+    uint32_t key = ts_node_start_byte(node);
+    const CBMType *cached = py_type_cache_lookup(ctx, key);
+    if (cached)
+        return cached;
+
+    const CBMType *result = py_eval_expr_type_uncached(ctx, node);
+    py_type_cache_insert(ctx, key, result);
+    return result;
+}
 
 static void py_process_statement(PyLSPContext *ctx, TSNode node) {
     if (!ctx || ts_node_is_null(node))
@@ -1766,8 +1875,8 @@ static void py_emit_call_for(PyLSPContext *ctx, TSNode call_node) {
                     // Skip if mod is already rooted under the project to avoid
                     // "<root>.<root>.mod".
                     if (!(strncmp(mod, ctx->module_qn, root_len) == 0 && mod[root_len] == '.')) {
-                        char *qual_mod = (char *)cbm_arena_alloc(ctx->arena, root_len + 1 +
-                                                                                strlen(mod) + 1);
+                        char *qual_mod =
+                            (char *)cbm_arena_alloc(ctx->arena, root_len + 1 + strlen(mod) + 1);
                         if (qual_mod) {
                             memcpy(qual_mod, ctx->module_qn, root_len);
                             qual_mod[root_len] = '.';

internal/cbm/lsp/py_lsp.h

@@ -24,6 +24,15 @@ typedef struct {
     const char *target_qn;   // e.g. "test.main.foo"
 } CBMDictLiteralEntry;
 
+// Memoization entry for py_eval_expr_type. Keyed by the expression node's
+// start byte offset, which is a unique identity for a node
+// within a single file's parse tree (two distinct expressions never start
+// at the same byte)
+typedef struct {
+    uint32_t node_start_byte;
+    const CBMType *result;
+} CBMPyTypeCacheEntry;
+
 // PyLSPContext holds state for one Python file's type evaluation.
 typedef struct {
     CBMArena *arena;
@@ -32,6 +41,13 @@ typedef struct {
     const CBMTypeRegistry *registry;
     CBMScope *current_scope;
 
+    // Memoization cache for py_eval_expr_type (issue #710: deeply chained
+    // method calls caused exponential reevaluation of shared subchains without this, Open
+    // addressing, linear probe, grown geometrically
+    CBMPyTypeCacheEntry *type_cache;
+    int type_cache_count;
+    int type_cache_cap;
+
     // Import map: local_name -> module_qn (arena-allocated, NULL-terminated).
     // Built from CBMFileResult.imports.
     const char **import_local_names;

tests/test_py_lsp.c

@@ -1243,6 +1243,39 @@ TEST(pylsp_round6_property_access_chains) {
     PASS();
 }
 
+TEST(pylsp_issue710_deep_call_chain_no_hang) {
+    /* Regression for #710: a deeply chained builder expression made
+     * py_eval_expr_type re-evaluate shared sub-chains at every level,
+     * O(2^depth). On the pre-fix code this hangs the definitions pass
+     * for 800+ seconds on a single file; with per-node memoization it
+     * completes in milliseconds. This test builds a ~40-link chain and
+     * asserts extraction terminates and stays addressable — if the
+     * exponential blowup regresses, the suite hangs here (caught by the
+     * runner's timeout) rather than passing silently. */
+    CBMFileResult *r = extract_py(
+        "from typing import Self\n"
+        "class G:\n"
+        "    def add(self, x) -> Self:\n"
+        "        return self\n"
+        "    def compile(self):\n"
+        "        return 1\n"
+        "def build():\n"
+        "    return (\n"
+        "        G()\n"
+        "        .add(1).add(2).add(3).add(4).add(5).add(6).add(7).add(8)\n"
+        "        .add(9).add(10).add(11).add(12).add(13).add(14).add(15)\n"
+        "        .add(16).add(17).add(18).add(19).add(20).add(21).add(22)\n"
+        "        .add(23).add(24).add(25).add(26).add(27).add(28).add(29)\n"
+        "        .add(30).add(31).add(32).add(33).add(34).add(35).add(36)\n"
+        "        .add(37).add(38).add(39).add(40)\n"
+        "        .compile()\n"
+        "    )\n");
+    ASSERT_NOT_NULL(r);
+    ASSERT_GTE(require_resolved(r, "build", "compile"), 0);
+    cbm_free_result(r);
+    PASS();
+}
+
 /* ── Suite ─────────────────────────────────────────────────────── */
 
 SUITE(py_lsp) {
@@ -1293,6 +1326,7 @@ SUITE(py_lsp) {
     RUN_TEST(pylsp_round1_assert_type_passthrough);
     RUN_TEST(pylsp_round1_forward_reference);
     RUN_TEST(pylsp_round1_self_return_chains);
+    RUN_TEST(pylsp_issue710_deep_call_chain_no_hang);
     RUN_TEST(pylsp_round1_generic_subscript_annotation);
     /* Round 2 — narrowing */
     RUN_TEST(pylsp_round2_isinstance_narrow);