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

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: