Merge pull request #2419 from Fieldnote-Echo/perf/py-lexical-ancestor-kind

perf(python): climb the walk cursor for scope-directive and default checks

internal/cbm/extract_usages.c

@@ -1789,9 +1789,25 @@ static uint32_t active_lexical_scope_id(const WalkState *state) {
 static uint32_t lexical_ancestor_of_kind(const WalkState *state, uint32_t start_id,
                                          bool want_function, bool want_block);
 
-static bool python_default_value_reference(TSNode node) {
-    for (TSNode parent = ts_node_parent(node); !ts_node_is_null(parent);
-         parent = ts_node_parent(parent)) {
+/* One step up: on the occurrence cursor when there is one, else with
+ * ts_node_parent, which descends from the ROOT, so a climb costs O(depth^2). */
+static TSNode usage_parent_step(TSTreeCursor *cursor, TSNode current) {
+    if (cursor) {
+        return ts_tree_cursor_goto_parent(cursor) ? ts_tree_cursor_current_node(cursor)
+                                                  : (TSNode){0};
+    }
+    usage_slow_parent_fallback_test_note();
+    return ts_node_parent(current);
+}
+
+/* Asked for every Python usage, and a deep expression puts O(depth) parents
+ * above each one. A tree without errors climbs them on the walk cursor, like
+ * python_scope_directive_kind; a tree with errors keeps ts_node_parent. */
+static bool python_default_value_reference(CBMExtractCtx *ctx, WalkState *state, TSNode node) {
+    TSTreeCursor *cursor =
+        ts_node_has_error(ctx->root) ? NULL : reset_occurrence_cursor(state, node);
+    for (TSNode parent = usage_parent_step(cursor, node); !ts_node_is_null(parent);
+         parent = usage_parent_step(cursor, parent)) {
         const char *kind = ts_node_type(parent);
         /* A nested executable scope inside the default owns its own lookups;
          * only a direct default expression is evaluated in the declaring
@@ -1812,10 +1828,10 @@ static bool python_default_value_reference(TSNode node) {
     return false;
 }
 
-static uint32_t usage_lexical_scope_id_for_node(CBMExtractCtx *ctx, const WalkState *state,
-                                                TSNode node) {
+static uint32_t usage_lexical_scope_id_for_node(CBMExtractCtx *ctx, WalkState *state, TSNode node) {
     uint32_t scope_id = active_lexical_scope_id(state);
-    if (!ctx || ctx->language != CBM_LANG_PYTHON || !python_default_value_reference(node)) {
+    if (!ctx || ctx->language != CBM_LANG_PYTHON ||
+        !python_default_value_reference(ctx, state, node)) {
         return scope_id;
     }
     uint32_t function_id = lexical_ancestor_of_kind(state, scope_id, true, false);
@@ -1891,8 +1907,8 @@ static bool ensure_lexical_binding_capacity(WalkState *state) {
 }
 
 static bool lexical_ancestor_kind(TSNode node, const char *kind) {
-    for (TSNode parent = ts_node_parent(node); !ts_node_is_null(parent);
-         parent = ts_node_parent(parent)) {
+    for (TSNode parent = usage_parent_step(NULL, node); !ts_node_is_null(parent);
+         parent = usage_parent_step(NULL, parent)) {
         if (strcmp(ts_node_type(parent), kind) == 0) {
             return true;
         }
@@ -1900,6 +1916,34 @@ static bool lexical_ancestor_kind(TSNode node, const char *kind) {
     return false;
 }
 
+/* Is `node` a name listed by a Python global/nonlocal statement? The grammar
+ * makes those names direct children of the statement, so in a tree without
+ * errors one parent step on the walk cursor answers it. The full climb above
+ * steps with ts_node_parent, which descends from the ROOT each time: O(depth^2)
+ * per identifier, asked for every Python identifier, so one deep arithmetic
+ * expression went cubic. Error recovery can put a name under an ERROR child of
+ * the statement, so a tree with errors keeps the full climb. */
+static CBMPythonDirectiveKind python_scope_directive_kind(CBMExtractCtx *ctx, WalkState *state,
+                                                          TSNode node) {
+    if (ts_node_has_error(ctx->root)) {
+        if (lexical_ancestor_kind(node, "global_statement")) {
+            return CBM_PYTHON_DIRECTIVE_GLOBAL;
+        }
+        return lexical_ancestor_kind(node, "nonlocal_statement") ? CBM_PYTHON_DIRECTIVE_NONLOCAL
+                                                                 : 0;
+    }
+    TSNode parent;
+    const char *field;
+    if (!occurrence_parent(reset_occurrence_cursor(state, node), node, &parent, &field)) {
+        return 0;
+    }
+    const char *kind = ts_node_type(parent);
+    if (strcmp(kind, "global_statement") == 0) {
+        return CBM_PYTHON_DIRECTIVE_GLOBAL;
+    }
+    return strcmp(kind, "nonlocal_statement") == 0 ? CBM_PYTHON_DIRECTIVE_NONLOCAL : 0;
+}
+
 static TSNode declared_function_name_owner(TSNode node, const CBMLangSpec *spec) {
     for (TSNode parent = ts_node_parent(node); !ts_node_is_null(parent);
          parent = ts_node_parent(parent)) {
@@ -2246,13 +2290,11 @@ static void record_lexical_binding(CBMExtractCtx *ctx, WalkState *state, TSNode
     TSNode function_declaration = declared_function_name_owner(node, spec);
     TSNode class_declaration = declared_class_name_owner(node, spec);
 
-    if (ctx->language == CBM_LANG_PYTHON && (lexical_ancestor_kind(node, "global_statement") ||
-                                             lexical_ancestor_kind(node, "nonlocal_statement"))) {
+    CBMPythonDirectiveKind scope_directive =
+        ctx->language == CBM_LANG_PYTHON ? python_scope_directive_kind(ctx, state, node) : 0;
+    if (scope_directive) {
         uint32_t function_id = lexical_ancestor_of_kind(state, current_id, true, false);
-        CBMPythonDirectiveKind directive = lexical_ancestor_kind(node, "global_statement")
-                                               ? CBM_PYTHON_DIRECTIVE_GLOBAL
-                                               : CBM_PYTHON_DIRECTIVE_NONLOCAL;
-        record_python_directive(state, function_id, name, directive);
+        record_python_directive(state, function_id, name, scope_directive);
         return;
     }
 
@@ -2719,8 +2761,7 @@ void handle_usages(CBMExtractCtx *ctx, TSNode node, const CBMLangSpec *spec, Wal
     }
 
     bool python_scope_directive =
-        ctx->language == CBM_LANG_PYTHON && (lexical_ancestor_kind(node, "global_statement") ||
-                                             lexical_ancestor_kind(node, "nonlocal_statement"));
+        ctx->language == CBM_LANG_PYTHON && python_scope_directive_kind(ctx, state, node) != 0;
     if (python_scope_directive || is_binding_occurrence(ctx, node, spec, state)) {
         char *binding_name = reference_name(ctx, node);
         if (binding_name && binding_name[0] && !cbm_is_keyword(binding_name, ctx->language)) {

tests/repro/repro_lexical_binding_precision.c

@@ -963,6 +963,59 @@ TEST(repro_python_global_assignment_is_module_binding) {
     PASS();
 }
 
+/* A global/nonlocal name is a directive, not a read (each name of `global a,
+ * b`), and the rebinding lands where it says: the module from that point on,
+ * or the enclosing function. Source 1 is an error tree, whose directive and
+ * default-value checks take the ts_node_parent fallback. */
+TEST(repro_python_scope_directives_bind_the_named_scope) {
+    static const char *const sources[] = {
+        "def handler():\n    pass\n"
+        "def outer():\n    handler = None\n"
+        "    def inner():\n        nonlocal handler\n        handler = outer\n"
+        "def early():\n    accept_early(handler)\n"
+        "def caller():\n    def nested():\n        global handler, other\n"
+        "        handler = outer\n    accept_outer(handler)\n"
+        "def sibling():\n    accept_sibling(handler)\n",
+        "def handler():\n    pass\n"
+        "def make(handler=handler):\n    return handler\n"
+        "def caller():\n    global handler other\n    handler = make\n    make(handler)\n",
+    };
+    static const struct {
+        int source;
+        const char *marker, *name;
+        int count, blocked, shadow;
+    } sites[] = {
+        {0, "global handler", "handler", 0, 0, 0},
+        {0, ", other", "other", 0, 0, 0},
+        {0, "nonlocal handler", "handler", 0, 0, 0},
+        {0, "accept_early(handler)", "handler", 1, 0, 0},
+        {0, "accept_outer(handler)", "handler", 1, 1, 0},
+        {0, "accept_sibling(handler)", "handler", 1, 1, 0},
+        {1, "global handler", "handler", 0, 0, 0},
+        {1, "=handler", "handler", 1, 0, 0},
+        {1, "make(handler)", "handler", 1, 1, 0},
+    };
+    int failures = 0;
+    for (int i = 0; i < (int)(sizeof(sites) / sizeof(sites[0])); i++) {
+        const char *source = sources[sites[i].source];
+        uint32_t start = lb_identifier_offset(source, sites[i].marker, sites[i].name);
+        ASSERT_NEQ(start, UINT32_MAX);
+        CBMFileResult *result = lb_extract(source, CBM_LANG_PYTHON, "directives.py");
+        ASSERT_NOT_NULL(result);
+        ASSERT_EQ(result->parse_incomplete, sites[i].source);
+        LBBindingSite got = lb_binding_site(result, NULL, sites[i].name, start);
+        cbm_free_result(result);
+        if (got.count != sites[i].count || got.blocked_count != sites[i].blocked ||
+            got.local_shadow_count != sites[i].shadow) {
+            printf("  %s: count %d blocked %d shadow %d\n", sites[i].marker, got.count,
+                   got.blocked_count, got.local_shadow_count);
+            failures++;
+        }
+    }
+    ASSERT_EQ(failures, 0);
+    PASS();
+}
+
 /* Module assignments participate in name lookup too. A later function must
  * not raw-name-fallback through a module variable to an older same-named
  * callable declaration. */
@@ -1055,6 +1108,7 @@ SUITE(repro_lexical_binding_precision) {
     RUN_TEST(repro_objectscript_udl_set_target_binds_method_local);
     RUN_TEST(repro_python_default_expression_uses_enclosing_namespace);
     RUN_TEST(repro_python_global_assignment_is_module_binding);
+    RUN_TEST(repro_python_scope_directives_bind_the_named_scope);
     RUN_TEST(repro_python_module_assignment_blocks_callable_fallback);
     RUN_TEST(repro_go_function_literal_parameter_does_not_leak);
 }

tests/test_extraction.c

@@ -10933,6 +10933,49 @@ TEST(extract_rescript_let_bindings_use_the_walk_cursor) {
     ASSERT_EQ(big_fallbacks, 0);
     PASS();
 }
+
+/* Every Python usage asks whether it is a global/nonlocal name and whether it
+ * sits in a parameter default. Both climbed with ts_node_parent, which
+ * descends from the ROOT: O(d^2) for a usage d deep, so one deep expression
+ * went cubic (2.2 s / 17 s / >90 s at 500/1000/2000 terms). Count every such
+ * root-descending step at N and 2N terms. */
+TEST(extract_python_deep_default_parent_climbs_are_linear) {
+    enum { SMALL = 200 };
+    uint64_t slow_parents[2];
+    for (int round = 0; round < 2; round++) {
+        /* def f(x=x + x + ... + x): return x -- a binary_operator chain. */
+        int terms = SMALL << round;
+        size_t capacity = 32U + (size_t)terms * 4U;
+        char *source = malloc(capacity);
+        ASSERT_NOT_NULL(source);
+        int length = snprintf(source, capacity, "def f(x=x");
+        for (int i = 1; i < terms; i++) {
+            length += snprintf(source + length, capacity - (size_t)length, " + x");
+        }
+        length += snprintf(source + length, capacity - (size_t)length, "):\n    return x\n");
+        cbm_usage_field_lookup_test_reset();
+        CBMFileResult *result =
+            cbm_extract_file(source, length, CBM_LANG_PYTHON, "proj", "deep.py", 0, NULL, NULL);
+        free(source);
+        ASSERT_NOT_NULL(result);
+        slow_parents[round] = cbm_usage_slow_parent_fallback_test_count();
+        /* Anti-vacuous: every default x climbed its chain to the parameter and
+         * reads the enclosing namespace; only `return x` is shadowed by it. */
+        int defaults = 0;
+        for (int i = 0; i < result->usages.count; i++) {
+            const CBMUsage *usage = &result->usages.items[i];
+            defaults += usage->ref_name && strcmp(usage->ref_name, "x") == 0 &&
+                        !usage->semantic_reference_local_shadow;
+        }
+        cbm_free_result(result);
+        ASSERT_EQ(defaults, terms);
+    }
+    fprintf(stderr, "  [python-deep-default] slow_parents(%d)=%llu slow_parents(%d)=%llu\n", SMALL,
+            (unsigned long long)slow_parents[0], 2 * SMALL, (unsigned long long)slow_parents[1]);
+    /* Twice the terms may cost about twice the steps; the cubic climb took 4x. */
+    ASSERT_LTE(slow_parents[1], slow_parents[0] * 3U + 64U);
+    PASS();
+}
 #endif
 
 /* ===================================================================
@@ -13455,6 +13498,7 @@ SUITE(extraction) {
     RUN_TEST(extract_wide_flat_reference_fields_are_linear);
     RUN_TEST(extract_csharp_argument_values_use_the_walk_cursor);
     RUN_TEST(extract_rescript_let_bindings_use_the_walk_cursor);
+    RUN_TEST(extract_python_deep_default_parent_climbs_are_linear);
 #endif
 
     /* Perl call-graph noise (#459 follow-up) */