Build grammar expressions directly in a shared finite-state machine

[Ubospica]

Changes

A finite-state machine (FSM) represents valid paths using states and transitions. Previously, the implementation built a small FSM for every sequence element and every choice branch, then copied those FSMs into a larger one. Large grammars therefore repeatedly created, copied, and destroyed short-lived objects.

This PR builds arbitrarily nested grammar expressions directly into a shared target FSM:

• Byte strings, rule references, character classes, repetitions, token edges, and excluded-token edges can all be appended directly.
• Sequence subexpressions reuse start and end states, while choice branches stream into the same target.
• The construction order deliberately preserves the state IDs, edge insertion order, start state, and end state of the previous output.

Correctness

• All 427 Python tests covering grammars, compilers, parsers, matchers, and the new FSM implementation pass.
• All 61 C++ tests pass.
• The main branch and this branch compiled the same set of 17 grammars; their printed FSMs and fully serialized grammars matched byte-for-byte across 177,330 output bytes.
• Added 24 real behavior cases and 17 exact-layout cases covering nested expressions, multibyte text, recursive rule references, token edges, and consistency of the allowed-token set at every step.

Performance

Standalone FSM construction benchmarks:

• A long sequence of 10,000 literal and character-class pairs: 12.656 -> 5.518 ms, a 56.4% reduction and 2.29x speedup.
• A single choice with 8,000 branches: 74.397 -> 66.402 ms, a 10.7% reduction and 1.12x speedup.
• Full construction of a cycle containing 8,000 rules: 46.813 -> 37.766 ms, a 19.3% reduction and 1.24x speedup.

Across four inputs, full grammar optimization time decreased by 10.2% to 20.2%. For a JSON Schema with 50,000 fields, the first full compilation dropped from 4440.5 ms to 4237.0 ms, a 4.6% reduction. This PR does not change allowed-token-set generation during decoding.

[Copilot]

Pull request overview

This PR refactors FSM construction during grammar compilation to stream sequence elements and choice branches directly into a single target FSM, reducing intermediate FSM allocations/copies and improving compile-time performance while preserving resulting FSM structure.

Changes:

• Refactor C++ FSM builder to build sequences/choices by appending directly into a caller-owned target FSM and returning end-state ID sets.
• Update character-class range building helpers to operate directly on FSM (instead of FSMWithStartEnd) to support in-place construction.
• Add a new Python correctness test suite covering sequence construction across element types, recursion, UTF-8, and token/bitmask behavior.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
tests/python/test_fsm_sequence_build.py	Adds extensive correctness coverage for streamed sequence/choice FSM construction and token-bitmask behavior.
cpp/grammar_functor.cc	Refactors sequence/choices FSM building to stream into a shared target FSM; adapts character-range helpers for in-place FSM mutation.

Comments suppressed due to low confidence (1)

tests/python/test_fsm_sequence_build.py:98

• The return type annotation -> (str, str) is not a valid typing return type for mypy (it’s a runtime tuple value). Since the project supports Python 3.8 and uses mypy strict, prefer Tuple[str, str].

def _build_long_sequence_grammar(num_segments: int) -> (str, str):

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.