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

`_normalize_ts_import_types` filtered each `import(...)` regex match against
every tree-sitter `type_arguments` byte range with a nested linear scan, in
both the pre-mask and post-mask passes:

    matches = [m for m in matches
               if not any(start <= m.start() < end
                          for start, end in original_type_ranges)]

That is O(matches x ranges), and `m.start()` is re-evaluated inside `any()`
for every range — which is why profiles blamed `re.Match.start`/`_sre` and
`sample` appeared to show a pathological regex. The regexes are all linear;
the quadratic caller was the real cost.

Both counts grow with file size in TypeScript that mixes generic calls
(`fn<A, B>(x)` — one `type_arguments` range each) with `import(...)` types,
which is the shape of generated and barrel modules in a large monorepo. On
the filter alone: a 0.8MB/10k-import file took ~21s, 2.1MB/25k ~2.6min, and
4.2MB/50k ~11.8min. With several such files `extract` never finished, and a
worker sat at 100% CPU holding no open files because the work is entirely
in-memory after the read.

Replace both scans with a binary search over start-sorted ranges plus a
prefix-maximum of `end` values, so a lookup touches nesting depth rather
than every range. Type-argument ranges nest but never partially overlap, so
the prefix maximum is a correct stop condition; a *closed* inner sibling
must not terminate the leftward walk, which a naive running-max break gets
wrong. Verified against the old predicate with 500k differential cases,
including non-nested random ranges: zero mismatches.

Normalization output is unchanged — byte-length and line-count preserving,
same masking decisions, and the runtime-vs-type disambiguation from 6baa5e0
still holds.

Also restore the emptiness guard in the second pass: the range list is now
an index tuple, which is always truthy, so it checks the spans it holds.

The new test asserts on scaling (2x input must not ~4x the time) rather than
wall-clock, so it is not machine-dependent. It fails at ~3.4x before this
change and passes after.

(cherry picked from commit b662361d8b)
This commit is contained in:
Sagexd08 authored and safishamsi committed 2026-09-07 20:21:35 +01:00
1 parent 98d62e231c
commit 1dcb1e1f3a
2 files changed
+98 -8

No files matched your search

+60 -8
View File
@@ -8,6 +8,7 @@ import os
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
@@ -1361,6 +1362,58 @@ 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.
_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] = []
@@ -1451,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
@@ -1476,19 +1529,18 @@ def _normalize_ts_import_types(source: bytes, *, tsx: bool = False) -> bytes | N
# 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
+38
View File
@@ -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)"
)