grammar.py

"""Tiny grammar engine for acrostic-style constraints.

A faithful Python port of src/grammar.js.

Primitives:
  - Atoms: {"kind": "lit", "allowed": set[str]} (consumes exactly one char from
    `allowed`) or {"kind": "body", "max": int} (consumes 0..max non-newline
    chars).
  - Atom sequences are concatenation-only; with the body/newline structure we
    use, transitions are deterministic, so state packs into one int:
    atom_idx * stride + count.

Builders:
  - compile_acrostic(secret, ...) — list-mode or prose-mode acrostic.
  - compile_literal(text) — exact-text matcher (used by the classifier).
  - union_grammars([g1, g2, ...]) — accept if any branch is alive.

The dead-state sentinel is -1 everywhere, matching the JS original.
"""

from __future__ import annotations

# Spaces in the secret are treated as "word breaks" — they don't pin the line to
# any particular letter, but they still produce a line, and the line must start
# with a punctuation character so the acrostic reads naturally
# ("HI WORLD" -> H... / I... / <punct>... / W... / O... / R... / L... / D...).
PUNCT_FOR_SPACE = set(
    list(".,;:!?-")
    + list("()[]{}")
    + list("~<>")
    + ['"', "'", "`"]
    + list("@#$%&+=/\\|_^")
)


class AtomGrammar:
    """A single concatenation-only atom-sequence grammar (an NFA packed into ints)."""

    def __init__(self, atoms):
        self.atoms = atoms
        max_body_max = 0
        for a in atoms:
            if a["kind"] == "body" and a["max"] > max_body_max:
                max_body_max = a["max"]
        self.stride = max_body_max + 2
        self.PAST_END = len(atoms) * self.stride
        self.state_count = self.PAST_END + 1

        # Precompute accepting states: (a, c) accepts iff atom `a` can be
        # epsilon-skipped at count `c` AND (a+1, 0) is accepting.
        accepting = bytearray(self.state_count)
        accepting[self.PAST_END] = 1
        next_accepting = True
        for a in range(len(atoms) - 1, -1, -1):
            atom = atoms[a]
            if next_accepting:
                mn = 1 if atom["kind"] == "lit" else 0
                mx = 1 if atom["kind"] == "lit" else atom["max"]
                for c in range(mn, mx + 1):
                    accepting[a * self.stride + c] = 1
            next_accepting = accepting[a * self.stride + 0] == 1
        self.accepting = accepting

        self.initial = 0

    def _consume_at(self, a, ch):
        atoms = self.atoms
        while a < len(atoms):
            atom = atoms[a]
            if atom["kind"] == "lit":
                if ch in atom["allowed"]:
                    return self.PAST_END if a + 1 >= len(atoms) else (a + 1) * self.stride
                return -1
            if ch != "\n":
                return a * self.stride + 1
            a += 1
        return -1

    def advance(self, state, s):
        stride = self.stride
        atoms = self.atoms
        cur = state
        for ch in s:
            if cur == self.PAST_END:
                return -1
            a = cur // stride
            c = cur - a * stride
            atom = atoms[a]
            if atom["kind"] == "lit":
                if c < 1 and ch in atom["allowed"]:
                    nxt = self.PAST_END if a + 1 >= len(atoms) else (a + 1) * stride
                else:
                    return -1
            else:
                if c < atom["max"] and ch != "\n":
                    nxt = a * stride + (c + 1)
                else:
                    nxt = self._consume_at(a + 1, ch)
            if nxt == -1:
                return -1
            cur = nxt
        return cur

    def accepts(self, state):
        return state is not None and 0 <= state < self.state_count and self.accepting[state] == 1


def compile_acrostic(secret, list_prefix=" * ", max_line=80, case_insensitive=False, first_line_prefix=True):
    if not secret:
        raise ValueError("secret must be non-empty")
    atoms = []
    for i, letter in enumerate(secret):
        want_prefix = i > 0 or first_line_prefix
        if want_prefix:
            for c in list_prefix:
                atoms.append({"kind": "lit", "allowed": {c}})
        if letter == " ":
            allowed = set(PUNCT_FOR_SPACE)
        elif case_insensitive:
            allowed = {letter.upper(), letter.lower()}
        else:
            allowed = {letter}
        atoms.append({"kind": "lit", "allowed": allowed})
        atoms.append({"kind": "body", "max": max_line})
        if i < len(secret) - 1:
            atoms.append({"kind": "lit", "allowed": {"\n"}})
    return AtomGrammar(atoms)


def compile_literal(text):
    if not text:
        raise ValueError("literal must be non-empty")
    atoms = [{"kind": "lit", "allowed": {c}} for c in text]
    return AtomGrammar(atoms)


class UnionGrammar:
    """Run several grammars in parallel; a token is alive iff at least one branch
    is alive. State is a list of per-branch ints (-1 = dead branch). When every
    branch is dead, advance returns -1 (the single-grammar dead sentinel)."""

    def __init__(self, grammars):
        self.grammars = grammars
        self.initial = [g.initial for g in grammars]

    def advance(self, state, s):
        nxt = [-1] * len(self.grammars)
        any_live = False
        for i, g in enumerate(self.grammars):
            if state[i] == -1:
                nxt[i] = -1
                continue
            r = g.advance(state[i], s)
            nxt[i] = r
            if r != -1:
                any_live = True
        return nxt if any_live else -1

    def accepts(self, state):
        if state == -1:
            return False
        for i, g in enumerate(self.grammars):
            if state[i] != -1 and g.accepts(state[i]):
                return True
        return False


def union_grammars(grammars):
    return UnionGrammar(grammars)