diff --git a/graphify/extract.py b/graphify/extract.py index 7cc9e62c90..cc2bab2dee 100644 --- a/graphify/extract.py +++ b/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 diff --git a/tests/test_ts_import_type_arguments.py b/tests/test_ts_import_type_arguments.py index 883cca19d2..097caf1ca5 100644 --- a/tests/test_ts_import_type_arguments.py +++ b/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(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)" + )