fix(extract): remove quadratic scan in TS import-type normalization (#3359)

graphify/extract.py

@@ -8,6 +8,7 @@
 import re
 import sys
 import textwrap
+from bisect import bisect_right
 from collections import Counter
 from dataclasses import dataclass, field
 from pathlib import Path, PurePath
@@ -1360,6 +1361,59 @@ def _ts_type_argument_ranges(root: Any, *, call_only: bool) -> list[tuple[int, i
     return ranges
 
 
+# Sorted type-argument ranges prepared for lookup: (starts, spans, prefix_max_end).
+# See _ts_build_range_index. A real alias, not a string: quoting the right-hand
+# side would make this a `str` value, which type checkers reject as a type.
+_TsRangeIndex = tuple[list[int], list[tuple[int, int]], list[int]]
+
+
+def _ts_ranges_containing(
+    ranges: _TsRangeIndex, offset: int
+) -> list[tuple[int, int]]:
+    """Return the ranges in ``ranges`` that contain ``offset``.
+
+    Scanning the full range list for every match made
+    :func:`_normalize_ts_import_types` O(matches x ranges); on a large monorepo
+    file with thousands of generic calls and thousands of ``import(...)`` types
+    that quadratic blowup pinned a worker at 100% CPU with no output and no file
+    I/O, so extraction never finished (#3359).
+
+    Every range containing ``offset`` must start at or before it, so binary-search
+    to the last such start and walk left. A range that ends before ``offset`` is
+    not itself a hit but must not stop the walk -- a closed sibling can sit inside
+    an enclosing range that does contain ``offset``. The index's running maximum of
+    ends (``prefix_max_end``) gives the correct stop: once no range at or before
+    this position reaches past ``offset``, none of the remaining ones can.
+    """
+    starts, spans, prefix_max_end = ranges
+    hits: list[tuple[int, int]] = []
+    index = bisect_right(starts, offset) - 1
+    while index >= 0 and prefix_max_end[index] > offset:
+        start, end = spans[index]
+        if end > offset:
+            hits.append((start, end))
+        index -= 1
+    hits.reverse()
+    return hits
+
+
+def _ts_build_range_index(root: Any, *, call_only: bool) -> _TsRangeIndex:
+    """Index :func:`_ts_type_argument_ranges` for :func:`_ts_ranges_containing`.
+
+    Returns the ranges sorted by start, split into a bare ``starts`` list for
+    :func:`bisect_right`, plus a prefix-maximum of ``end`` values so a lookup can
+    tell when no earlier range can still reach a given offset.
+    """
+    spans = sorted(_ts_type_argument_ranges(root, call_only=call_only))
+    starts = [start for start, _ in spans]
+    prefix_max_end: list[int] = []
+    running = -1
+    for _, end in spans:
+        running = max(running, end)
+        prefix_max_end.append(running)
+    return starts, spans, prefix_max_end
+
+
 def _ts_error_nodes(root: Any) -> list[Any]:
     """Return parser error nodes without depending on a grammar's error name."""
     errors: list[Any] = []
@@ -1450,10 +1504,10 @@ def _normalize_ts_import_types(source: bytes, *, tsx: bool = False) -> bytes | N
     # its syntax is parseable as written. This includes the grammar's deliberate
     # comparison-vs-generic ambiguity (`a < b, import("...") > (d)`); retaining
     # that source is essential because it is a runtime expression, not a type.
-    original_type_ranges = _ts_type_argument_ranges(original_root, call_only=False)
+    original_type_ranges = _ts_build_range_index(original_root, call_only=False)
     matches = [
         match for match in matches
-        if not any(start <= match.start() < end for start, end in original_type_ranges)
+        if not _ts_ranges_containing(original_type_ranges, match.start())
     ]
     if not matches:
         return None
@@ -1475,19 +1529,18 @@ def placeholder(match: "re.Match[bytes]") -> bytes:
     # type annotations already parse ``import(...)`` correctly and should keep
     # their native AST shape. A range from an outer call's type_arguments also
     # covers nested generic arguments.
-    type_argument_ranges = _ts_type_argument_ranges(root, call_only=True)
+    type_argument_ranges = _ts_build_range_index(root, call_only=True)
     errors = _ts_error_nodes(original_root)
 
-    if not type_argument_ranges:
+    # `type_argument_ranges` is an index tuple, so test the spans it holds --
+    # the tuple itself is always truthy.
+    if not type_argument_ranges[1]:
         return None
 
     norm = bytearray(source)
     changed = False
     for match in matches:
-        containing_ranges = [
-            candidate for candidate in type_argument_ranges
-            if candidate[0] <= match.start() < candidate[1]
-        ]
+        containing_ranges = _ts_ranges_containing(type_argument_ranges, match.start())
         if containing_ranges and any(
             _ts_mask_candidate_is_malformed(original_root, candidate, errors)
             for candidate in containing_ranges

tests/test_ts_import_type_arguments.py

@@ -285,3 +285,41 @@ def test_ts_normalizer_does_not_mask_import_text_in_literals_or_comments():
     assert b'import(\\\"./text\\\")' in normalized
     assert b'import("./comment")' in normalized
     assert b'import("./types")' not in normalized
+
+
+def test_ts_normalizer_scales_linearly_on_large_files():
+    """#3359: the match/type-argument-range filters must not be O(matches x ranges).
+
+    A large monorepo file mixes thousands of generic calls (each a
+    ``type_arguments`` range) with thousands of ``import(...)`` types. Filtering
+    each match by scanning every range made extraction quadratic, so `extract`
+    spun at 100% CPU in the regex/scan path and never finished.
+
+    Assert on scaling, not wall-clock, so the test is not machine-dependent:
+    doubling the input must not roughly quadruple the time.
+    """
+    import time
+
+    def build(n: int) -> bytes:
+        lines = []
+        for i in range(n):
+            lines.append(f"const a{i} = fn<Foo{i}, Bar{i}>(x);")
+            lines.append(f"type T{i} = import('./m{i}').Thing;")
+        return "\n".join(lines).encode()
+
+    def timed(n: int) -> float:
+        source = build(n)
+        start = time.perf_counter()
+        _normalize_ts_import_types(source)
+        return time.perf_counter() - start
+
+    timed(200)  # warm the grammar/parser import off the measured path
+    small = min(timed(1000) for _ in range(3))
+    large = min(timed(2000) for _ in range(3))
+
+    # Linear work doubles (~2x). Quadratic work quadruples (~4x). A generous
+    # 3x ceiling separates the two without being flaky under load.
+    assert large < small * 3, (
+        f"scaling looks super-linear: {small:.4f}s -> {large:.4f}s "
+        f"({large / small:.1f}x for 2x input)"
+    )