diff --git a/bex/grammar.py b/bex/grammar.py index 44e5dae..ab4db66 100644 --- a/bex/grammar.py +++ b/bex/grammar.py @@ -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 diff --git a/tests/test_grammar.py b/tests/test_grammar.py index 2539376..248769f 100644 --- a/tests/test_grammar.py +++ b/tests/test_grammar.py @@ -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: