Merge pull request #2420 from Fieldnote-Echo/perf/py-cross-registry-index

perf(py-lsp): O(1) type membership in the shared Python cross-registry build

internal/cbm/lsp/py_lsp.c

@@ -6179,10 +6179,36 @@ static void py_split_field_defs(CBMArena *arena, const char *field_defs, const c
     *types_out = types;
 }
 
+/* Qualified names of reg->types[0..noted), for a caller that registers into a
+ * registry it does not finalize until the end (cbm_py_build_cross_registry):
+ * that registry's own lookup scans every type. qns == NULL probes the
+ * registry instead. */
+typedef struct {
+    CBMHashTable *qns;
+    int noted;
+} PyKnownTypes;
+
+/* Note the types added since the last call. Run after every add_type, so an
+ * add that failed silently cannot desync set and registry. If the set cannot
+ * grow it is dropped and probes fall back to the registry. */
+static void py_note_known_types(PyKnownTypes *known, const CBMTypeRegistry *reg) {
+    for (; known && known->qns && known->noted < reg->type_count; known->noted++) {
+        const char *qn = reg->types[known->noted].qualified_name;
+        if (!qn || cbm_ht_has(known->qns, qn))
+            continue;
+        cbm_ht_set(known->qns, qn, (void *)(uintptr_t)(known->noted + 1));
+        if (!cbm_ht_has(known->qns, qn)) {
+            cbm_ht_free(known->qns);
+            known->qns = NULL;
+        }
+    }
+}
+
 /* Build a registry from CBMLSPDef[] supplied by the caller — covers both
- * the source file's own defs and cross-file referenced defs. */
+ * the source file's own defs and cross-file referenced defs. known is NULL
+ * except for the shared Tier-2 build. */
 static bool py_register_lsp_defs(CBMArena *arena, CBMArena *idx_arena, CBMTypeRegistry *reg,
-                                 CBMLSPDef *defs, int def_count) {
+                                 CBMLSPDef *defs, int def_count, PyKnownTypes *known) {
     /* Pass 1: types only — the method pass probes the registry per Method def
      * (receiver auto-registration), which is a LINEAR scan pre-finalize:
      * O(methods x types) per file (same quadratic as php_register_lsp_defs;
@@ -6206,6 +6232,7 @@ static bool py_register_lsp_defs(CBMArena *arena, CBMArena *idx_arena, CBMTypeRe
             }
             py_split_field_defs(arena, d->field_defs, &rt.field_names, &rt.field_types);
             cbm_registry_add_type(reg, rt);
+            py_note_known_types(known, reg);
         }
     }
     /* idx_arena == NULL skips the mid-build finalize: the tier-2 builder
@@ -6259,13 +6286,17 @@ static bool py_register_lsp_defs(CBMArena *arena, CBMArena *idx_arena, CBMTypeRe
 
             if (strcmp(d->label, "Method") == 0 && d->receiver_type && d->receiver_type[0]) {
                 rf.receiver_type = d->receiver_type; /* borrowed */
-                if (!cbm_registry_lookup_type(reg, rf.receiver_type)) {
+                bool recv_known = known && known->qns
+                                      ? cbm_ht_has(known->qns, rf.receiver_type)
+                                      : cbm_registry_lookup_type(reg, rf.receiver_type) != NULL;
+                if (!recv_known) {
                     CBMRegisteredType auto_t;
                     memset(&auto_t, 0, sizeof(auto_t));
                     auto_t.qualified_name = rf.receiver_type;
                     const char *dot = strrchr(d->receiver_type, '.');
                     auto_t.short_name = dot ? dot + 1 : rf.receiver_type; /* borrowed substring */
                     cbm_registry_add_type(reg, auto_t);
+                    py_note_known_types(known, reg);
                 }
             }
             cbm_registry_add_func(reg, rf);
@@ -6277,7 +6308,7 @@ static bool py_register_lsp_defs(CBMArena *arena, CBMArena *idx_arena, CBMTypeRe
 #ifdef CBM_ENABLE_TEST_SEAMS
 bool cbm_py_lsp_test_register_defs(CBMArena *arena, CBMTypeRegistry *reg, CBMLSPDef *defs,
                                    int def_count) {
-    return py_register_lsp_defs(arena, NULL, reg, defs, def_count);
+    return py_register_lsp_defs(arena, NULL, reg, defs, def_count, NULL);
 }
 #endif
 
@@ -6774,7 +6805,7 @@ CBMLSPStatus cbm_run_py_lsp_cross_status(CBMArena *arena, const char *source, in
     py_lsp_init_for_stage(&ctx, arena, source, source_len, &reg, module_qn, out, "cross");
     if (ctx.eval_failure)
         goto cross_cleanup;
-    if (!py_register_lsp_defs(arena, &idx_arena, &reg, defs, def_count)) {
+    if (!py_register_lsp_defs(arena, &idx_arena, &reg, defs, def_count, NULL)) {
         ctx.eval_failure = CBM_LSP_SCOPE_FAILED;
         goto cross_cleanup;
     }
@@ -6813,15 +6844,27 @@ CBMTypeRegistry *cbm_py_build_cross_registry(CBMArena *arena, CBMLSPDef *defs, i
     cbm_registry_init(reg, arena);
     cbm_python_stdlib_register(reg, arena);
 
+    /* The registry is not finalized until the end, so probing it for each
+     * Method's receiver scanned every type: O(methods x types), single-threaded
+     * (~54 s of lsp_cross_prepare on a large Python tree). Probe a set of
+     * the registered type QNs instead; registration order is unchanged. Sized
+     * for the stdlib seed, not def_count (the mixed all-language count). NULL
+     * on allocation failure: the registry probe runs as before. */
+    PyKnownTypes known = {cbm_ht_create((uint32_t)reg->type_count), 0};
+    py_note_known_types(&known, reg);
+
     /* Filter to Python defs only — defs[] is mixed-language all_defs. */
     for (int i = 0; i < def_count; i++) {
         CBMLSPDef *d = &defs[i];
         if (d->lang != CBM_LANG_PYTHON)
             continue;
         /* Reuse the existing register fn on a single-def slice (n=1 inline). */
-        if (!py_register_lsp_defs(arena, NULL, reg, d, 1))
+        if (!py_register_lsp_defs(arena, NULL, reg, d, 1, &known)) {
+            cbm_ht_free(known.qns);
             return NULL; /* caller owns the shared arena and other registries */
+        }
     }
+    cbm_ht_free(known.qns);
 
     py_mark_ambiguous_callable_bindings(reg, arena);
     cbm_registry_build_type_qn_order(reg, arena);

tests/test_py_lsp.c

@@ -751,6 +751,91 @@ TEST(pylsp_fused_self_attr_chain_via_overlay) {
     PASS();
 }
 
+/* Characterization of the shared Tier-2 build order. Defs register one at a
+ * time in defs[] order, so a Method seen before its Class auto-registers a
+ * placeholder receiver type that stays ahead of the Class entry, and a Method
+ * whose receiver is already known adds no type. Any change to how the build
+ * probes for known receivers must keep this exact sequence. */
+TEST(pylsp_cross_registry_interleaved_order) {
+    CBMArena arena;
+    cbm_arena_init(&arena);
+    CBMTypeRegistry *base = cbm_py_build_cross_registry(&arena, NULL, 0);
+    ASSERT_NOT_NULL(base);
+    int tb = base->type_count, fb = base->func_count;
+
+    CBMLSPDef defs[] = {
+        {.qualified_name = "t.m.A.m",
+         .short_name = "m",
+         .label = "Method",
+         .receiver_type = "t.m.A",
+         .def_module_qn = "t.m",
+         .lang = CBM_LANG_PYTHON},
+        {.qualified_name = "t.m.A",
+         .short_name = "A",
+         .label = "Class",
+         .def_module_qn = "t.m",
+         .lang = CBM_LANG_PYTHON},
+        {.qualified_name = "t.m.B",
+         .short_name = "B",
+         .label = "Class",
+         .def_module_qn = "t.m",
+         .lang = CBM_LANG_PYTHON},
+        {.qualified_name = "t.m.B.f",
+         .short_name = "f",
+         .label = "Method",
+         .receiver_type = "t.m.B",
+         .def_module_qn = "t.m",
+         .lang = CBM_LANG_PYTHON},
+        {.qualified_name = "t.m.C.g",
+         .short_name = "g",
+         .label = "Method",
+         .receiver_type = "t.m.C",
+         .def_module_qn = "t.m",
+         .lang = CBM_LANG_PYTHON},
+        /* Receiver auto-registered by C.g just above: no second placeholder. */
+        {.qualified_name = "t.m.C.h",
+         .short_name = "h",
+         .label = "Method",
+         .receiver_type = "t.m.C",
+         .def_module_qn = "t.m",
+         .lang = CBM_LANG_PYTHON},
+        /* Non-Python defs are skipped by the builder. */
+        {.qualified_name = "t.m.D",
+         .short_name = "D",
+         .label = "Class",
+         .def_module_qn = "t.m",
+         .lang = CBM_LANG_GO},
+    };
+    CBMTypeRegistry *reg =
+        cbm_py_build_cross_registry(&arena, defs, (int)(sizeof(defs) / sizeof(defs[0])));
+    ASSERT_NOT_NULL(reg);
+
+    /* Types: placeholder A (from A.m), Class A, Class B, placeholder C. */
+    ASSERT_EQ(reg->type_count, tb + 4);
+    ASSERT_TRUE(reg->types[tb].qualified_name == defs[0].receiver_type);
+    ASSERT_STR_EQ(reg->types[tb].short_name, "A");
+    ASSERT_TRUE(reg->types[tb + 1].qualified_name == defs[1].qualified_name);
+    ASSERT_TRUE(reg->types[tb + 2].qualified_name == defs[2].qualified_name);
+    ASSERT_TRUE(reg->types[tb + 3].qualified_name == defs[4].receiver_type);
+    ASSERT_STR_EQ(reg->types[tb + 3].short_name, "C");
+
+    /* Funcs keep defs[] order. */
+    ASSERT_EQ(reg->func_count, fb + 4);
+    ASSERT_STR_EQ(reg->funcs[fb].qualified_name, "t.m.A.m");
+    ASSERT_STR_EQ(reg->funcs[fb + 1].qualified_name, "t.m.B.f");
+    ASSERT_STR_EQ(reg->funcs[fb + 2].qualified_name, "t.m.C.g");
+    ASSERT_STR_EQ(reg->funcs[fb + 2].receiver_type, "t.m.C");
+    ASSERT_STR_EQ(reg->funcs[fb + 3].qualified_name, "t.m.C.h");
+
+    const CBMRegisteredType *a = cbm_registry_lookup_type(reg, "t.m.A");
+    ASSERT_NOT_NULL(a);
+    ASSERT_STR_EQ(a->qualified_name, "t.m.A");
+    ASSERT_NULL(cbm_registry_lookup_type(reg, "t.m.D"));
+
+    cbm_arena_destroy(&arena);
+    PASS();
+}
+
 /* Issue #228: a class/static method invoked directly on a CROSS-FILE imported
  * class name — ActionRecordX.build_from_text(...) — produced no CALLS edge, so
  * the method showed in/out degree 0 and was flagged as dead code. Distinct from
@@ -2448,6 +2533,59 @@ TEST(pylsp_scale_work_is_linear) {
     PASS();
 }
 
+/* #2413 (O9: no clock): the shared Tier-2 build registers one def at a time
+ * into a registry it finalizes only at the end, so a receiver probe that scans
+ * the types registered so far makes it O(methods x types). The registry
+ * primitives count every candidate they examine (CBM_LSP_WORK): 4x the classes
+ * must cost at most ~4x the work; the scanning probe lands at ~16x. */
+static uint64_t cross_registry_work_for(int n_classes, int *out_types) {
+    CBMArena arena;
+    cbm_arena_init(&arena);
+    int n_defs = n_classes * 3;
+    CBMLSPDef *defs = (CBMLSPDef *)cbm_arena_calloc(&arena, (size_t)n_defs * sizeof(*defs));
+    if (!defs) {
+        cbm_arena_destroy(&arena);
+        *out_types = -1;
+        return 0;
+    }
+    /* Extraction order: each class followed by its two methods. */
+    for (int i = 0; i < n_classes; i++) {
+        CBMLSPDef *cls = &defs[i * 3];
+        cls->qualified_name = cbm_arena_sprintf(&arena, "scale.mod.Cls%d", i);
+        cls->short_name = cls->qualified_name + strlen("scale.mod.");
+        cls->label = "Class";
+        cls->def_module_qn = "scale.mod";
+        cls->lang = CBM_LANG_PYTHON;
+        for (int j = 1; j <= 2; j++) {
+            CBMLSPDef *m = &defs[i * 3 + j];
+            m->qualified_name = cbm_arena_sprintf(&arena, "%s.m%d", cls->qualified_name, j);
+            m->short_name = j == 1 ? "m1" : "m2";
+            m->label = "Method";
+            m->receiver_type = cls->qualified_name;
+            m->def_module_qn = "scale.mod";
+            m->lang = CBM_LANG_PYTHON;
+        }
+    }
+    (void)cbm_lsp_work_take();
+    CBMTypeRegistry *reg = cbm_py_build_cross_registry(&arena, defs, n_defs);
+    uint64_t work = cbm_lsp_work_take();
+    *out_types = reg ? reg->type_count : -1;
+    cbm_arena_destroy(&arena);
+    return work;
+}
+
+TEST(pylsp_cross_registry_build_work_is_linear) {
+    int t1 = 0, t4 = 0;
+    uint64_t w1 = cross_registry_work_for(1000, &t1);
+    uint64_t w4 = cross_registry_work_for(4000, &t4);
+    printf("    cross registry work: 1000=%llu (types=%d)  4000=%llu (types=%d)\n",
+           (unsigned long long)w1, t1, (unsigned long long)w4, t4);
+    /* Every class registered exactly once (no duplicate placeholder). */
+    ASSERT_EQ(t4 - t1, 3000);
+    ASSERT(w4 <= 5 * w1 + 1000);
+    PASS();
+}
+
 extern const TSLanguage *tree_sitter_python(void);
 
 TEST(pylsp_issue1527_memo_allocation_failure_is_reported) {
@@ -4850,6 +4988,7 @@ SUITE(py_lsp) {
     RUN_TEST(pylsp_issue1527_long_binary_chain_is_not_nesting);
     RUN_TEST(pylsp_issue1527_long_receiver_chain_is_not_nesting);
     RUN_TEST(pylsp_scale_work_is_linear);
+    RUN_TEST(pylsp_cross_registry_build_work_is_linear);
     /* Phase 2 — smoke */
     RUN_TEST(pylsp_smoke_empty);
     RUN_TEST(pylsp_smoke_one_function);
@@ -4886,6 +5025,7 @@ SUITE(py_lsp) {
     RUN_TEST(pylsp_crossfile_receiver_through_typed_field);
     RUN_TEST(pylsp_crossfile_typed_field_shared_registry);
     RUN_TEST(pylsp_fused_self_attr_chain_via_overlay);
+    RUN_TEST(pylsp_cross_registry_interleaved_order);
     RUN_TEST(pylsp_crossfile_classmethod_on_class_issue228);
     RUN_TEST(pylsp_crossfile_inheritance);
     RUN_TEST(pylsp_batch_two_files);