52f286a073
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.
222 lines
6.2 KiB
Python
222 lines
6.2 KiB
Python
"""Tests for grammar AST, matching, counting, and GBNF rendering."""
|
|
|
|
import pytest
|
|
from bex.grammar import (
|
|
Symbol, Concat, Alt, Plus, Optional, Star, Epsilon, Empty,
|
|
alphabet, count_words, lang_size, model_cost, match,
|
|
)
|
|
from bex.gbnf import to_gbnf, to_gbnf_with_rules, grammar_structure_score
|
|
|
|
|
|
# ── AST Construction ──
|
|
|
|
class TestASTConstruction:
|
|
def test_symbol(self):
|
|
assert Symbol('a') == Symbol('a')
|
|
assert Symbol('a') != Symbol('b')
|
|
|
|
def test_concat(self):
|
|
c = Concat([Symbol('a'), Symbol('b')])
|
|
assert c.parts == [Symbol('a'), Symbol('b')]
|
|
|
|
def test_alt(self):
|
|
a = Alt([Symbol('a'), Symbol('b')])
|
|
assert a.parts == [Symbol('a'), Symbol('b')]
|
|
|
|
def test_plus(self):
|
|
p = Plus(Symbol('a'))
|
|
assert p.child == Symbol('a')
|
|
|
|
def test_optional(self):
|
|
o = Optional(Symbol('a'))
|
|
assert o.child == Symbol('a')
|
|
|
|
def test_star(self):
|
|
s = Star(Symbol('a'))
|
|
assert s.child == Symbol('a')
|
|
|
|
def test_epsilon(self):
|
|
assert Epsilon() == Epsilon()
|
|
|
|
def test_empty(self):
|
|
assert Empty() == Empty()
|
|
|
|
|
|
# ── Alphabet ──
|
|
|
|
class TestAlphabet:
|
|
def test_symbol(self):
|
|
assert alphabet(Symbol('a')) == {'a'}
|
|
|
|
def test_concat(self):
|
|
assert alphabet(Concat([Symbol('a'), Symbol('b')])) == {'a', 'b'}
|
|
|
|
def test_nested(self):
|
|
g = Concat([Symbol('a'), Plus(Alt([Symbol('b'), Symbol('c')]))])
|
|
assert alphabet(g) == {'a', 'b', 'c'}
|
|
|
|
|
|
# ── Matching ──
|
|
|
|
class TestMatch:
|
|
def test_symbol(self):
|
|
assert match(Symbol('a'), ['a'])
|
|
assert not match(Symbol('a'), ['b'])
|
|
assert not match(Symbol('a'), [])
|
|
|
|
def test_concat(self):
|
|
g = Concat([Symbol('a'), Symbol('b'), Symbol('c')])
|
|
assert match(g, ['a', 'b', 'c'])
|
|
assert not match(g, ['a', 'b'])
|
|
assert not match(g, ['a', 'b', 'c', 'd'])
|
|
|
|
def test_alt(self):
|
|
g = Alt([Symbol('a'), Symbol('b')])
|
|
assert match(g, ['a'])
|
|
assert match(g, ['b'])
|
|
assert not match(g, ['c'])
|
|
|
|
def test_plus(self):
|
|
g = Plus(Symbol('a'))
|
|
assert match(g, ['a'])
|
|
assert match(g, ['a', 'a'])
|
|
assert match(g, ['a', 'a', 'a'])
|
|
assert not match(g, [])
|
|
assert not match(g, ['b'])
|
|
|
|
def test_optional(self):
|
|
g = Optional(Symbol('a'))
|
|
assert match(g, ['a'])
|
|
assert match(g, [])
|
|
assert not match(g, ['b'])
|
|
|
|
def test_star(self):
|
|
g = Star(Symbol('a'))
|
|
assert match(g, [])
|
|
assert match(g, ['a'])
|
|
assert match(g, ['a', 'a'])
|
|
assert not match(g, ['b'])
|
|
|
|
def test_complex_grammar(self):
|
|
g = Concat([Plus(Symbol('init')), Plus(Symbol('capability')),
|
|
Concat([Plus(Symbol('invoke')), Symbol('request')])])
|
|
assert match(g, ['init', 'capability', 'invoke', 'request'])
|
|
assert match(g, ['init', 'init', 'capability', 'capability', 'invoke', 'request'])
|
|
assert not match(g, ['init'])
|
|
|
|
|
|
# ── Counting ──
|
|
|
|
class TestCountWords:
|
|
def test_symbol(self):
|
|
assert count_words(Symbol('a'), 1) == 1
|
|
assert count_words(Symbol('a'), 0) == 0
|
|
|
|
def test_epsilon(self):
|
|
assert count_words(Epsilon(), 0) == 1
|
|
assert count_words(Epsilon(), 1) == 0
|
|
|
|
def test_empty(self):
|
|
assert count_words(Empty(), 0) == 0
|
|
|
|
def test_concat(self):
|
|
g = Concat([Symbol('a'), Symbol('b'), Symbol('c')])
|
|
assert count_words(g, 3) == 1
|
|
assert count_words(g, 2) == 0
|
|
|
|
def test_alt(self):
|
|
g = Alt([Symbol('a'), Symbol('b'), Symbol('c')])
|
|
assert count_words(g, 1) == 3
|
|
|
|
def test_plus(self):
|
|
g = Plus(Symbol('a'))
|
|
assert count_words(g, 1) == 1
|
|
assert count_words(g, 2) == 1
|
|
assert count_words(g, 0) == 0
|
|
|
|
def test_optional(self):
|
|
g = Optional(Symbol('a'))
|
|
assert count_words(g, 0) == 1
|
|
assert count_words(g, 1) == 1
|
|
|
|
def test_star(self):
|
|
g = Star(Symbol('a'))
|
|
assert count_words(g, 0) == 1
|
|
assert count_words(g, 1) == 1
|
|
|
|
def test_disjunction_plus(self):
|
|
g = Plus(Alt([Symbol('a'), Symbol('b'), Symbol('c')]))
|
|
assert count_words(g, 1) == 3
|
|
assert count_words(g, 2) == 9
|
|
|
|
def test_info_plus(self):
|
|
g = Plus(Symbol('info'))
|
|
for l in range(1, 8):
|
|
assert count_words(g, l) == 1
|
|
|
|
|
|
# ── Model cost ──
|
|
|
|
class TestModelCost:
|
|
def test_symbol(self):
|
|
assert model_cost(Symbol('a')) == 1
|
|
|
|
def test_concat(self):
|
|
assert model_cost(Concat([Symbol('a'), Symbol('b'), Symbol('c')])) == 3
|
|
|
|
def test_plus(self):
|
|
assert model_cost(Plus(Symbol('a'))) == 1
|
|
|
|
def test_nested(self):
|
|
g = Concat([Symbol('a'), Plus(Alt([Symbol('b'), Symbol('c')]))])
|
|
assert model_cost(g) == 3
|
|
|
|
|
|
# ── GBNF Rendering ──
|
|
|
|
def test_to_gbnf():
|
|
g = Plus(Symbol('init'))
|
|
result = to_gbnf(g)
|
|
assert result == 'root ::= "init"+'
|
|
|
|
|
|
def test_to_gbnf_with_rules():
|
|
g = Alt([Symbol('a'), Symbol('b')])
|
|
result = to_gbnf_with_rules(g, name='choice')
|
|
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:
|
|
def test_empty(self):
|
|
assert grammar_structure_score(Empty()) == 0.0
|
|
|
|
def test_symbol(self):
|
|
assert grammar_structure_score(Symbol('a')) == 0.0
|
|
|
|
def test_concat(self):
|
|
g = Concat([Symbol('a'), Symbol('b'), Symbol('c')])
|
|
assert grammar_structure_score(g) > 0.0
|
|
|
|
def test_plus(self):
|
|
g = Plus(Symbol('a'))
|
|
assert grammar_structure_score(g) > 0.0
|