perf(resolver): look up TypeScript proof imports and Rust uses by name

The TypeScript proof scanned a file's imports for each callee and annotation, and every narrowing of a path for each receiver on it. TypeScript's own checker.ts names over a thousand imports in 54,000 lines. Imports now use the file's rows, and a narrowing lookup is one binary search: entries sorted by path and start carry the furthest end so far, and one covers a position exactly when that end passes it.

A Rust module's uses were walked for each local item found, each imported name and each trait a method call asked about. Long lists now answer from rows by local name and by last segment, in list order.

Each of these went from about 4x to about 2x work when the input doubled. Answers are unchanged.

Signed-off-by: Martin Vogel <martin.vogel.tech@gmail.com>

internal/cbm/lsp/ts_proof.c

@@ -28,7 +28,8 @@ typedef struct {
 typedef struct {
     const char *path;
     uint32_t start, end;
-    TSNode guard; /* The callee; null for instanceof. */
+    TSNode guard;   /* The callee; null for instanceof. */
+    uint32_t reach; /* Once sorted: the furthest end among this path's entries up to this one. */
 } ts_proof_narrowing_t;
 
 struct CBMTSProofAudit {
@@ -1010,6 +1011,12 @@ static void settle_guards(TSLSPContext *ctx) {
     audit->narrowing_count = kept;
     if (kept > 1)
         qsort(audit->narrowings, kept, sizeof(*audit->narrowings), narrowing_compare);
+    for (size_t i = 0; i < kept; i++) {
+        ts_proof_narrowing_t *entry = &audit->narrowings[i];
+        const ts_proof_narrowing_t *prior = i ? &audit->narrowings[i - 1] : NULL;
+        bool same_path = prior && !strcmp(prior->path, entry->path);
+        entry->reach = same_path && prior->reach > entry->end ? prior->reach : entry->end;
+    }
 }
 
 /* The checker selects on what the value was narrowed to, and this proof reads declarations, so
@@ -1026,18 +1033,21 @@ static bool receiver_narrowed(TSLSPContext *ctx, TSNode receiver, uint32_t at) {
     const char *path = narrowing_path(ctx, receiver, 0);
     if (!path)
         return false;
+    /* This path's entries starting at or before `at` are a prefix of its sorted entries, so one of
+     * them covers `at` exactly when the furthest end among them passes it. */
+    CBM_LSP_WORK(1);
     size_t low = 0, high = audit->narrowing_count;
     while (low < high) {
         size_t mid = low + (high - low) / 2;
-        if (strcmp(audit->narrowings[mid].path, path) < 0)
+        int order = strcmp(audit->narrowings[mid].path, path);
+        if (order < 0 || (order == 0 && audit->narrowings[mid].start <= at))
             low = mid + 1;
         else
             high = mid;
     }
-    for (; low < audit->narrowing_count && !strcmp(audit->narrowings[low].path, path); low++)
-        if (at >= audit->narrowings[low].start && at < audit->narrowings[low].end)
-            return true;
-    return false;
+    if (!low || strcmp(audit->narrowings[low - 1].path, path))
+        return false;
+    return at < audit->narrowings[low - 1].reach;
 }
 
 /* The registered function a call's callee names: a module function, an import, or a method of a
@@ -1067,11 +1077,8 @@ static const CBMRegisteredFunc *callee_function(TSLSPContext *ctx, TSNode callee
                 return NULL;
             qn = cbm_arena_sprintf(ctx->arena, "%s.%s", ctx->module_qn, name);
         } else {
-            for (int i = 0; i < ctx->import_count; i++) {
-                const char *key = ctx->import_local_names[i];
+            for (int i = ts_import_first(ctx, name); i >= 0; i = ts_import_next(ctx, name, i)) {
                 const char *value = ctx->import_module_qns[i];
-                if (!key || strcmp(key, name))
-                    continue;
                 if (!value || !*value || (qn && strcmp(qn, value)))
                     return NULL;
                 qn = value;
@@ -1123,10 +1130,7 @@ static bool annotation_names_local_type(TSLSPContext *ctx, TSNode annotation, ui
     (void)find_binding(ctx, name, at, &occupied);
     if (occupied)
         return true;
-    for (int i = 0; i < ctx->import_count; i++)
-        if (ctx->import_local_names[i] && !strcmp(ctx->import_local_names[i], name))
-            return true;
-    return false;
+    return ts_import_first(ctx, name) >= 0;
 }
 
 /* The class the elements of this value are: `items[i]` and `for (const x of items)`. */

internal/cbm/resolve/rust_flow.c

@@ -722,9 +722,9 @@ static bool rs_inst_binds_trait(rs_walk_t *w, rs_inst_t *inst, const rs_item_t *
     const rs_mod_t *syn = inst->syn;
     if (trait->owner == syn)
         return true;
-    for (const rs_use_t *use = syn->uses; use; use = use->next) {
-        if (!use->nseg || strcmp(use->segs[use->nseg - 1], trait->name) != 0)
-            continue;
+    int row = -1;
+    for (const rs_use_t *use = rs_use_ending(syn, NULL, trait->name, &row); use;
+         use = rs_use_ending(syn, use, trait->name, &row)) {
         int consumed = 0;
         rs_ent_t ent = rs_resolve_path(&w->q, inst, NULL, 0, use->segs, use->nseg, use->absolute,
                                        RS_NS_TYPE, true, NULL, &consumed);

internal/cbm/resolve/rust_names.c

@@ -981,8 +981,71 @@ static void rs_mod_index_impls(CBMArena *arena, rs_mod_t *mod) {
     mod->nimpl_rows = count;
 }
 
+static void rs_mod_index_uses(CBMArena *arena, rs_mod_t *mod) {
+    int count = 0;
+    for (const rs_use_t *use = mod->uses; use; use = use->next)
+        count++;
+    if (count < CBM_NAME_ROWS_MIN)
+        return;
+    const rs_use_t **by_name = cbm_arena_alloc(arena, (size_t)count * sizeof(*by_name));
+    const rs_use_t **by_tail = cbm_arena_alloc(arena, (size_t)count * sizeof(*by_tail));
+    if (!by_name || !by_tail)
+        return; /* the list still answers */
+    int names = 0, tails = 0;
+    for (const rs_use_t *use = mod->uses; use; use = use->next) {
+        if (use->local) {
+            by_name[names] = use;
+            if (!cbm_name_rows_add(arena, &mod->use_name_rows, NULL, use->local, names++))
+                return;
+        }
+        if (use->nseg) {
+            by_tail[tails] = use;
+            if (!cbm_name_rows_add(arena, &mod->use_tail_rows, NULL, use->segs[use->nseg - 1],
+                                   tails++))
+                return;
+        }
+    }
+    mod->uses_by_name = by_name;
+    mod->uses_by_tail = by_tail;
+}
+
+/* The first or next use binding this local name, in list order. */
+const rs_use_t *rs_use_named(const rs_mod_t *mod, const rs_use_t *after, const char *name,
+                             int *row) {
+    cbm_resolve_stats_add(CBM_RESOLVE_STAT_LOOKUP_PROBES, 1);
+    if (mod->uses_by_name) {
+        *row = after ? mod->use_name_rows.next[*row]
+                     : cbm_name_rows_lookup(&mod->use_name_rows, NULL, name);
+        return *row >= 0 ? mod->uses_by_name[*row] : NULL;
+    }
+    for (const rs_use_t *use = after ? after->next : mod->uses; use; use = use->next) {
+        cbm_resolve_stats_add(CBM_RESOLVE_STAT_LOOKUP_PROBES, 1);
+        if (use->local && !strcmp(use->local, name))
+            return use;
+    }
+    return NULL;
+}
+
+/* The first or next use whose path ends in this name, in list order. */
+const rs_use_t *rs_use_ending(const rs_mod_t *mod, const rs_use_t *after, const char *name,
+                              int *row) {
+    cbm_resolve_stats_add(CBM_RESOLVE_STAT_LOOKUP_PROBES, 1);
+    if (mod->uses_by_tail) {
+        *row = after ? mod->use_tail_rows.next[*row]
+                     : cbm_name_rows_lookup(&mod->use_tail_rows, NULL, name);
+        return *row >= 0 ? mod->uses_by_tail[*row] : NULL;
+    }
+    for (const rs_use_t *use = after ? after->next : mod->uses; use; use = use->next) {
+        cbm_resolve_stats_add(CBM_RESOLVE_STAT_LOOKUP_PROBES, 1);
+        if (use->nseg && !strcmp(use->segs[use->nseg - 1], name))
+            return use;
+    }
+    return NULL;
+}
+
 void rs_mod_seal(CBMArena *arena, rs_mod_t *mod) {
     rs_mod_index_impls(arena, mod);
+    rs_mod_index_uses(arena, mod);
     for (rs_impl_t *impl = mod->impls; impl; impl = impl->next)
         rs_consts_seal(arena, impl->consts, &impl->sorted_consts, &impl->nconsts);
     for (rs_item_t *item = mod->item_list; item; item = item->next)
@@ -2339,9 +2402,9 @@ static rs_ent_t rs_named_uses(rs_query_t *q, rs_inst_t *inst, const char *name,
                               rs_inst_t *from) {
     rs_ent_t found = rs_ent(RS_ENT_NONE, RS_VIA_NONE);
     bool hidden = false, found_cfg = false;
-    for (const rs_use_t *use = inst->syn->uses; use; use = use->next) {
-        if (!use->local || strcmp(use->local, name) != 0)
-            continue;
+    int row = -1;
+    for (const rs_use_t *use = rs_use_named(inst->syn, NULL, name, &row); use;
+         use = rs_use_named(inst->syn, use, name, &row)) {
         if (!rs_visible(use->vis, inst->syn, from)) {
             hidden = true;
             continue;
@@ -2484,8 +2547,10 @@ static rs_ent_t rs_lookup_steps(rs_query_t *q, rs_inst_t *inst, const char *name
         /* `#[cfg(a)] use x::f; #[cfg(not(a))] fn f() {}`: one name, a different declaration under
          * each configuration, and nothing here says which one applies. Without cfg this does
          * not compile, so only a conditional use can stand beside a local item. */
-        for (const rs_use_t *use = inst->syn->uses; use; use = use->next) {
-            if (!use->has_cfg || !use->local || strcmp(use->local, name) != 0)
+        int row = -1;
+        for (const rs_use_t *use = rs_use_named(inst->syn, NULL, name, &row); use;
+             use = rs_use_named(inst->syn, use, name, &row)) {
+            if (!use->has_cfg)
                 continue;
             rs_ent_t imported = rs_named_uses(q, inst, name, ns, from);
             if (rs_ent_found(&imported) && !rs_ent_same(&ent, &imported))

internal/cbm/resolve/rust_names_internal.h

@@ -216,6 +216,12 @@ typedef struct rs_use {
     struct rs_use *next;
 } rs_use_t;
 
+/* The first or next use of a module binding a local name, and whose path ends in a name. */
+const rs_use_t *rs_use_named(const rs_mod_t *mod, const rs_use_t *after, const char *name,
+                             int *row);
+const rs_use_t *rs_use_ending(const rs_mod_t *mod, const rs_use_t *after, const char *name,
+                              int *row);
+
 /* One syntactic module: a file's top level, an inline `mod x { }`, or the
  * item scope of a block. The same file reached from two crate roots has one
  * syntactic module and two instances. */
@@ -230,6 +236,10 @@ struct rs_mod {
     rs_item_t **items; /* sorted by name after sealing the module */
     int nitems;
     rs_use_t *uses;  /* named leaves */
+    /* A long uses list by local name and by last segment, in list order: every name looked up
+     * and every trait asked about would walk it. */
+    const rs_use_t **uses_by_name, **uses_by_tail;
+    CBMNameRows use_name_rows, use_tail_rows;
     rs_use_t *globs; /* glob leaves */
     rs_impl_t *impls;
     /* A long impls list by start byte, then list position: each impl walked looks itself up. */

tests/test_resolver_rust.c

@@ -1615,6 +1615,73 @@ TEST(resolver_rust_trait_on_external_type_delegating_to_inherent) {
     PASS();
 }
 
+/* Enough impls on outside types to answer from rows by path and method: the same answers as the
+ * chains give, for a crate's Result alias, a delegating impl, two traits on one node, and two
+ * traits on two nodes. The first source stays under the rows' threshold, the second is over it. */
+#define RUST_OUTSIDE_IMPLS                                                                        \
+    "use std::process::Child as StdChild;\n"                                                      \
+    "mod a { pub trait One { fn apart(&self) -> u32; }\n"                                         \
+    " impl One for String { fn apart(&self) -> u32 { 1 } } }\n"                                   \
+    "mod b { pub trait Two { fn apart(&self) -> u32; }\n"                                         \
+    " impl Two for String { fn apart(&self) -> u32 { 2 } } }\n"                                   \
+    "use a::One;\nuse b::Two;\n"                                                                  \
+    "trait WithIndex<T> { fn with_index(self, i: usize) -> Option<T>; }\n"                        \
+    "impl<T, E> WithIndex<T> for Result<T, E> { fn with_index(self, i: usize) -> Option<T> { "    \
+    "self.ok() } }\n"                                                                             \
+    "trait Shout { fn shout(&self) -> String; }\n"                                                \
+    "impl Shout for String { fn shout(&self) -> String { self.clone() } }\n"                      \
+    "trait Wait { fn id(&self) -> u32; fn extra(&self) -> u32; }\n"                               \
+    "impl Wait for StdChild { fn id(&self) -> u32 { self.id() } fn extra(&self) -> u32 { 0 } }\n" \
+    "trait Three { fn near(&self) -> u32; }\n"                                                    \
+    "trait Four { fn near(&self) -> u32; }\n"                                                     \
+    "impl Three for String { fn near(&self) -> u32 { 3 } }\n"                                     \
+    "impl Four for String { fn near(&self) -> u32 { 4 } }\n"                                      \
+    "fn from_alias(r: heed::Result<u32>) { r.with_index(1); }\n"                                  \
+    "fn on_string(s: String) { s.shout(); }\n"                                                    \
+    "fn other_method(s: String) { s.len(); }\n"                                                   \
+    "fn inherent(c: StdChild) { c.id(); }\n"                                                      \
+    "fn trait_only(c: StdChild) { c.extra(); }\n"                                                 \
+    "fn one_node(s: String) { s.near(); }\n"                                                      \
+    "fn two_nodes(s: String) { s.apart(); }\n"                                                    \
+    "fn main() {}\n"
+#define RUST_OUTSIDE_IMPLS_MANIFEST \
+    "[package]\nname = \"demo\"\nedition = \"2021\"\n[dependencies]\nheed = \"1\"\n"
+#define RUST_OUTSIDE_IMPLS_ANSWERS(store)                                                          \
+    do {                                                                                           \
+        ASSERT_NOT_NULL(store);                                                                    \
+        ASSERT_EQ(rust_edge(store, "CALLS", "src.main.from_alias", "src.main.Result.with_index"),  \
+                  1);                                                                              \
+        ASSERT_EQ(rust_edge(store, "CALLS", "src.main.on_string", "src.main.String.shout"), 1);    \
+        ASSERT_EQ(rust_edges_from(store, "CALLS", "src.main.other_method"), 0);                    \
+        ASSERT_EQ(rust_edges_from(store, "CALLS", "src.main.inherent"), 0);                        \
+        ASSERT_EQ(rust_edge(store, "CALLS", "src.main.trait_only", "src.main.StdChild.extra"), 1); \
+        ASSERT_EQ(rust_edges_from(store, "CALLS", "src.main.one_node"), 1);                        \
+        ASSERT_EQ(rust_edges_from(store, "CALLS", "src.main.two_nodes"), 0);                       \
+        ASSERT_TRUE(                                                                               \
+            rust_occurrence_reason(store, "CALLS", "c.id", "rust_external_receiver_method"));      \
+        ASSERT_TRUE(                                                                               \
+            rust_occurrence_reason(store, "CALLS", "s.apart", "rust_target_qn_collision"));        \
+    } while (0)
+TEST(resolver_rust_many_outside_impls_answer_as_few_do) {
+    static const rust_proof_file_t few[] = {{"Cargo.toml", RUST_OUTSIDE_IMPLS_MANIFEST},
+                                            {"src/main.rs", RUST_OUTSIDE_IMPLS}};
+    static const rust_proof_file_t many[] = {
+        {"Cargo.toml", RUST_OUTSIDE_IMPLS_MANIFEST},
+        {"src/main.rs",
+         "trait Pad { fn p0(&self); fn p1(&self); fn p2(&self); fn p3(&self); fn p4(&self);"
+         " fn p5(&self); fn p6(&self); fn p7(&self); fn p8(&self); }\n"
+         "impl Pad for u8 { fn p0(&self) {} fn p1(&self) {} fn p2(&self) {} fn p3(&self) {}"
+         " fn p4(&self) {} fn p5(&self) {} fn p6(&self) {} fn p7(&self) {} fn p8(&self) {} "
+         "}\n" RUST_OUTSIDE_IMPLS}};
+    RUST_ROUTES(store, root, few) {
+        RUST_OUTSIDE_IMPLS_ANSWERS(store);
+    }
+    RUST_ROUTES(store, root, many) {
+        RUST_OUTSIDE_IMPLS_ANSWERS(store);
+    }
+    PASS();
+}
+
 /* tokio's blocking.rs: `cfg_rt! { pub(crate) use crate::runtime::spawn_blocking; }` beside
  * `cfg_not_rt! { pub(crate) fn spawn_blocking() {} }`. Each configuration has a different declaration,
  * so neither is the answer. A plain #[cfg] twin is the same case. */
@@ -2629,6 +2696,7 @@ SUITE(resolver_rust) {
     RUN_TEST(resolver_rust_unrepaired_syntax_error_stays_unresolved);
     RUN_TEST(resolver_rust_spelled_external_type_takes_a_project_trait);
     RUN_TEST(resolver_rust_trait_on_external_type_delegating_to_inherent);
+    RUN_TEST(resolver_rust_many_outside_impls_answer_as_few_do);
     RUN_TEST(resolver_rust_conditional_use_beside_a_local_item_is_ambiguous);
     RUN_TEST(resolver_rust_wrapper_keeps_into_inner_for_itself);
     RUN_TEST(resolver_rust_pin_project_declares_its_struct);

tests/test_resolver_ts.c

@@ -665,6 +665,27 @@ TEST(resolver_ts_duplicate_method_declarations_abstain) {
     PASS();
 }
 
+/* Sorted by start, the nested function's narrowing is the last one before the final call, and it
+ * ended before that call. The outer one still covers it; nothing covers the function after them. */
+TEST(resolver_ts_narrowing_outlives_a_nested_one) {
+    static const char source[] = "class C { run() {} } class D extends C { run() {} }\n"
+                                 "function outer(a: C) {\n"
+                                 "  if (a instanceof D) { a.run(); }\n"
+                                 "  function inner(a: C) { if (a instanceof D) { a.run(); } }\n"
+                                 "  a.run();\n"
+                                 "}\n"
+                                 "function after(a: C) { a.run(); }\n";
+    CBMFileResult *file = cbm_extract_file(source, (int)strlen(source), CBM_LANG_TYPESCRIPT, "test",
+                                           "main.ts", 0, NULL, NULL);
+    ASSERT_NOT_NULL(file);
+    int declared = ts_proved_count(&file->resolved_calls, "test.main.C.run");
+    int narrower = ts_proved_count(&file->resolved_calls, "test.main.D.run");
+    cbm_free_result(file);
+    ASSERT_EQ(declared, 1);
+    ASSERT_EQ(narrower, 0);
+    PASS();
+}
+
 SUITE(resolver_ts) {
     RUN_TEST(resolver_ts_declared_receiver_selects_exact_occurrence);
     RUN_TEST(resolver_ts_unknown_union_and_later_shadow_abstain);
@@ -684,6 +705,7 @@ SUITE(resolver_ts) {
     RUN_TEST(resolver_ts_annotation_type_refs_are_exact_occurrences);
     RUN_TEST(resolver_ts_this_of_a_literal_method_is_not_the_class);
     RUN_TEST(resolver_ts_imported_type_requires_export_binding);
+    RUN_TEST(resolver_ts_narrowing_outlives_a_nested_one);
     RUN_TEST(resolver_ts_overlay_copies_coalesce_only_same_declaration);
     RUN_TEST(resolver_ts_duplicate_method_declarations_abstain);
 }