fix: restore memoization on _count_concat to stop count_words explosion

During the AST migration, _count_concat lost its @lru_cache (it was
replaced by an uncached recursive version). Distributing a length L
across concat parts then revisited the same (remaining_parts, length)
states exponentially -> 16M function calls for a single 4-method group,
and several RAGSAK groups took 18-52s (looked hung at the tail).

Restore memoization: _count_concat is now @lru_cache-keyed on
(tuple(parts), length). Also fix a stray blank line between the
@lru_cache decorator and count_words.

Impact (RAGSAK, --slice package --min-structure 0.5):
  straggler groups 18-52s -> <0.5s; full run 54.9s -> 6.7s.
Adds a regression test asserting a length-20 (a|b)* string counts in <2s.
This commit is contained in:
tobjend 2026-07-13 00:03:40 +02:00
parent f8dda557d2
commit 52f286a073
2 changed files with 26 additions and 6 deletions

View file

@ -173,7 +173,6 @@ def _match_rep(child, seq, pos, min_rep):
_COUNT_CAP = 10**12
@lru_cache(maxsize=None)
def count_words(node, length):
"""Count how many words of exactly `length` are in L(node).
@ -189,7 +188,7 @@ def count_words(node, length):
if isinstance(node, Empty):
return 0
if isinstance(node, Concat):
return _count_concat(tuple(id(p) for p in node.parts), node, length, 0)
return _count_concat(tuple(node.parts), length)
if isinstance(node, Alt):
total = 0
for p in node.parts:
@ -206,14 +205,17 @@ def count_words(node, length):
return 0
def _count_concat(part_ids, node, length, idx):
if idx >= len(node.parts):
@lru_cache(maxsize=None)
def _count_concat(parts, length):
if not parts:
return 1 if length == 0 else 0
first = parts[0]
rest = parts[1:]
total = 0
for take in range(length + 1):
cnt = count_words(node.parts[idx], take)
cnt = count_words(first, take)
if cnt:
total += cnt * _count_concat(part_ids, node, length - take, idx + 1)
total += cnt * _count_concat(rest, length - take)
if total >= _COUNT_CAP:
return _COUNT_CAP
return total

View file

@ -186,6 +186,24 @@ def test_to_gbnf_with_rules():
assert result == 'choice ::= "a" | "b"'
def test_count_words_concat_is_memoized():
"""Regression: _count_concat must be memoized.
An uncached implementation revisits (remaining_parts, length) states
exponentially; a long sequence over a multi-part Concat explodes into
millions of calls. With memoization it completes instantly.
"""
import time
# length-L string of (a|b) alternatives -> exactly 2^L words of length L
L = 20
g = Concat([Alt([Symbol('a'), Symbol('b')]) for _ in range(L)])
t0 = time.time()
n = count_words(g, L)
dt = time.time() - t0
assert dt < 2.0, f"count_words was too slow ({dt:.2f}s) — memoization lost"
assert n == 2 ** L, f"expected 2^{L}={2**L} words, got {n}"
# ── Structure score ──
class TestStructureScore: