From e61bc1e339265b4ced19755a0413d939c47a5920 Mon Sep 17 00:00:00 2001 From: Ubospica Date: Mon, 20 Jul 2026 14:15:57 -0400 Subject: [PATCH 1/8] perf: stream FSM sequence concatenation --- cpp/grammar_functor.cc | 64 +++++++++++++++++++++++++++--------------- 1 file changed, 42 insertions(+), 22 deletions(-) diff --git a/cpp/grammar_functor.cc b/cpp/grammar_functor.cc index eae30fe71..da36f0b20 100644 --- a/cpp/grammar_functor.cc +++ b/cpp/grammar_functor.cc @@ -1539,45 +1539,34 @@ std::optional GrammarFSMBuilderImpl::TokenTagDispatch( std::optional GrammarFSMBuilderImpl::Sequence( const GrammarExpr& expr, const Grammar& grammar ) { - std::vector fsm_lists; - - // Build the fsm of sub-expressions. - for (const auto& sequence_id : expr) { - const auto& sequence_expr = grammar->GetGrammarExpr(sequence_id); - switch (sequence_expr.type) { + auto build_element = [](const GrammarExpr& element_expr) -> std::optional { + switch (element_expr.type) { case (ExprType::kByteString): { - fsm_lists.push_back(ByteString(sequence_expr)); - break; + return ByteString(element_expr); } case (ExprType::kRuleRef): { - fsm_lists.push_back(RuleRef(sequence_expr)); - break; + return RuleRef(element_expr); } case (ExprType::kCharacterClass): case (ExprType::kCharacterClassStar): { - fsm_lists.push_back(CharacterClass(sequence_expr)); - break; + return CharacterClass(element_expr); } case (ExprType::kRepeat): { - fsm_lists.push_back(Repeat(sequence_expr)); - break; + return Repeat(element_expr); } case (ExprType::kToken): { - fsm_lists.push_back(Token(sequence_expr)); - break; + return Token(element_expr); } case (ExprType::kExcludeToken): { - fsm_lists.push_back(ExcludeToken(sequence_expr)); - break; + return ExcludeToken(element_expr); } default: { return std::nullopt; } } - } + }; - // Check if the sequence is empty. - if (fsm_lists.empty()) { + if (expr.size() == 0) { FSMWithStartEnd empty_fsm; empty_fsm.AddState(); empty_fsm.SetStartState(0); @@ -1585,7 +1574,38 @@ std::optional GrammarFSMBuilderImpl::Sequence( return empty_fsm; } - return FSMWithStartEnd::Concat(fsm_lists); + if (expr.size() == 1) { + return build_element(grammar->GetGrammarExpr(expr[0])); + } + + // Concatenate the element FSMs as they are built so temporary FSMs can be released + // immediately instead of retaining one allocation-heavy FSM object per sequence element. + FSM result_fsm; + int start = -1; + std::vector previous_ends; + std::vector state_mapping; + for (int32_t sequence_id : expr) { + auto element_fsm = build_element(grammar->GetGrammarExpr(sequence_id)); + if (!element_fsm.has_value()) { + return std::nullopt; + } + result_fsm.AddFSM(element_fsm->GetFsm(), &state_mapping); + int element_start = state_mapping[element_fsm->GetStart()]; + if (start == -1) { + start = element_start; + } else { + for (int32_t previous_end : previous_ends) { + result_fsm.AddEpsilonEdge(previous_end, element_start); + } + } + previous_ends.clear(); + previous_ends.reserve(element_fsm->GetEnds().size()); + for (int32_t end : element_fsm->GetEnds()) { + previous_ends.push_back(state_mapping[end]); + } + } + + return FSMWithStartEnd(result_fsm, start, std::move(previous_ends)); } FSMWithStartEnd GrammarFSMBuilderImpl::RuleRef(const GrammarExpr& expr) { From 2d6b6ff68956bb0d68f07a259c4e979b60a3f0e0 Mon Sep 17 00:00:00 2001 From: Ubospica Date: Mon, 20 Jul 2026 14:19:49 -0400 Subject: [PATCH 2/8] perf: append simple FSM sequence elements directly --- cpp/grammar_functor.cc | 74 +++++++++++++++++++++++++++++++++--------- 1 file changed, 59 insertions(+), 15 deletions(-) diff --git a/cpp/grammar_functor.cc b/cpp/grammar_functor.cc index da36f0b20..e77b88eb9 100644 --- a/cpp/grammar_functor.cc +++ b/cpp/grammar_functor.cc @@ -1582,30 +1582,74 @@ std::optional GrammarFSMBuilderImpl::Sequence( // immediately instead of retaining one allocation-heavy FSM object per sequence element. FSM result_fsm; int start = -1; - std::vector previous_ends; + int previous_end = -1; std::vector state_mapping; for (int32_t sequence_id : expr) { - auto element_fsm = build_element(grammar->GetGrammarExpr(sequence_id)); - if (!element_fsm.has_value()) { - return std::nullopt; + const auto& element_expr = grammar->GetGrammarExpr(sequence_id); + int element_start = result_fsm.NumStates(); + int element_end = -1; + switch (element_expr.type) { + case ExprType::kByteString: { + int current_state = result_fsm.AddState(); + for (int32_t byte : element_expr) { + int next_state = result_fsm.AddState(); + result_fsm.AddEdge( + current_state, next_state, static_cast(byte), static_cast(byte) + ); + current_state = next_state; + } + element_end = current_state; + break; + } + case ExprType::kRuleRef: { + result_fsm.AddState(); + element_end = result_fsm.AddState(); + result_fsm.AddRuleEdge(element_start, element_end, element_expr[0]); + break; + } + case ExprType::kRepeat: { + result_fsm.AddState(); + element_end = result_fsm.AddState(); + result_fsm.AddRepeatEdge( + element_start, element_end, element_expr[0], element_expr[1], element_expr[2] + ); + break; + } + case ExprType::kToken: + case ExprType::kExcludeToken: { + result_fsm.AddState(); + element_end = result_fsm.AddState(); + std::vector token_ids(element_expr.begin(), element_expr.end()); + if (element_expr.type == ExprType::kToken) { + result_fsm.AddTokenEdge(element_start, element_end, token_ids); + } else { + result_fsm.AddExcludeTokenEdge(element_start, element_end, token_ids); + } + break; + } + case ExprType::kCharacterClass: + case ExprType::kCharacterClassStar: { + auto element_fsm = CharacterClass(element_expr); + result_fsm.AddFSM(element_fsm.GetFsm(), &state_mapping); + element_start = state_mapping[element_fsm.GetStart()]; + XGRAMMAR_DCHECK(element_fsm.GetEnds().size() == 1); + element_end = state_mapping[element_fsm.GetEnds()[0]]; + break; + } + default: { + return std::nullopt; + } } - result_fsm.AddFSM(element_fsm->GetFsm(), &state_mapping); - int element_start = state_mapping[element_fsm->GetStart()]; + if (start == -1) { start = element_start; } else { - for (int32_t previous_end : previous_ends) { - result_fsm.AddEpsilonEdge(previous_end, element_start); - } - } - previous_ends.clear(); - previous_ends.reserve(element_fsm->GetEnds().size()); - for (int32_t end : element_fsm->GetEnds()) { - previous_ends.push_back(state_mapping[end]); + result_fsm.AddEpsilonEdge(previous_end, element_start); } + previous_end = element_end; } - return FSMWithStartEnd(result_fsm, start, std::move(previous_ends)); + return FSMWithStartEnd(result_fsm, start, {previous_end}); } FSMWithStartEnd GrammarFSMBuilderImpl::RuleRef(const GrammarExpr& expr) { From c492d4019d87dda1bee8cad68c892caf41808dfe Mon Sep 17 00:00:00 2001 From: Ubospica Date: Thu, 23 Jul 2026 05:03:28 -0400 Subject: [PATCH 3/8] test: cover FSM sequence building correctness in Python Add behavior tests for sequence FSM construction: exhaustive comparisons against reference regexes, UTF-8 byte strings, long mixed sequences, recursive rule references, large repetition ranges, token edges, and bitmask consistency with string acceptance. --- tests/python/test_fsm_sequence_build.py | 244 ++++++++++++++++++++++++ 1 file changed, 244 insertions(+) create mode 100644 tests/python/test_fsm_sequence_build.py diff --git a/tests/python/test_fsm_sequence_build.py b/tests/python/test_fsm_sequence_build.py new file mode 100644 index 000000000..8d5891420 --- /dev/null +++ b/tests/python/test_fsm_sequence_build.py @@ -0,0 +1,244 @@ +"""Correctness tests for building FSMs from grammar sequences. + +The FSM builder constructs sequence FSMs by streaming each element (byte string, rule +reference, character class, repetition, token edge) directly into one target FSM. These +tests verify the resulting matcher behavior on real grammars covering every element type, +including exhaustive comparisons against reference regexes. +""" + +import itertools +import re +import sys +from typing import List + +import pytest + +import xgrammar as xgr +from xgrammar.testing import ( + _get_masked_tokens_from_bitmask, + _get_matcher_from_grammar_and_tokenizer_info, +) + + +def _make_string_matcher(grammar_str: str) -> xgr.GrammarMatcher: + tokenizer_info = xgr.TokenizerInfo([]) + compiler = xgr.GrammarCompiler(tokenizer_info, cache_enabled=False) + compiled = compiler.compile_grammar(grammar_str) + return xgr.GrammarMatcher(compiled, terminate_without_stop_token=True) + + +def _matcher_accepts(matcher: xgr.GrammarMatcher, input_str: str) -> bool: + matcher.reset() + return matcher.accept_string(input_str) and matcher.is_terminated() + + +# --- Exhaustive comparison against reference regexes --- + +# Each entry: (grammar, equivalent regex, alphabet, max enumerated length) +grammar_regex_alphabet_cases = [ + # Byte strings separated by character classes: the sequence FSM path with + # multiple simple elements. + ('root ::= "a" [0-9] "b"', r"a[0-9]b", "a0b", 4), + ('root ::= "ab" [xy] "cd"', r"ab[xy]cd", "abxycd", 6), + # Character class star inside a sequence. + ('root ::= "<" [a-c]* ">"', r"<[a-c]*>", "", 5), + # Negated character class inside a sequence. + ('root ::= "<" [^>] ">"', r"<[^>]>", "", 4), + # Rule reference between byte strings, including an empty alternative. + ('root ::= "a" mid "c"\nmid ::= "b" | "x" | ""', r"a(b|x|)c", "abxc", 4), + # Repetition range in a sequence. + ('root ::= "a" [xy]{1, 3} "z"', r"a[xy]{1,3}z", "axyz", 6), + # Nested alternation followed by sequence elements. + ('root ::= ("a" | "bb") [0-9] "z"', r"(a|bb)[0-9]z", "ab09z", 5), + # Single-element sequences of each type. + ('root ::= "abc"', r"abc", "abc", 4), + ("root ::= [a-z]", r"[a-z]", "az", 3), + ("root ::= [a-b]*", r"[a-b]*", "ab", 4), +] + + +@pytest.mark.parametrize("grammar_str, regex, alphabet, max_len", grammar_regex_alphabet_cases) +def test_sequence_matches_reference_regex( + grammar_str: str, regex: str, alphabet: str, max_len: int +): + """Enumerate every string up to max_len over the alphabet and compare the matcher + against Python's regex engine.""" + matcher = _make_string_matcher(grammar_str) + pattern = re.compile(regex) + checked = 0 + for length in range(max_len + 1): + for candidate_chars in itertools.product(alphabet, repeat=length): + candidate = "".join(candidate_chars) + expected = pattern.fullmatch(candidate) is not None + assert _matcher_accepts(matcher, candidate) == expected, ( + f"Mismatch for input {candidate!r}: grammar={grammar_str!r} regex={regex!r} " + f"expected={expected}" + ) + checked += 1 + assert checked > 1 + + +# --- UTF-8 multi-byte content in sequences --- + + +def test_sequence_with_utf8_byte_strings(): + grammar_str = 'root ::= "你" [好世] "界"' + matcher = _make_string_matcher(grammar_str) + assert _matcher_accepts(matcher, "你好界") + assert _matcher_accepts(matcher, "你世界") + assert not _matcher_accepts(matcher, "你界") + assert not _matcher_accepts(matcher, "你好世界") + assert not _matcher_accepts(matcher, "好界") + assert not _matcher_accepts(matcher, "你好") + + +# --- Long mixed sequences stress the streaming concatenation loop --- + + +def _build_long_sequence_grammar(num_segments: int) -> (str, str): + """A single rule whose body alternates byte strings and character classes.""" + elements: List[str] = [] + valid_parts: List[str] = [] + for index in range(num_segments): + literal = f"s{index:02d}" + elements.append(f'"{literal}"') + elements.append("[0-9]") + valid_parts.append(literal) + valid_parts.append(str(index % 10)) + grammar_str = "root ::= " + " ".join(elements) + return grammar_str, "".join(valid_parts) + + +def test_long_mixed_sequence(): + grammar_str, valid_input = _build_long_sequence_grammar(64) + matcher = _make_string_matcher(grammar_str) + assert _matcher_accepts(matcher, valid_input) + # Truncations must be rejected. + assert not _matcher_accepts(matcher, valid_input[:-1]) + assert not _matcher_accepts(matcher, valid_input[: len(valid_input) // 2]) + # Any extension must be rejected. + assert not _matcher_accepts(matcher, valid_input + "0") + # Mutating one character at several positions must be rejected. + for position in range(0, len(valid_input), 17): + original_char = valid_input[position] + replacement = "x" if original_char != "x" else "y" + mutated = valid_input[:position] + replacement + valid_input[position + 1 :] + assert not _matcher_accepts(matcher, mutated), f"position {position}" + + +# --- Recursive rule references inside sequences --- + + +def test_recursive_rule_ref_in_sequence(): + grammar_str = 'root ::= "(" root ")" | ""' + matcher = _make_string_matcher(grammar_str) + for depth in (0, 1, 2, 8, 32): + assert _matcher_accepts(matcher, "(" * depth + ")" * depth) + assert not _matcher_accepts(matcher, "(") + assert not _matcher_accepts(matcher, "(()") + assert not _matcher_accepts(matcher, "())") + + +# --- Empty sequence --- + + +def test_empty_sequence(): + matcher = _make_string_matcher('root ::= ""') + assert _matcher_accepts(matcher, "") + assert not _matcher_accepts(matcher, "a") + + +# --- Large repetition ranges that stay as repeat edges --- + + +def test_large_repetition_in_sequence(): + grammar_str = 'root ::= "a" [xy]{2, 100} "z"' + matcher = _make_string_matcher(grammar_str) + assert not _matcher_accepts(matcher, "axz") + assert _matcher_accepts(matcher, "a" + "xy" * 1 + "z") + assert _matcher_accepts(matcher, "a" + "x" * 100 + "z") + assert not _matcher_accepts(matcher, "a" + "x" * 101 + "z") + assert not _matcher_accepts(matcher, "a" + "x" * 50) + + +# --- Token and ExcludeToken edges inside sequences --- + +TOKEN_TEST_VOCAB = ["", "", "aa", "bb", "cc", "dd"] +# 0 1 2 3 4 5 +STOP_TOKEN_ID = 1 + + +def _make_token_matcher(grammar_str: str) -> xgr.GrammarMatcher: + tokenizer_info = xgr.TokenizerInfo(TOKEN_TEST_VOCAB) + grammar = xgr.Grammar.from_ebnf(grammar_str) + return _get_matcher_from_grammar_and_tokenizer_info(grammar, tokenizer_info) + + +def test_token_edge_between_byte_strings(): + matcher = _make_token_matcher('root ::= "aa" Token(3, 4) "dd"\n') + assert matcher.accept_token(2) # "aa" + assert not matcher.accept_token(5) # "dd" not in Token(3, 4) + assert matcher.accept_token(3) # "bb" + assert matcher.accept_token(5) # "dd" + assert matcher.accept_token(STOP_TOKEN_ID) + assert matcher.is_terminated() + + +def test_exclude_token_edge_in_sequence(): + matcher = _make_token_matcher('root ::= "aa" ExcludeToken(3) "dd"\n') + assert matcher.accept_token(2) # "aa" + assert not matcher.accept_token(3) # excluded + assert matcher.accept_token(4) # "cc" + assert matcher.accept_token(5) # "dd" + assert matcher.accept_token(STOP_TOKEN_ID) + assert matcher.is_terminated() + + +# --- Token bitmask consistency with string acceptance --- + +mask_consistency_grammars = [ + 'root ::= "ab" [0-9] "cd"', + 'root ::= "a" mid "c"\nmid ::= "b" | "x" | ""', + 'root ::= "<" [a-c]* ">"', + 'root ::= "a" [xy]{1, 3} "z"', +] + + +@pytest.mark.parametrize("grammar_str", mask_consistency_grammars) +def test_bitmask_matches_string_acceptance(grammar_str: str): + """At every step of a generation, the bitmask must allow exactly the vocabulary + pieces the matcher would accept as the next characters.""" + vocab = ["a", "b", "ab", "0", "5", "c", "d", "cd", "x", "y", "z", "<", ">", "bc"] + tokenizer_info = xgr.TokenizerInfo(vocab) + compiler = xgr.GrammarCompiler(tokenizer_info, cache_enabled=False) + compiled = compiler.compile_grammar(grammar_str) + matcher = xgr.GrammarMatcher(compiled, terminate_without_stop_token=True) + bitmask = xgr.allocate_token_bitmask(1, tokenizer_info.vocab_size) + + valid_inputs = { + 'root ::= "ab" [0-9] "cd"': "ab5cd", + 'root ::= "a" mid "c"\nmid ::= "b" | "x" | ""': "axc", + 'root ::= "<" [a-c]* ">"': "", + 'root ::= "a" [xy]{1, 3} "z"': "axyz", + } + remaining = valid_inputs[grammar_str] + + while True: + matcher.fill_next_token_bitmask(bitmask) + rejected = set(_get_masked_tokens_from_bitmask(bitmask, tokenizer_info.vocab_size)) + for token_id, piece in enumerate(vocab): + fork = matcher.fork() + piece_accepted = fork.accept_string(piece) + assert piece_accepted == (token_id not in rejected), ( + f"Bitmask disagrees with accept_string for piece {piece!r} " + f"after consuming {valid_inputs[grammar_str][: -len(remaining) or None]!r}" + ) + if not remaining: + break + assert matcher.accept_string(remaining[0]) + remaining = remaining[1:] + assert matcher.is_terminated() + + +if __name__ == "__main__": + pytest.main(sys.argv) From 80d7c4241b9e8287b60e9456f9bd7fa5bed1c153 Mon Sep 17 00:00:00 2001 From: Ubospica Date: Tue, 28 Jul 2026 05:17:53 -0400 Subject: [PATCH 4/8] refactor: extract BuildElement as a standalone function --- cpp/grammar_functor.cc | 60 ++++++++++++++++++++++-------------------- 1 file changed, 32 insertions(+), 28 deletions(-) diff --git a/cpp/grammar_functor.cc b/cpp/grammar_functor.cc index e77b88eb9..4235db87e 100644 --- a/cpp/grammar_functor.cc +++ b/cpp/grammar_functor.cc @@ -1163,6 +1163,9 @@ class GrammarFSMBuilderImpl { static FSMWithStartEnd ExcludeToken(const GrammarExpr& expr); static std::optional TokenTagDispatch(const Grammar::Impl::TokenTagDispatch& ttd ); + /*! \brief Build the FSM of a single sequence element. Returns std::nullopt for unsupported + * element types. */ + static std::optional BuildElement(const GrammarExpr& element_expr); static std::optional Sequence(const GrammarExpr& expr, const Grammar& grammar); static std::optional Choices(const GrammarExpr& expr, const Grammar& grammar); static std::optional TagDispatch(const Grammar::Impl::TagDispatch& tag_dispatch); @@ -1536,36 +1539,37 @@ std::optional GrammarFSMBuilderImpl::TokenTagDispatch( return FSMWithStartEnd(fsm, start, ends); } -std::optional GrammarFSMBuilderImpl::Sequence( - const GrammarExpr& expr, const Grammar& grammar +std::optional GrammarFSMBuilderImpl::BuildElement(const GrammarExpr& element_expr ) { - auto build_element = [](const GrammarExpr& element_expr) -> std::optional { - switch (element_expr.type) { - case (ExprType::kByteString): { - return ByteString(element_expr); - } - case (ExprType::kRuleRef): { - return RuleRef(element_expr); - } - case (ExprType::kCharacterClass): - case (ExprType::kCharacterClassStar): { - return CharacterClass(element_expr); - } - case (ExprType::kRepeat): { - return Repeat(element_expr); - } - case (ExprType::kToken): { - return Token(element_expr); - } - case (ExprType::kExcludeToken): { - return ExcludeToken(element_expr); - } - default: { - return std::nullopt; - } + switch (element_expr.type) { + case (ExprType::kByteString): { + return ByteString(element_expr); } - }; + case (ExprType::kRuleRef): { + return RuleRef(element_expr); + } + case (ExprType::kCharacterClass): + case (ExprType::kCharacterClassStar): { + return CharacterClass(element_expr); + } + case (ExprType::kRepeat): { + return Repeat(element_expr); + } + case (ExprType::kToken): { + return Token(element_expr); + } + case (ExprType::kExcludeToken): { + return ExcludeToken(element_expr); + } + default: { + return std::nullopt; + } + } +} +std::optional GrammarFSMBuilderImpl::Sequence( + const GrammarExpr& expr, const Grammar& grammar +) { if (expr.size() == 0) { FSMWithStartEnd empty_fsm; empty_fsm.AddState(); @@ -1575,7 +1579,7 @@ std::optional GrammarFSMBuilderImpl::Sequence( } if (expr.size() == 1) { - return build_element(grammar->GetGrammarExpr(expr[0])); + return BuildElement(grammar->GetGrammarExpr(expr[0])); } // Concatenate the element FSMs as they are built so temporary FSMs can be released From e61314602ef8dfff6c0b569e7a43a902a77448c4 Mon Sep 17 00:00:00 2001 From: Ubospica Date: Tue, 28 Jul 2026 06:18:38 -0400 Subject: [PATCH 5/8] perf: build grammar fragments in a shared FSM Stream sequence and choice fragments into one target FSM so element builders can reuse caller-provided start states without temporary FSM copies. --- cpp/grammar_functor.cc | 483 ++++++++++++++++++++--------------------- 1 file changed, 240 insertions(+), 243 deletions(-) diff --git a/cpp/grammar_functor.cc b/cpp/grammar_functor.cc index 4235db87e..cba1ddff0 100644 --- a/cpp/grammar_functor.cc +++ b/cpp/grammar_functor.cc @@ -1089,6 +1089,8 @@ class GrammarFSMBuilderImpl { const static uint32_t kMin4BytesUnicode = 0xF0808080; const static uint32_t kMax4BytesUnicode = 0xF7BFBFBF; + explicit GrammarFSMBuilderImpl(FSM* target_fsm = nullptr) : target_fsm_(target_fsm) {} + void Apply(Grammar* grammar) { FSM complete_fsm; std::vector> per_rule_fsms((*grammar)->NumRules()); @@ -1158,33 +1160,50 @@ class GrammarFSMBuilderImpl { static FSMWithStartEnd RuleRef(const GrammarExpr& expr); static FSMWithStartEnd CharacterClass(const GrammarExpr& expr); static FSMWithStartEnd ByteString(const GrammarExpr& expr); - static FSMWithStartEnd Repeat(const GrammarExpr& expr); static FSMWithStartEnd Token(const GrammarExpr& expr); static FSMWithStartEnd ExcludeToken(const GrammarExpr& expr); static std::optional TokenTagDispatch(const Grammar::Impl::TokenTagDispatch& ttd ); - /*! \brief Build the FSM of a single sequence element. Returns std::nullopt for unsupported - * element types. */ - static std::optional BuildElement(const GrammarExpr& element_expr); static std::optional Sequence(const GrammarExpr& expr, const Grammar& grammar); static std::optional Choices(const GrammarExpr& expr, const Grammar& grammar); static std::optional TagDispatch(const Grammar::Impl::TagDispatch& tag_dispatch); static Result Regex(const std::string& regex, bool json_string = false); - static void AddCharacterRange(FSMWithStartEnd& fsm, int from, int to, uint32_t min, uint32_t max); /* Building tool functions.*/ static std::optional BuildTagDispatch( const std::vector>& string_trigger_rules, bool loop_after_dispatch, const std::vector& excluded_strings ); - static FSMWithStartEnd BuildNegativeCharacterClass(const GrammarExpr& expr); + + private: + static FSMWithStartEnd BuildSingleElement(const GrammarExpr& expr); + bool BuildElement( + const GrammarExpr& element_expr, int start_state, std::vector* end_states + ); + bool BuildSequence( + const GrammarExpr& expr, + const Grammar& grammar, + int start_state, + std::vector* end_states + ); + bool BuildChoices( + const GrammarExpr& expr, + const Grammar& grammar, + int start_state, + std::vector* end_states + ); + void BuildCharacterClass( + const GrammarExpr& expr, int start_state, std::vector* end_states + ); + void BuildNegativeCharacterClass(const GrammarExpr& expr, int start_state, int end_state); + void AddCharacterRange(int from, int to, uint32_t min, uint32_t max); + + FSM* target_fsm_; }; // This function will add a range [min, max] of characters to the FSM, and the length // of the characters are the same. -void AddSameLengthCharacterRange( - FSMWithStartEnd& fsm, int from, int to, uint32_t min, uint32_t max -) { +void AddSameLengthCharacterRange(FSM& fsm, int from, int to, uint32_t min, uint32_t max) { uint8_t byte_min[4] = { static_cast(min & 0xFF), static_cast(min >> 8), @@ -1200,7 +1219,7 @@ void AddSameLengthCharacterRange( // ASCII. if (byte_max[1] == 0) { - fsm.GetFsm().AddEdge(from, to, byte_min[0], byte_max[0]); + fsm.AddEdge(from, to, byte_min[0], byte_max[0]); return; } @@ -1208,7 +1227,7 @@ void AddSameLengthCharacterRange( // 4-byte unicode. if (byte_max[3] == byte_min[3]) { int tmp_state = fsm.AddState(); - fsm.GetFsm().AddEdge(from, tmp_state, byte_min[3], byte_max[3]); + fsm.AddEdge(from, tmp_state, byte_min[3], byte_max[3]); min = (min & 0x00FFFFFF); max = (max & 0x00FFFFFF); AddSameLengthCharacterRange(fsm, tmp_state, to, min, max); @@ -1216,14 +1235,14 @@ void AddSameLengthCharacterRange( } if ((min & 0x00FFFFFF) != 0x808080) { int tmp_state_min = fsm.AddState(); - fsm.GetFsm().AddEdge(from, tmp_state_min, byte_min[3], byte_min[3]); + fsm.AddEdge(from, tmp_state_min, byte_min[3], byte_min[3]); AddSameLengthCharacterRange(fsm, tmp_state_min, to, (min & 0x00FFFFFF), 0x00BFBFBF); } else { byte_min[3]--; } if ((max & 0x00FFFFFF) != 0xBFBFBF) { int tmp_state_max = fsm.AddState(); - fsm.GetFsm().AddEdge(from, tmp_state_max, byte_max[3], byte_max[3]); + fsm.AddEdge(from, tmp_state_max, byte_max[3], byte_max[3]); AddSameLengthCharacterRange(fsm, tmp_state_max, to, 0x00808080, (max & 0x00FFFFFF)); } else { byte_max[3]++; @@ -1231,15 +1250,15 @@ void AddSameLengthCharacterRange( if (byte_max[3] - byte_min[3] > 1) { int tmp_state_mid = fsm.AddState(); // First byte. - fsm.GetFsm().AddEdge(from, tmp_state_mid, byte_min[3] + 1, byte_max[3] - 1); + fsm.AddEdge(from, tmp_state_mid, byte_min[3] + 1, byte_max[3] - 1); int tmp_state_mid2 = fsm.AddState(); // Second byte. - fsm.GetFsm().AddEdge(tmp_state_mid, tmp_state_mid2, 0x80, 0xBF); + fsm.AddEdge(tmp_state_mid, tmp_state_mid2, 0x80, 0xBF); int tmp_state_mid3 = fsm.AddState(); // Third byte. - fsm.GetFsm().AddEdge(tmp_state_mid2, tmp_state_mid3, 0x80, 0xBF); + fsm.AddEdge(tmp_state_mid2, tmp_state_mid3, 0x80, 0xBF); // Last byte. - fsm.GetFsm().AddEdge(tmp_state_mid3, to, 0x80, 0xBF); + fsm.AddEdge(tmp_state_mid3, to, 0x80, 0xBF); } return; } @@ -1247,7 +1266,7 @@ void AddSameLengthCharacterRange( // 3 byte unicode. if (byte_max[2] == byte_min[2]) { int tmp_state = fsm.AddState(); - fsm.GetFsm().AddEdge(from, tmp_state, byte_min[2], byte_max[2]); + fsm.AddEdge(from, tmp_state, byte_min[2], byte_max[2]); min = (min & 0x00FFFF); max = (max & 0x00FFFF); AddSameLengthCharacterRange(fsm, tmp_state, to, min, max); @@ -1255,14 +1274,14 @@ void AddSameLengthCharacterRange( } if ((min & 0x00FFFF) != 0x8080) { int tmp_state_min = fsm.AddState(); - fsm.GetFsm().AddEdge(from, tmp_state_min, byte_min[2], byte_min[2]); + fsm.AddEdge(from, tmp_state_min, byte_min[2], byte_min[2]); AddSameLengthCharacterRange(fsm, tmp_state_min, to, (min & 0x00FFFF), 0x00BFBF); } else { byte_min[2]--; } if ((max & 0x00FFFF) != 0xBFBF) { int tmp_state_max = fsm.AddState(); - fsm.GetFsm().AddEdge(from, tmp_state_max, byte_max[2], byte_max[2]); + fsm.AddEdge(from, tmp_state_max, byte_max[2], byte_max[2]); AddSameLengthCharacterRange(fsm, tmp_state_max, to, 0x0080, (max & 0x00FFFF)); } else { byte_max[2]++; @@ -1270,12 +1289,12 @@ void AddSameLengthCharacterRange( if (byte_max[2] - byte_min[2] > 1) { int tmp_state_mid = fsm.AddState(); // First byte. - fsm.GetFsm().AddEdge(from, tmp_state_mid, byte_min[2] + 1, byte_max[2] - 1); + fsm.AddEdge(from, tmp_state_mid, byte_min[2] + 1, byte_max[2] - 1); int tmp_state_mid2 = fsm.AddState(); // Second byte. - fsm.GetFsm().AddEdge(tmp_state_mid, tmp_state_mid2, 0x80, 0xBF); + fsm.AddEdge(tmp_state_mid, tmp_state_mid2, 0x80, 0xBF); // Last byte. - fsm.GetFsm().AddEdge(tmp_state_mid2, to, 0x80, 0xBF); + fsm.AddEdge(tmp_state_mid2, to, 0x80, 0xBF); } return; } @@ -1283,7 +1302,7 @@ void AddSameLengthCharacterRange( // 2 byte unicode. if (byte_max[1] == byte_min[1]) { int tmp_state = fsm.AddState(); - fsm.GetFsm().AddEdge(from, tmp_state, byte_min[1], byte_max[1]); + fsm.AddEdge(from, tmp_state, byte_min[1], byte_max[1]); min = (min & 0x00FF); max = (max & 0x00FF); AddSameLengthCharacterRange(fsm, tmp_state, to, min, max); @@ -1291,14 +1310,14 @@ void AddSameLengthCharacterRange( } if ((min & 0x00FF) != 0x80) { int tmp_state_min = fsm.AddState(); - fsm.GetFsm().AddEdge(from, tmp_state_min, byte_min[1], byte_min[1]); + fsm.AddEdge(from, tmp_state_min, byte_min[1], byte_min[1]); AddSameLengthCharacterRange(fsm, tmp_state_min, to, (min & 0x00FF), 0x00BF); } else { byte_min[1]--; } if ((max & 0x00FF) != 0xBF) { int tmp_state_max = fsm.AddState(); - fsm.GetFsm().AddEdge(from, tmp_state_max, byte_max[1], byte_max[1]); + fsm.AddEdge(from, tmp_state_max, byte_max[1], byte_max[1]); AddSameLengthCharacterRange(fsm, tmp_state_max, to, 0x0080, (max & 0x00FF)); } else { byte_max[1]++; @@ -1306,16 +1325,15 @@ void AddSameLengthCharacterRange( if (byte_max[1] - byte_min[1] > 1) { int tmp_state_mid = fsm.AddState(); // First byte. - fsm.GetFsm().AddEdge(from, tmp_state_mid, byte_min[1] + 1, byte_max[1] - 1); - fsm.GetFsm().AddEdge(tmp_state_mid, to, 0x80, 0xBF); + fsm.AddEdge(from, tmp_state_mid, byte_min[1] + 1, byte_max[1] - 1); + fsm.AddEdge(tmp_state_mid, to, 0x80, 0xBF); } return; } // This function will add a range [min, max] of unicode characters to the FSM. -void GrammarFSMBuilderImpl::AddCharacterRange( - FSMWithStartEnd& fsm, int from, int to, uint32_t min, uint32_t max -) { +void GrammarFSMBuilderImpl::AddCharacterRange(int from, int to, uint32_t min, uint32_t max) { + XGRAMMAR_DCHECK(target_fsm_ != nullptr); XGRAMMAR_CHECK(min <= max) << "Invalid character range: min (" << min << ") > max (" << max << ")"; // Ensure max and min are valid unicode value. @@ -1349,55 +1367,58 @@ void GrammarFSMBuilderImpl::AddCharacterRange( // Step2. Divide the range into several ranges, which contain characters with different lengths. if (max <= kMax1ByteUnicode) { - AddSameLengthCharacterRange(fsm, from, to, min, max); + AddSameLengthCharacterRange(*target_fsm_, from, to, min, max); return; } if (max <= kMax2BytesUnicode) { if (min >= kMin2BytesUnicode) { - AddSameLengthCharacterRange(fsm, from, to, min, max); + AddSameLengthCharacterRange(*target_fsm_, from, to, min, max); } else { - AddSameLengthCharacterRange(fsm, from, to, min, kMax1ByteUnicode); - AddSameLengthCharacterRange(fsm, from, to, kMin2BytesUnicode, max); + AddSameLengthCharacterRange(*target_fsm_, from, to, min, kMax1ByteUnicode); + AddSameLengthCharacterRange(*target_fsm_, from, to, kMin2BytesUnicode, max); } return; } if (max <= kMax3BytesUnicode) { if (min >= kMin3BytesUnicode) { - AddSameLengthCharacterRange(fsm, from, to, min, max); + AddSameLengthCharacterRange(*target_fsm_, from, to, min, max); } else if (min >= kMin2BytesUnicode) { - AddSameLengthCharacterRange(fsm, from, to, min, kMax2BytesUnicode); - AddSameLengthCharacterRange(fsm, from, to, kMin3BytesUnicode, max); + AddSameLengthCharacterRange(*target_fsm_, from, to, min, kMax2BytesUnicode); + AddSameLengthCharacterRange(*target_fsm_, from, to, kMin3BytesUnicode, max); } else { - AddSameLengthCharacterRange(fsm, from, to, min, kMax1ByteUnicode); - AddSameLengthCharacterRange(fsm, from, to, kMin2BytesUnicode, kMax2BytesUnicode); - AddSameLengthCharacterRange(fsm, from, to, kMin3BytesUnicode, max); + AddSameLengthCharacterRange(*target_fsm_, from, to, min, kMax1ByteUnicode); + AddSameLengthCharacterRange(*target_fsm_, from, to, kMin2BytesUnicode, kMax2BytesUnicode); + AddSameLengthCharacterRange(*target_fsm_, from, to, kMin3BytesUnicode, max); } return; } XGRAMMAR_CHECK(max <= kMax4BytesUnicode); if (min >= kMin4BytesUnicode) { - AddSameLengthCharacterRange(fsm, from, to, min, max); + AddSameLengthCharacterRange(*target_fsm_, from, to, min, max); } else if (min >= kMin3BytesUnicode) { - AddSameLengthCharacterRange(fsm, from, to, min, kMax3BytesUnicode); - AddSameLengthCharacterRange(fsm, from, to, kMin4BytesUnicode, max); + AddSameLengthCharacterRange(*target_fsm_, from, to, min, kMax3BytesUnicode); + AddSameLengthCharacterRange(*target_fsm_, from, to, kMin4BytesUnicode, max); } else if (min >= kMin2BytesUnicode) { - AddSameLengthCharacterRange(fsm, from, to, min, kMax2BytesUnicode); - AddSameLengthCharacterRange(fsm, from, to, kMin3BytesUnicode, kMax3BytesUnicode); - AddSameLengthCharacterRange(fsm, from, to, kMin4BytesUnicode, max); + AddSameLengthCharacterRange(*target_fsm_, from, to, min, kMax2BytesUnicode); + AddSameLengthCharacterRange(*target_fsm_, from, to, kMin3BytesUnicode, kMax3BytesUnicode); + AddSameLengthCharacterRange(*target_fsm_, from, to, kMin4BytesUnicode, max); } else { - AddSameLengthCharacterRange(fsm, from, to, min, kMax1ByteUnicode); - AddSameLengthCharacterRange(fsm, from, to, kMin2BytesUnicode, kMax2BytesUnicode); - AddSameLengthCharacterRange(fsm, from, to, kMin3BytesUnicode, kMax3BytesUnicode); - AddSameLengthCharacterRange(fsm, from, to, kMin4BytesUnicode, max); + AddSameLengthCharacterRange(*target_fsm_, from, to, min, kMax1ByteUnicode); + AddSameLengthCharacterRange(*target_fsm_, from, to, kMin2BytesUnicode, kMax2BytesUnicode); + AddSameLengthCharacterRange(*target_fsm_, from, to, kMin3BytesUnicode, kMax3BytesUnicode); + AddSameLengthCharacterRange(*target_fsm_, from, to, kMin4BytesUnicode, max); } return; } -FSMWithStartEnd GrammarFSMBuilderImpl::BuildNegativeCharacterClass(const GrammarExpr& expr) { +void GrammarFSMBuilderImpl::BuildNegativeCharacterClass( + const GrammarExpr& expr, int start_state, int end_state +) { XGRAMMAR_DCHECK( expr.type == ExprType::kCharacterClass || expr.type == ExprType::kCharacterClassStar ); XGRAMMAR_DCHECK(expr[0]); // Negative character class should be true. + XGRAMMAR_DCHECK(target_fsm_ != nullptr); std::bitset<128> char_set; for (int i = 1; i < static_cast(expr.size()); i += 2) { uint8_t byte_min = static_cast(expr[i]); @@ -1412,18 +1433,6 @@ FSMWithStartEnd GrammarFSMBuilderImpl::BuildNegativeCharacterClass(const Grammar } } - // Construct the basic FSM. - FSMWithStartEnd result_fsm; - int start_state = result_fsm.AddState(); - bool is_star = expr.type == ExprType::kCharacterClassStar; - result_fsm.SetStartState(start_state); - int end_state = -1; - if (is_star) { - end_state = start_state; - } else { - end_state = result_fsm.AddState(); - } - result_fsm.AddEndState(end_state); int left_bound = -1; for (int i = 0; i < 128; ++i) { if (!char_set[i]) { @@ -1432,7 +1441,7 @@ FSMWithStartEnd GrammarFSMBuilderImpl::BuildNegativeCharacterClass(const Grammar while (right_bound < 128 && !char_set[right_bound]) { right_bound++; } - result_fsm.GetFsm().AddEdge( + target_fsm_->AddEdge( start_state, end_state, static_cast(left_bound), @@ -1441,65 +1450,60 @@ FSMWithStartEnd GrammarFSMBuilderImpl::BuildNegativeCharacterClass(const Grammar i = right_bound; } } - AddCharacterRange(result_fsm, start_state, end_state, kMin2BytesUnicode, kMax4BytesUnicode); - return result_fsm; + AddCharacterRange(start_state, end_state, kMin2BytesUnicode, kMax4BytesUnicode); } -FSMWithStartEnd GrammarFSMBuilderImpl::CharacterClass(const GrammarExpr& expr) { +void GrammarFSMBuilderImpl::BuildCharacterClass( + const GrammarExpr& expr, int start_state, std::vector* end_states +) { + XGRAMMAR_DCHECK(target_fsm_ != nullptr); + end_states->clear(); + int end_state = + expr.type == ExprType::kCharacterClassStar ? start_state : target_fsm_->AddState(); bool is_negative = expr[0]; - FSMWithStartEnd result_fsm; if (is_negative) { - result_fsm = BuildNegativeCharacterClass(expr); - return result_fsm; - } - int start_state = result_fsm.AddState(); - result_fsm.SetStartState(start_state); - bool is_star = expr.type == ExprType::kCharacterClassStar; - int end_state = -1; - if (is_star) { - end_state = start_state; + BuildNegativeCharacterClass(expr, start_state, end_state); } else { - end_state = result_fsm.AddState(); + for (int i = 1; i < static_cast(expr.size()); i += 2) { + uint32_t codepoint_min = static_cast(expr[i]); + uint32_t codepoint_max = static_cast(expr[i + 1]); + // Convert Unicode codepoints to packed UTF-8 format for AddCharacterRange + uint32_t packed_min = CodepointToPackedUTF8(codepoint_min); + uint32_t packed_max = CodepointToPackedUTF8(codepoint_max); + AddCharacterRange(start_state, end_state, packed_min, packed_max); + } } - result_fsm.AddEndState(end_state); - for (int i = 1; i < static_cast(expr.size()); i += 2) { - uint32_t codepoint_min = static_cast(expr[i]); - uint32_t codepoint_max = static_cast(expr[i + 1]); - // Convert Unicode codepoints to packed UTF-8 format for AddCharacterRange - uint32_t packed_min = CodepointToPackedUTF8(codepoint_min); - uint32_t packed_max = CodepointToPackedUTF8(codepoint_max); - AddCharacterRange(result_fsm, start_state, end_state, packed_min, packed_max); - } - return result_fsm; + end_states->push_back(end_state); } -FSMWithStartEnd GrammarFSMBuilderImpl::Repeat(const GrammarExpr& expr) { - int32_t rule_id = expr[0]; - int32_t lower = expr[1]; - int32_t upper = expr[2]; - FSMWithStartEnd repeat_fsm; - repeat_fsm.AddState(); - repeat_fsm.AddState(); - repeat_fsm.SetStartState(0); - repeat_fsm.AddEndState(1); - repeat_fsm.GetFsm().AddRepeatEdge(0, 1, rule_id, lower, upper); - return repeat_fsm; +FSMWithStartEnd GrammarFSMBuilderImpl::BuildSingleElement(const GrammarExpr& expr) { + FSM result_fsm; + int start_state = result_fsm.AddState(); + std::vector end_states; + GrammarFSMBuilderImpl builder(&result_fsm); + XGRAMMAR_CHECK(builder.BuildElement(expr, start_state, &end_states)) + << "Unsupported grammar expression type: " << static_cast(expr.type); + return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); +} + +FSMWithStartEnd GrammarFSMBuilderImpl::RuleRef(const GrammarExpr& expr) { + return BuildSingleElement(expr); +} + +FSMWithStartEnd GrammarFSMBuilderImpl::CharacterClass(const GrammarExpr& expr) { + return BuildSingleElement(expr); +} + +FSMWithStartEnd GrammarFSMBuilderImpl::ByteString(const GrammarExpr& expr) { + return BuildSingleElement(expr); } FSMWithStartEnd GrammarFSMBuilderImpl::Token(const GrammarExpr& expr) { - XGRAMMAR_DCHECK(expr.type == ExprType::kToken); - std::vector token_ids(expr.begin(), expr.end()); - FSM fsm(2); - fsm.AddTokenEdge(0, 1, token_ids); - return FSMWithStartEnd(fsm, 0, {1}); + return BuildSingleElement(expr); } FSMWithStartEnd GrammarFSMBuilderImpl::ExcludeToken(const GrammarExpr& expr) { - XGRAMMAR_DCHECK(expr.type == ExprType::kExcludeToken); - std::vector token_ids(expr.begin(), expr.end()); - FSM fsm(2); - fsm.AddExcludeTokenEdge(0, 1, token_ids); - return FSMWithStartEnd(fsm, 0, {1}); + return BuildSingleElement(expr); } std::optional GrammarFSMBuilderImpl::TokenTagDispatch( @@ -1539,186 +1543,179 @@ std::optional GrammarFSMBuilderImpl::TokenTagDispatch( return FSMWithStartEnd(fsm, start, ends); } -std::optional GrammarFSMBuilderImpl::BuildElement(const GrammarExpr& element_expr +bool GrammarFSMBuilderImpl::BuildElement( + const GrammarExpr& element_expr, int start_state, std::vector* end_states ) { + XGRAMMAR_DCHECK(target_fsm_ != nullptr); + XGRAMMAR_DCHECK(start_state >= 0 && start_state < target_fsm_->NumStates()); + end_states->clear(); switch (element_expr.type) { case (ExprType::kByteString): { - return ByteString(element_expr); + int current_state = start_state; + for (int32_t byte : element_expr) { + int next_state = target_fsm_->AddState(); + target_fsm_->AddEdge( + current_state, next_state, static_cast(byte), static_cast(byte) + ); + current_state = next_state; + } + end_states->push_back(current_state); + return true; } case (ExprType::kRuleRef): { - return RuleRef(element_expr); + int end_state = target_fsm_->AddState(); + target_fsm_->AddRuleEdge(start_state, end_state, element_expr[0]); + end_states->push_back(end_state); + return true; } case (ExprType::kCharacterClass): case (ExprType::kCharacterClassStar): { - return CharacterClass(element_expr); + BuildCharacterClass(element_expr, start_state, end_states); + return true; } case (ExprType::kRepeat): { - return Repeat(element_expr); - } - case (ExprType::kToken): { - return Token(element_expr); + int end_state = target_fsm_->AddState(); + target_fsm_->AddRepeatEdge( + start_state, end_state, element_expr[0], element_expr[1], element_expr[2] + ); + end_states->push_back(end_state); + return true; } + case (ExprType::kToken): case (ExprType::kExcludeToken): { - return ExcludeToken(element_expr); + int end_state = target_fsm_->AddState(); + std::vector token_ids(element_expr.begin(), element_expr.end()); + if (element_expr.type == ExprType::kToken) { + target_fsm_->AddTokenEdge(start_state, end_state, token_ids); + } else { + target_fsm_->AddExcludeTokenEdge(start_state, end_state, token_ids); + } + end_states->push_back(end_state); + return true; } default: { - return std::nullopt; + return false; } } } -std::optional GrammarFSMBuilderImpl::Sequence( - const GrammarExpr& expr, const Grammar& grammar +bool GrammarFSMBuilderImpl::BuildSequence( + const GrammarExpr& expr, + const Grammar& grammar, + int start_state, + std::vector* end_states ) { + XGRAMMAR_DCHECK(target_fsm_ != nullptr); + XGRAMMAR_DCHECK(start_state >= 0 && start_state < target_fsm_->NumStates()); + end_states->clear(); if (expr.size() == 0) { - FSMWithStartEnd empty_fsm; - empty_fsm.AddState(); - empty_fsm.SetStartState(0); - empty_fsm.AddEndState(0); - return empty_fsm; + end_states->push_back(start_state); + return true; } - if (expr.size() == 1) { - return BuildElement(grammar->GetGrammarExpr(expr[0])); + if (!BuildElement(grammar->GetGrammarExpr(expr[0]), start_state, end_states)) { + return false; } - // Concatenate the element FSMs as they are built so temporary FSMs can be released - // immediately instead of retaining one allocation-heavy FSM object per sequence element. - FSM result_fsm; - int start = -1; - int previous_end = -1; - std::vector state_mapping; - for (int32_t sequence_id : expr) { - const auto& element_expr = grammar->GetGrammarExpr(sequence_id); - int element_start = result_fsm.NumStates(); - int element_end = -1; - switch (element_expr.type) { - case ExprType::kByteString: { - int current_state = result_fsm.AddState(); - for (int32_t byte : element_expr) { - int next_state = result_fsm.AddState(); - result_fsm.AddEdge( - current_state, next_state, static_cast(byte), static_cast(byte) - ); - current_state = next_state; - } - element_end = current_state; - break; - } - case ExprType::kRuleRef: { - result_fsm.AddState(); - element_end = result_fsm.AddState(); - result_fsm.AddRuleEdge(element_start, element_end, element_expr[0]); - break; - } - case ExprType::kRepeat: { - result_fsm.AddState(); - element_end = result_fsm.AddState(); - result_fsm.AddRepeatEdge( - element_start, element_end, element_expr[0], element_expr[1], element_expr[2] - ); - break; - } - case ExprType::kToken: - case ExprType::kExcludeToken: { - result_fsm.AddState(); - element_end = result_fsm.AddState(); - std::vector token_ids(element_expr.begin(), element_expr.end()); - if (element_expr.type == ExprType::kToken) { - result_fsm.AddTokenEdge(element_start, element_end, token_ids); - } else { - result_fsm.AddExcludeTokenEdge(element_start, element_end, token_ids); - } - break; - } - case ExprType::kCharacterClass: - case ExprType::kCharacterClassStar: { - auto element_fsm = CharacterClass(element_expr); - result_fsm.AddFSM(element_fsm.GetFsm(), &state_mapping); - element_start = state_mapping[element_fsm.GetStart()]; - XGRAMMAR_DCHECK(element_fsm.GetEnds().size() == 1); - element_end = state_mapping[element_fsm.GetEnds()[0]]; - break; - } - default: { - return std::nullopt; - } + std::vector next_end_states; + for (int index = 1; index < static_cast(expr.size()); ++index) { + int element_start = target_fsm_->AddState(); + if (!BuildElement(grammar->GetGrammarExpr(expr[index]), element_start, &next_end_states)) { + return false; } - - if (start == -1) { - start = element_start; - } else { - result_fsm.AddEpsilonEdge(previous_end, element_start); + for (int32_t previous_end_state : *end_states) { + target_fsm_->AddEpsilonEdge(previous_end_state, element_start); } - previous_end = element_end; + end_states->swap(next_end_states); } - return FSMWithStartEnd(result_fsm, start, {previous_end}); -} - -FSMWithStartEnd GrammarFSMBuilderImpl::RuleRef(const GrammarExpr& expr) { - FSMWithStartEnd result_fsm; - result_fsm.AddState(); - result_fsm.AddState(); - result_fsm.SetStartState(0); - result_fsm.AddEndState(1); - result_fsm.GetFsm().AddRuleEdge(0, 1, expr[0]); - return result_fsm; + return true; } -FSMWithStartEnd GrammarFSMBuilderImpl::ByteString(const GrammarExpr& expr) { - XGRAMMAR_DCHECK(expr.type == ExprType::kByteString); - FSMWithStartEnd result_fsm; - int current_state = result_fsm.AddState(); - result_fsm.SetStartState(current_state); - for (const auto& byte : expr) { - int next_state = result_fsm.AddState(); - result_fsm.GetFsm().AddEdge( - current_state, next_state, static_cast(byte), static_cast(byte) - ); - current_state = next_state; +std::optional GrammarFSMBuilderImpl::Sequence( + const GrammarExpr& expr, const Grammar& grammar +) { + FSM result_fsm; + int start_state = result_fsm.AddState(); + std::vector end_states; + GrammarFSMBuilderImpl builder(&result_fsm); + if (!builder.BuildSequence(expr, grammar, start_state, &end_states)) { + return std::nullopt; } - result_fsm.AddEndState(current_state); - return result_fsm; + return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); } -std::optional GrammarFSMBuilderImpl::Choices( - const GrammarExpr& expr, const Grammar& grammar +bool GrammarFSMBuilderImpl::BuildChoices( + const GrammarExpr& expr, + const Grammar& grammar, + int start_state, + std::vector* end_states ) { XGRAMMAR_DCHECK(expr.type == ExprType::kChoices); - std::vector fsm_list; + XGRAMMAR_DCHECK(target_fsm_ != nullptr); + XGRAMMAR_DCHECK(start_state >= 0 && start_state < target_fsm_->NumStates()); + end_states->clear(); + + int sequence_count = 0; bool nullable = false; - for (const auto& choice_id : expr) { + for (int32_t choice_id : expr) { const auto& choice_expr = grammar->GetGrammarExpr(choice_id); if (choice_expr.type == ExprType::kEmptyStr) { nullable = true; - continue; + } else { + XGRAMMAR_DCHECK(choice_expr.type == ExprType::kSequence); + ++sequence_count; } - XGRAMMAR_DCHECK(choice_expr.type == ExprType::kSequence); - auto fsm_result = Sequence(choice_expr, grammar); - if (!fsm_result.has_value()) { - return std::nullopt; + } + + if (sequence_count == 0) { + end_states->push_back(start_state); + return true; + } + + if (sequence_count == 1 && !nullable) { + for (int32_t choice_id : expr) { + const auto& choice_expr = grammar->GetGrammarExpr(choice_id); + if (choice_expr.type != ExprType::kEmptyStr) { + return BuildSequence(choice_expr, grammar, start_state, end_states); + } } - fsm_list.push_back(std::move(fsm_result.value())); + XGRAMMAR_UNREACHABLE(); } - if (fsm_list.empty()) { - // It's an empty rule. - FSMWithStartEnd empty_fsm; - empty_fsm.AddState(); - empty_fsm.SetStartState(0); - empty_fsm.AddEndState(0); - return empty_fsm; + std::vector branch_end_states; + for (int32_t choice_id : expr) { + const auto& choice_expr = grammar->GetGrammarExpr(choice_id); + if (choice_expr.type == ExprType::kEmptyStr) { + continue; + } + int branch_start_state = target_fsm_->AddState(); + if (!BuildSequence(choice_expr, grammar, branch_start_state, &branch_end_states)) { + return false; + } + target_fsm_->AddEpsilonEdge(start_state, branch_start_state); + end_states->insert(end_states->end(), branch_end_states.begin(), branch_end_states.end()); } + if (nullable) { - FSMWithStartEnd null_fsm; - null_fsm.AddState(); - null_fsm.SetStartState(0); - null_fsm.AddEndState(0); - fsm_list.push_back(std::move(null_fsm)); + int nullable_branch_state = target_fsm_->AddState(); + target_fsm_->AddEpsilonEdge(start_state, nullable_branch_state); + end_states->push_back(nullable_branch_state); } + return true; +} - auto result = FSMWithStartEnd::Union(fsm_list); +std::optional GrammarFSMBuilderImpl::Choices( + const GrammarExpr& expr, const Grammar& grammar +) { + FSM result_fsm; + int start_state = result_fsm.AddState(); + std::vector end_states; + GrammarFSMBuilderImpl builder(&result_fsm); + if (!builder.BuildChoices(expr, grammar, start_state, &end_states)) { + return std::nullopt; + } + FSMWithStartEnd result(result_fsm, start_state, std::move(end_states)); result = result.SimplifyEpsilon(); result = result.MergeEquivalentStates(); return result; From 0b901002f583a80aa8d25e448cb8768719889269 Mon Sep 17 00:00:00 2001 From: Ubospica Date: Tue, 28 Jul 2026 07:49:24 -0400 Subject: [PATCH 6/8] refactor: clarify grammar FSM builder lifecycles Require fragment builders to hold a valid target FSM and centralize per-rule dispatch, while preserving existing standalone builder entry points. Add golden structure checks that detect any change in state numbering, edge order, or complete-FSM layout. --- cpp/grammar_functor.cc | 173 ++++++++++--------- tests/python/test_fsm_structure_stability.py | 150 ++++++++++++++++ 2 files changed, 240 insertions(+), 83 deletions(-) create mode 100644 tests/python/test_fsm_structure_stability.py diff --git a/cpp/grammar_functor.cc b/cpp/grammar_functor.cc index cba1ddff0..dc72ed34d 100644 --- a/cpp/grammar_functor.cc +++ b/cpp/grammar_functor.cc @@ -1089,43 +1089,16 @@ class GrammarFSMBuilderImpl { const static uint32_t kMin4BytesUnicode = 0xF0808080; const static uint32_t kMax4BytesUnicode = 0xF7BFBFBF; - explicit GrammarFSMBuilderImpl(FSM* target_fsm = nullptr) : target_fsm_(target_fsm) {} + explicit GrammarFSMBuilderImpl(FSM& target_fsm) : target_fsm_(target_fsm) {} - void Apply(Grammar* grammar) { + static void Apply(Grammar* grammar) { FSM complete_fsm; std::vector> per_rule_fsms((*grammar)->NumRules()); std::vector state_mapping; for (int i = 0; i < (*grammar)->NumRules(); ++i) { - auto rule = (*grammar)->GetRule(i); - auto grammar_expr = (*grammar)->GetGrammarExpr(rule.body_expr_id); - if (grammar_expr.type == Grammar::Impl::GrammarExprType::kTagDispatch) { - auto rule_fsm = TagDispatch((*grammar)->GetTagDispatch(grammar_expr)); - XGRAMMAR_CHECK(rule_fsm.has_value()) << "Failed to build tag dispatch fsm for rule " << i; - per_rule_fsms[i] = rule_fsm->AddToCompleteFSM(&complete_fsm, &state_mapping); - } else if (grammar_expr.type == Grammar::Impl::GrammarExprType::kTokenTagDispatch) { - auto rule_fsm = TokenTagDispatch((*grammar)->GetTokenTagDispatch(grammar_expr)); - XGRAMMAR_CHECK(rule_fsm.has_value()) - << "Failed to build token tag dispatch fsm for rule " << i; - per_rule_fsms[i] = rule_fsm->AddToCompleteFSM(&complete_fsm, &state_mapping); - } else if (grammar_expr.type == Grammar::Impl::GrammarExprType::kRegex) { - // Every regex rule must have an automaton. - auto regex_str = (*grammar)->GetRegexString(grammar_expr); - auto rule_fsm_result = Regex(regex_str, (*grammar)->GetRegexIsJSONString(grammar_expr)); - if (rule_fsm_result.IsErr()) { - XGRAMMAR_LOG(FATAL) << "Failed to build the automaton for rule " - << (*grammar)->GetRule(i).name << " with regex " << regex_str << ": " - << std::move(rule_fsm_result).UnwrapErr().what(); - } - auto rule_fsm = std::move(rule_fsm_result).Unwrap(); - per_rule_fsms[i] = rule_fsm.AddToCompleteFSM(&complete_fsm, &state_mapping); - } else { - XGRAMMAR_DCHECK(grammar_expr.type == Grammar::Impl::GrammarExprType::kChoices); - auto rule_fsm = Choices(grammar_expr, *grammar); - if (rule_fsm.has_value()) { - per_rule_fsms[i] = rule_fsm->AddToCompleteFSM(&complete_fsm, &state_mapping); - } - } + auto rule_fsm = BuildRuleFSM(*grammar, i); + per_rule_fsms[i] = rule_fsm.AddToCompleteFSM(&complete_fsm, &state_mapping); } for (int i = 0; i < (*grammar)->NumRules(); ++i) { @@ -1176,6 +1149,7 @@ class GrammarFSMBuilderImpl { ); private: + static FSMWithStartEnd BuildRuleFSM(Grammar& grammar, int rule_id); static FSMWithStartEnd BuildSingleElement(const GrammarExpr& expr); bool BuildElement( const GrammarExpr& element_expr, int start_state, std::vector* end_states @@ -1198,7 +1172,7 @@ class GrammarFSMBuilderImpl { void BuildNegativeCharacterClass(const GrammarExpr& expr, int start_state, int end_state); void AddCharacterRange(int from, int to, uint32_t min, uint32_t max); - FSM* target_fsm_; + FSM& target_fsm_; }; // This function will add a range [min, max] of characters to the FSM, and the length @@ -1333,7 +1307,6 @@ void AddSameLengthCharacterRange(FSM& fsm, int from, int to, uint32_t min, uint3 // This function will add a range [min, max] of unicode characters to the FSM. void GrammarFSMBuilderImpl::AddCharacterRange(int from, int to, uint32_t min, uint32_t max) { - XGRAMMAR_DCHECK(target_fsm_ != nullptr); XGRAMMAR_CHECK(min <= max) << "Invalid character range: min (" << min << ") > max (" << max << ")"; // Ensure max and min are valid unicode value. @@ -1367,46 +1340,46 @@ void GrammarFSMBuilderImpl::AddCharacterRange(int from, int to, uint32_t min, ui // Step2. Divide the range into several ranges, which contain characters with different lengths. if (max <= kMax1ByteUnicode) { - AddSameLengthCharacterRange(*target_fsm_, from, to, min, max); + AddSameLengthCharacterRange(target_fsm_, from, to, min, max); return; } if (max <= kMax2BytesUnicode) { if (min >= kMin2BytesUnicode) { - AddSameLengthCharacterRange(*target_fsm_, from, to, min, max); + AddSameLengthCharacterRange(target_fsm_, from, to, min, max); } else { - AddSameLengthCharacterRange(*target_fsm_, from, to, min, kMax1ByteUnicode); - AddSameLengthCharacterRange(*target_fsm_, from, to, kMin2BytesUnicode, max); + AddSameLengthCharacterRange(target_fsm_, from, to, min, kMax1ByteUnicode); + AddSameLengthCharacterRange(target_fsm_, from, to, kMin2BytesUnicode, max); } return; } if (max <= kMax3BytesUnicode) { if (min >= kMin3BytesUnicode) { - AddSameLengthCharacterRange(*target_fsm_, from, to, min, max); + AddSameLengthCharacterRange(target_fsm_, from, to, min, max); } else if (min >= kMin2BytesUnicode) { - AddSameLengthCharacterRange(*target_fsm_, from, to, min, kMax2BytesUnicode); - AddSameLengthCharacterRange(*target_fsm_, from, to, kMin3BytesUnicode, max); + AddSameLengthCharacterRange(target_fsm_, from, to, min, kMax2BytesUnicode); + AddSameLengthCharacterRange(target_fsm_, from, to, kMin3BytesUnicode, max); } else { - AddSameLengthCharacterRange(*target_fsm_, from, to, min, kMax1ByteUnicode); - AddSameLengthCharacterRange(*target_fsm_, from, to, kMin2BytesUnicode, kMax2BytesUnicode); - AddSameLengthCharacterRange(*target_fsm_, from, to, kMin3BytesUnicode, max); + AddSameLengthCharacterRange(target_fsm_, from, to, min, kMax1ByteUnicode); + AddSameLengthCharacterRange(target_fsm_, from, to, kMin2BytesUnicode, kMax2BytesUnicode); + AddSameLengthCharacterRange(target_fsm_, from, to, kMin3BytesUnicode, max); } return; } XGRAMMAR_CHECK(max <= kMax4BytesUnicode); if (min >= kMin4BytesUnicode) { - AddSameLengthCharacterRange(*target_fsm_, from, to, min, max); + AddSameLengthCharacterRange(target_fsm_, from, to, min, max); } else if (min >= kMin3BytesUnicode) { - AddSameLengthCharacterRange(*target_fsm_, from, to, min, kMax3BytesUnicode); - AddSameLengthCharacterRange(*target_fsm_, from, to, kMin4BytesUnicode, max); + AddSameLengthCharacterRange(target_fsm_, from, to, min, kMax3BytesUnicode); + AddSameLengthCharacterRange(target_fsm_, from, to, kMin4BytesUnicode, max); } else if (min >= kMin2BytesUnicode) { - AddSameLengthCharacterRange(*target_fsm_, from, to, min, kMax2BytesUnicode); - AddSameLengthCharacterRange(*target_fsm_, from, to, kMin3BytesUnicode, kMax3BytesUnicode); - AddSameLengthCharacterRange(*target_fsm_, from, to, kMin4BytesUnicode, max); + AddSameLengthCharacterRange(target_fsm_, from, to, min, kMax2BytesUnicode); + AddSameLengthCharacterRange(target_fsm_, from, to, kMin3BytesUnicode, kMax3BytesUnicode); + AddSameLengthCharacterRange(target_fsm_, from, to, kMin4BytesUnicode, max); } else { - AddSameLengthCharacterRange(*target_fsm_, from, to, min, kMax1ByteUnicode); - AddSameLengthCharacterRange(*target_fsm_, from, to, kMin2BytesUnicode, kMax2BytesUnicode); - AddSameLengthCharacterRange(*target_fsm_, from, to, kMin3BytesUnicode, kMax3BytesUnicode); - AddSameLengthCharacterRange(*target_fsm_, from, to, kMin4BytesUnicode, max); + AddSameLengthCharacterRange(target_fsm_, from, to, min, kMax1ByteUnicode); + AddSameLengthCharacterRange(target_fsm_, from, to, kMin2BytesUnicode, kMax2BytesUnicode); + AddSameLengthCharacterRange(target_fsm_, from, to, kMin3BytesUnicode, kMax3BytesUnicode); + AddSameLengthCharacterRange(target_fsm_, from, to, kMin4BytesUnicode, max); } return; } @@ -1418,7 +1391,6 @@ void GrammarFSMBuilderImpl::BuildNegativeCharacterClass( expr.type == ExprType::kCharacterClass || expr.type == ExprType::kCharacterClassStar ); XGRAMMAR_DCHECK(expr[0]); // Negative character class should be true. - XGRAMMAR_DCHECK(target_fsm_ != nullptr); std::bitset<128> char_set; for (int i = 1; i < static_cast(expr.size()); i += 2) { uint8_t byte_min = static_cast(expr[i]); @@ -1441,7 +1413,7 @@ void GrammarFSMBuilderImpl::BuildNegativeCharacterClass( while (right_bound < 128 && !char_set[right_bound]) { right_bound++; } - target_fsm_->AddEdge( + target_fsm_.AddEdge( start_state, end_state, static_cast(left_bound), @@ -1456,10 +1428,8 @@ void GrammarFSMBuilderImpl::BuildNegativeCharacterClass( void GrammarFSMBuilderImpl::BuildCharacterClass( const GrammarExpr& expr, int start_state, std::vector* end_states ) { - XGRAMMAR_DCHECK(target_fsm_ != nullptr); end_states->clear(); - int end_state = - expr.type == ExprType::kCharacterClassStar ? start_state : target_fsm_->AddState(); + int end_state = expr.type == ExprType::kCharacterClassStar ? start_state : target_fsm_.AddState(); bool is_negative = expr[0]; if (is_negative) { BuildNegativeCharacterClass(expr, start_state, end_state); @@ -1476,11 +1446,51 @@ void GrammarFSMBuilderImpl::BuildCharacterClass( end_states->push_back(end_state); } +FSMWithStartEnd GrammarFSMBuilderImpl::BuildRuleFSM(Grammar& grammar, int rule_id) { + auto rule = grammar->GetRule(rule_id); + auto grammar_expr = grammar->GetGrammarExpr(rule.body_expr_id); + switch (grammar_expr.type) { + case ExprType::kTagDispatch: { + auto rule_fsm = TagDispatch(grammar->GetTagDispatch(grammar_expr)); + XGRAMMAR_CHECK(rule_fsm.has_value()) + << "Failed to build tag dispatch fsm for rule " << rule_id; + return std::move(*rule_fsm); + } + case ExprType::kTokenTagDispatch: { + auto rule_fsm = TokenTagDispatch(grammar->GetTokenTagDispatch(grammar_expr)); + XGRAMMAR_CHECK(rule_fsm.has_value()) + << "Failed to build token tag dispatch fsm for rule " << rule_id; + return std::move(*rule_fsm); + } + case ExprType::kRegex: { + auto regex_string = grammar->GetRegexString(grammar_expr); + auto rule_fsm_result = Regex(regex_string, grammar->GetRegexIsJSONString(grammar_expr)); + if (rule_fsm_result.IsErr()) { + XGRAMMAR_LOG(FATAL) << "Failed to build the automaton for rule " << rule.name + << " with regex " << regex_string << ": " + << std::move(rule_fsm_result).UnwrapErr().what(); + } + return std::move(rule_fsm_result).Unwrap(); + } + case ExprType::kChoices: { + auto rule_fsm = Choices(grammar_expr, grammar); + XGRAMMAR_CHECK(rule_fsm.has_value()) + << "Failed to build choices fsm for rule " << rule_id << " (" << rule.name << ")"; + return std::move(*rule_fsm); + } + default: + XGRAMMAR_LOG(FATAL) << "Unsupported grammar expression type " + << static_cast(grammar_expr.type) << " for rule " << rule_id << " (" + << rule.name << ")"; + } + XGRAMMAR_UNREACHABLE(); +} + FSMWithStartEnd GrammarFSMBuilderImpl::BuildSingleElement(const GrammarExpr& expr) { FSM result_fsm; int start_state = result_fsm.AddState(); std::vector end_states; - GrammarFSMBuilderImpl builder(&result_fsm); + GrammarFSMBuilderImpl builder(result_fsm); XGRAMMAR_CHECK(builder.BuildElement(expr, start_state, &end_states)) << "Unsupported grammar expression type: " << static_cast(expr.type); return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); @@ -1546,15 +1556,14 @@ std::optional GrammarFSMBuilderImpl::TokenTagDispatch( bool GrammarFSMBuilderImpl::BuildElement( const GrammarExpr& element_expr, int start_state, std::vector* end_states ) { - XGRAMMAR_DCHECK(target_fsm_ != nullptr); - XGRAMMAR_DCHECK(start_state >= 0 && start_state < target_fsm_->NumStates()); + XGRAMMAR_DCHECK(start_state >= 0 && start_state < target_fsm_.NumStates()); end_states->clear(); switch (element_expr.type) { case (ExprType::kByteString): { int current_state = start_state; for (int32_t byte : element_expr) { - int next_state = target_fsm_->AddState(); - target_fsm_->AddEdge( + int next_state = target_fsm_.AddState(); + target_fsm_.AddEdge( current_state, next_state, static_cast(byte), static_cast(byte) ); current_state = next_state; @@ -1563,8 +1572,8 @@ bool GrammarFSMBuilderImpl::BuildElement( return true; } case (ExprType::kRuleRef): { - int end_state = target_fsm_->AddState(); - target_fsm_->AddRuleEdge(start_state, end_state, element_expr[0]); + int end_state = target_fsm_.AddState(); + target_fsm_.AddRuleEdge(start_state, end_state, element_expr[0]); end_states->push_back(end_state); return true; } @@ -1574,8 +1583,8 @@ bool GrammarFSMBuilderImpl::BuildElement( return true; } case (ExprType::kRepeat): { - int end_state = target_fsm_->AddState(); - target_fsm_->AddRepeatEdge( + int end_state = target_fsm_.AddState(); + target_fsm_.AddRepeatEdge( start_state, end_state, element_expr[0], element_expr[1], element_expr[2] ); end_states->push_back(end_state); @@ -1583,12 +1592,12 @@ bool GrammarFSMBuilderImpl::BuildElement( } case (ExprType::kToken): case (ExprType::kExcludeToken): { - int end_state = target_fsm_->AddState(); + int end_state = target_fsm_.AddState(); std::vector token_ids(element_expr.begin(), element_expr.end()); if (element_expr.type == ExprType::kToken) { - target_fsm_->AddTokenEdge(start_state, end_state, token_ids); + target_fsm_.AddTokenEdge(start_state, end_state, token_ids); } else { - target_fsm_->AddExcludeTokenEdge(start_state, end_state, token_ids); + target_fsm_.AddExcludeTokenEdge(start_state, end_state, token_ids); } end_states->push_back(end_state); return true; @@ -1605,8 +1614,7 @@ bool GrammarFSMBuilderImpl::BuildSequence( int start_state, std::vector* end_states ) { - XGRAMMAR_DCHECK(target_fsm_ != nullptr); - XGRAMMAR_DCHECK(start_state >= 0 && start_state < target_fsm_->NumStates()); + XGRAMMAR_DCHECK(start_state >= 0 && start_state < target_fsm_.NumStates()); end_states->clear(); if (expr.size() == 0) { end_states->push_back(start_state); @@ -1619,12 +1627,12 @@ bool GrammarFSMBuilderImpl::BuildSequence( std::vector next_end_states; for (int index = 1; index < static_cast(expr.size()); ++index) { - int element_start = target_fsm_->AddState(); + int element_start = target_fsm_.AddState(); if (!BuildElement(grammar->GetGrammarExpr(expr[index]), element_start, &next_end_states)) { return false; } for (int32_t previous_end_state : *end_states) { - target_fsm_->AddEpsilonEdge(previous_end_state, element_start); + target_fsm_.AddEpsilonEdge(previous_end_state, element_start); } end_states->swap(next_end_states); } @@ -1638,7 +1646,7 @@ std::optional GrammarFSMBuilderImpl::Sequence( FSM result_fsm; int start_state = result_fsm.AddState(); std::vector end_states; - GrammarFSMBuilderImpl builder(&result_fsm); + GrammarFSMBuilderImpl builder(result_fsm); if (!builder.BuildSequence(expr, grammar, start_state, &end_states)) { return std::nullopt; } @@ -1652,8 +1660,7 @@ bool GrammarFSMBuilderImpl::BuildChoices( std::vector* end_states ) { XGRAMMAR_DCHECK(expr.type == ExprType::kChoices); - XGRAMMAR_DCHECK(target_fsm_ != nullptr); - XGRAMMAR_DCHECK(start_state >= 0 && start_state < target_fsm_->NumStates()); + XGRAMMAR_DCHECK(start_state >= 0 && start_state < target_fsm_.NumStates()); end_states->clear(); int sequence_count = 0; @@ -1689,17 +1696,17 @@ bool GrammarFSMBuilderImpl::BuildChoices( if (choice_expr.type == ExprType::kEmptyStr) { continue; } - int branch_start_state = target_fsm_->AddState(); + int branch_start_state = target_fsm_.AddState(); if (!BuildSequence(choice_expr, grammar, branch_start_state, &branch_end_states)) { return false; } - target_fsm_->AddEpsilonEdge(start_state, branch_start_state); + target_fsm_.AddEpsilonEdge(start_state, branch_start_state); end_states->insert(end_states->end(), branch_end_states.begin(), branch_end_states.end()); } if (nullable) { - int nullable_branch_state = target_fsm_->AddState(); - target_fsm_->AddEpsilonEdge(start_state, nullable_branch_state); + int nullable_branch_state = target_fsm_.AddState(); + target_fsm_.AddEpsilonEdge(start_state, nullable_branch_state); end_states->push_back(nullable_branch_state); } return true; @@ -1711,7 +1718,7 @@ std::optional GrammarFSMBuilderImpl::Choices( FSM result_fsm; int start_state = result_fsm.AddState(); std::vector end_states; - GrammarFSMBuilderImpl builder(&result_fsm); + GrammarFSMBuilderImpl builder(result_fsm); if (!builder.BuildChoices(expr, grammar, start_state, &end_states)) { return std::nullopt; } @@ -2917,7 +2924,7 @@ Grammar StructureNormalizer::Apply(const Grammar& grammar) { /*************************** Forward grammar optimizers to their impl ***************************/ -void GrammarFSMBuilder::Apply(Grammar* grammar) { GrammarFSMBuilderImpl().Apply(grammar); } +void GrammarFSMBuilder::Apply(Grammar* grammar) { GrammarFSMBuilderImpl::Apply(grammar); } void RepetitionNormalizer::Apply(Grammar* grammar) { RepetitionNormalizerImpl().Apply(grammar); } diff --git a/tests/python/test_fsm_structure_stability.py b/tests/python/test_fsm_structure_stability.py new file mode 100644 index 000000000..984bd12a7 --- /dev/null +++ b/tests/python/test_fsm_structure_stability.py @@ -0,0 +1,150 @@ +"""Guard the exact FSM layout produced by grammar compilation.""" + +import hashlib +import json + +import pytest + +import xgrammar as xgr +from xgrammar.testing import _print_grammar_fsms + +FSM_STRUCTURE_CASES = [ + ( + "ebnf", + 'root ::= "a" [0-9] "b"', + "55b4f598dd003190ab972e2c89246b46b7ad102506fdfa785003c270866e8572", + ), + ( + "ebnf", + 'root ::= "hello" [a-zA-Z_] [0-9]* "world"', + "bf85be1292ab8bfcbdb08c09e4508ffe855d110cc19b36fd849c6724897c6a50", + ), + ( + "ebnf", + 'root ::= "<" [a-c]* ">"', + "19cef267fb8eddbc9c5ac8fba92c135600fc720258abf0bd58cb6c346fe75c5a", + ), + ( + "ebnf", + 'root ::= "x" [^0-9] "y" [^a-z]* "z"', + "fd5ab210c129fe4da573c175bfd0e97440d30a88d038a160e5bcbeb2f9fc00d0", + ), + ( + "ebnf", + 'root ::= "(" inner ")" inner\ninner ::= [0-9] [0-9]', + "4de2c028805f73a42e2eb3766defb94a7b6a9cfb735fd9d9a7bf8c610e4b4012", + ), + ( + "ebnf", + 'root ::= "a" item{2,5} "b"\nitem ::= [0-9]', + "eddead662499416bbd37a4b70d70912ace125deaeaa8fc1f48fbe61dde26b8c7", + ), + ( + "ebnf", + 'root ::= item{3,}\nitem ::= "ab" [xy]', + "e80cfab029429aa5111b5b1b95c002f7fa4bc7b0a4e88ecae2d507067dba7b6d", + ), + ( + "ebnf", + 'root ::= "ab" [0-9] | "cd" sub | sub sub\nsub ::= [a-f] "q"', + "29090db9590a3b07af82126f08ec25663f657fde2b659c7d172bf6da227861ba", + ), + ( + "ebnf", + 'root ::= "abcdefghijklmnopqrstuvwxyz" [0-9] "ABCDEFGHIJKLMNOPQRSTUVWXYZ"', + "58b89fec9137ba57ef63aba58efa201d7a64b9c949b80d0e4376467613779aab", + ), + ( + "ebnf", + 'root ::= "中文" [\\u4e00-\\u9fff] "端"', + "a2d7f36e1d3a0453d4267fb3937ed07a4abaac7361e051d3aaaaf7a6a58de462", + ), + ("ebnf", 'root ::= "only"', "7880ea9fd56d9cca1976659fa1ccd4b7b347ab2688800e96bdafcfaa591019e5"), + ("ebnf", "root ::= [0-9]", "e210303c507b9b23c732079d40b3f0b01af0de6deb86c8974148dbb92b92896c"), + ( + "ebnf", + 'root ::= "" "a" ""', + "8698019eb93735ab521cfd930bf7688a276b4ebfb33ccd7cea20f8139f99fbeb", + ), + ( + "ebnf", + 'root ::= a b c\na ::= "x" [0-9]*\nb ::= a "y" | [^xyz]\nc ::= b{1,3} "end"', + "692de2ea3a886cc7aca953da58fe7dc66c26bc39a71b01fdeee53cfb58e89561", + ), + ( + "json_schema", + { + "type": "object", + "properties": { + "name": {"type": "string"}, + "age": {"type": "integer"}, + "tags": {"type": "array", "items": {"type": "string"}}, + "nested": { + "type": "object", + "properties": {"value": {"type": "number"}}, + "required": ["value"], + }, + }, + "required": ["name", "age", "tags", "nested"], + "additionalProperties": False, + }, + "4d050454cf0d4a9a579a7f0d33803c6a98a0354649945be7e4830505c473ab1c", + ), + ( + "json_schema", + { + "type": "array", + "items": { + "type": "object", + "properties": { + "id": {"type": "string", "pattern": "[a-f0-9]{8}"}, + "kind": {"enum": ["alpha", "beta", "gamma"]}, + }, + "required": ["id", "kind"], + }, + "minItems": 1, + "maxItems": 4, + }, + "d0284724594dcee4009b06e2ef7d210b7614f062482caa5df925a10619af4323", + ), + ( + "structural_tag", + { + "type": "structural_tag", + "format": { + "type": "dispatch", + "rules": [ + ["", {"type": "const_string", "value": "A"}], + ["", {"type": "json_schema", "json_schema": {"type": "object"}}], + ], + "loop": True, + "excludes": [], + }, + }, + "61e2ccf44a3c589547de488dfbd6cd84a9a4fa39ce86ed3f71a11e388cc4050c", + ), +] + + +@pytest.fixture(scope="module") +def compiler(): + vocabulary = [chr(character) for character in range(33, 127)] + ["中", "文", "端"] + tokenizer_info = xgr.TokenizerInfo(vocabulary) + return xgr.GrammarCompiler(tokenizer_info, max_threads=1, cache_enabled=False) + + +@pytest.mark.parametrize("kind, source, expected_digest", FSM_STRUCTURE_CASES) +def test_compiled_fsm_structure_stable(compiler, kind, source, expected_digest): + """Detect changes to state numbers, edge order, endpoints, or complete-FSM layout.""" + if kind == "ebnf": + compiled = compiler.compile_grammar(source) + elif kind == "json_schema": + compiled = compiler.compile_json_schema( + json.dumps(source, separators=(",", ":")), any_whitespace=True + ) + else: + compiled = compiler.compile_structural_tag(source) + + printed_fsm = _print_grammar_fsms(compiled.grammar) + actual_digest = hashlib.sha256(printed_fsm.encode()).hexdigest() + assert actual_digest == expected_digest From cd5e551cdd4b9dc021dd829dba33c7b2dc4457f3 Mon Sep 17 00:00:00 2001 From: Ubospica Date: Tue, 28 Jul 2026 07:54:51 -0400 Subject: [PATCH 7/8] test: record FSM structure from the main build Regenerate the JSON-pattern case with the explicit main native library so the golden digest tracks the intended baseline rather than an installed package. --- tests/python/test_fsm_structure_stability.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/python/test_fsm_structure_stability.py b/tests/python/test_fsm_structure_stability.py index 984bd12a7..cfcf591f6 100644 --- a/tests/python/test_fsm_structure_stability.py +++ b/tests/python/test_fsm_structure_stability.py @@ -105,7 +105,7 @@ "minItems": 1, "maxItems": 4, }, - "d0284724594dcee4009b06e2ef7d210b7614f062482caa5df925a10619af4323", + "a7a376ee74bc9e59e36da46fd7dddc9c322555c96110120260779bd07a9d14c5", ), ( "structural_tag", From 23746d7ff52385fd4d261641cfdc46d641aa7e92 Mon Sep 17 00:00:00 2001 From: Ubospica Date: Tue, 28 Jul 2026 08:38:06 -0400 Subject: [PATCH 8/8] refactor: support nested expressions in grammar FSM builder Build rule state machines recursively so construction no longer depends on the normalized choice-sequence-element layout. --- cpp/grammar_functor.cc | 515 +++++++++++++++--------- cpp/grammar_impl.h | 8 +- cpp/tvm_ffi/tvm_ffi.cc | 20 + python/xgrammar/testing.py | 12 + tests/python/test_fsm_sequence_build.py | 40 ++ 5 files changed, 410 insertions(+), 185 deletions(-) diff --git a/cpp/grammar_functor.cc b/cpp/grammar_functor.cc index dc72ed34d..20cdd1af6 100644 --- a/cpp/grammar_functor.cc +++ b/cpp/grammar_functor.cc @@ -1089,7 +1089,8 @@ class GrammarFSMBuilderImpl { const static uint32_t kMin4BytesUnicode = 0xF0808080; const static uint32_t kMax4BytesUnicode = 0xF7BFBFBF; - explicit GrammarFSMBuilderImpl(FSM& target_fsm) : target_fsm_(target_fsm) {} + explicit GrammarFSMBuilderImpl(FSM& target_fsm, const std::string* rule_name = nullptr) + : target_fsm_(target_fsm), rule_name_(rule_name) {} static void Apply(Grammar* grammar) { FSM complete_fsm; @@ -1142,37 +1143,69 @@ class GrammarFSMBuilderImpl { static std::optional TagDispatch(const Grammar::Impl::TagDispatch& tag_dispatch); static Result Regex(const std::string& regex, bool json_string = false); /* Building tool functions.*/ - static std::optional BuildTagDispatch( + static std::optional BuildTagDispatchFSM( const std::vector>& string_trigger_rules, bool loop_after_dispatch, const std::vector& excluded_strings ); private: - static FSMWithStartEnd BuildRuleFSM(Grammar& grammar, int rule_id); - static FSMWithStartEnd BuildSingleElement(const GrammarExpr& expr); - bool BuildElement( - const GrammarExpr& element_expr, int start_state, std::vector* end_states + static FSMWithStartEnd BuildRuleFSM(const Grammar& grammar, int rule_id); + static FSMWithStartEnd BuildExpressionFSM( + const GrammarExpr& expr, const Grammar& grammar, const std::string* rule_name = nullptr ); - bool BuildSequence( + void BuildExpression( const GrammarExpr& expr, const Grammar& grammar, int start_state, std::vector* end_states ); - bool BuildChoices( + void BuildEmptyString(int start_state, std::vector* end_states); + void BuildByteString(const GrammarExpr& expr, int start_state, std::vector* end_states); + void BuildRuleRef(const GrammarExpr& expr, int start_state, std::vector* end_states); + void BuildCharacterClass( + const GrammarExpr& expr, int start_state, std::vector* end_states + ); + void BuildCharacterClassStar( + const GrammarExpr& expr, int start_state, std::vector* end_states + ); + void BuildRepeat(const GrammarExpr& expr, int start_state, std::vector* end_states); + void BuildToken(const GrammarExpr& expr, int start_state, std::vector* end_states); + void BuildExcludeToken( + const GrammarExpr& expr, int start_state, std::vector* end_states + ); + void BuildRegex( + const std::string& regex, bool json_string, int start_state, std::vector* end_states + ); + void BuildTagDispatch( + const Grammar::Impl::TagDispatch& tag_dispatch, + int start_state, + std::vector* end_states + ); + void BuildTokenTagDispatch( + const Grammar::Impl::TokenTagDispatch& token_tag_dispatch, + int start_state, + std::vector* end_states + ); + void BuildSequence( const GrammarExpr& expr, const Grammar& grammar, int start_state, std::vector* end_states ); - void BuildCharacterClass( - const GrammarExpr& expr, int start_state, std::vector* end_states + void BuildChoices( + const GrammarExpr& expr, + const Grammar& grammar, + int start_state, + std::vector* end_states ); + void AddCharacterClassTransitions(const GrammarExpr& expr, int start_state, int end_state); void BuildNegativeCharacterClass(const GrammarExpr& expr, int start_state, int end_state); + void AppendFSM(FSMWithStartEnd fsm, int start_state, std::vector* end_states); void AddCharacterRange(int from, int to, uint32_t min, uint32_t max); FSM& target_fsm_; + const std::string* rule_name_; }; // This function will add a range [min, max] of characters to the FSM, and the length @@ -1425,11 +1458,12 @@ void GrammarFSMBuilderImpl::BuildNegativeCharacterClass( AddCharacterRange(start_state, end_state, kMin2BytesUnicode, kMax4BytesUnicode); } -void GrammarFSMBuilderImpl::BuildCharacterClass( - const GrammarExpr& expr, int start_state, std::vector* end_states +void GrammarFSMBuilderImpl::AddCharacterClassTransitions( + const GrammarExpr& expr, int start_state, int end_state ) { - end_states->clear(); - int end_state = expr.type == ExprType::kCharacterClassStar ? start_state : target_fsm_.AddState(); + XGRAMMAR_DCHECK( + expr.type == ExprType::kCharacterClass || expr.type == ExprType::kCharacterClassStar + ); bool is_negative = expr[0]; if (is_negative) { BuildNegativeCharacterClass(expr, start_state, end_state); @@ -1443,172 +1477,215 @@ void GrammarFSMBuilderImpl::BuildCharacterClass( AddCharacterRange(start_state, end_state, packed_min, packed_max); } } +} + +void GrammarFSMBuilderImpl::BuildCharacterClass( + const GrammarExpr& expr, int start_state, std::vector* end_states +) { + XGRAMMAR_DCHECK(expr.type == ExprType::kCharacterClass); + end_states->clear(); + int end_state = target_fsm_.AddState(); + AddCharacterClassTransitions(expr, start_state, end_state); end_states->push_back(end_state); } -FSMWithStartEnd GrammarFSMBuilderImpl::BuildRuleFSM(Grammar& grammar, int rule_id) { - auto rule = grammar->GetRule(rule_id); - auto grammar_expr = grammar->GetGrammarExpr(rule.body_expr_id); - switch (grammar_expr.type) { - case ExprType::kTagDispatch: { - auto rule_fsm = TagDispatch(grammar->GetTagDispatch(grammar_expr)); - XGRAMMAR_CHECK(rule_fsm.has_value()) - << "Failed to build tag dispatch fsm for rule " << rule_id; - return std::move(*rule_fsm); - } - case ExprType::kTokenTagDispatch: { - auto rule_fsm = TokenTagDispatch(grammar->GetTokenTagDispatch(grammar_expr)); - XGRAMMAR_CHECK(rule_fsm.has_value()) - << "Failed to build token tag dispatch fsm for rule " << rule_id; - return std::move(*rule_fsm); - } - case ExprType::kRegex: { - auto regex_string = grammar->GetRegexString(grammar_expr); - auto rule_fsm_result = Regex(regex_string, grammar->GetRegexIsJSONString(grammar_expr)); - if (rule_fsm_result.IsErr()) { - XGRAMMAR_LOG(FATAL) << "Failed to build the automaton for rule " << rule.name - << " with regex " << regex_string << ": " - << std::move(rule_fsm_result).UnwrapErr().what(); - } - return std::move(rule_fsm_result).Unwrap(); - } - case ExprType::kChoices: { - auto rule_fsm = Choices(grammar_expr, grammar); - XGRAMMAR_CHECK(rule_fsm.has_value()) - << "Failed to build choices fsm for rule " << rule_id << " (" << rule.name << ")"; - return std::move(*rule_fsm); - } - default: - XGRAMMAR_LOG(FATAL) << "Unsupported grammar expression type " - << static_cast(grammar_expr.type) << " for rule " << rule_id << " (" - << rule.name << ")"; - } - XGRAMMAR_UNREACHABLE(); +void GrammarFSMBuilderImpl::BuildCharacterClassStar( + const GrammarExpr& expr, int start_state, std::vector* end_states +) { + XGRAMMAR_DCHECK(expr.type == ExprType::kCharacterClassStar); + end_states->clear(); + AddCharacterClassTransitions(expr, start_state, start_state); + end_states->push_back(start_state); +} + +FSMWithStartEnd GrammarFSMBuilderImpl::BuildRuleFSM(const Grammar& grammar, int rule_id) { + const auto& rule = grammar->GetRule(rule_id); + return BuildExpressionFSM(grammar->GetGrammarExpr(rule.body_expr_id), grammar, &rule.name); } -FSMWithStartEnd GrammarFSMBuilderImpl::BuildSingleElement(const GrammarExpr& expr) { +FSMWithStartEnd GrammarFSMBuilderImpl::BuildExpressionFSM( + const GrammarExpr& expr, const Grammar& grammar, const std::string* rule_name +) { FSM result_fsm; int start_state = result_fsm.AddState(); std::vector end_states; - GrammarFSMBuilderImpl builder(result_fsm); - XGRAMMAR_CHECK(builder.BuildElement(expr, start_state, &end_states)) - << "Unsupported grammar expression type: " << static_cast(expr.type); - return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); + GrammarFSMBuilderImpl builder(result_fsm, rule_name); + builder.BuildExpression(expr, grammar, start_state, &end_states); + FSMWithStartEnd result(result_fsm, start_state, std::move(end_states)); + if (expr.type != ExprType::kTagDispatch && expr.type != ExprType::kTokenTagDispatch) { + result = result.SimplifyEpsilon(); + result = result.MergeEquivalentStates(); + } + return result; } FSMWithStartEnd GrammarFSMBuilderImpl::RuleRef(const GrammarExpr& expr) { - return BuildSingleElement(expr); + FSM result_fsm; + int start_state = result_fsm.AddState(); + std::vector end_states; + GrammarFSMBuilderImpl builder(result_fsm); + builder.BuildRuleRef(expr, start_state, &end_states); + return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); } FSMWithStartEnd GrammarFSMBuilderImpl::CharacterClass(const GrammarExpr& expr) { - return BuildSingleElement(expr); + FSM result_fsm; + int start_state = result_fsm.AddState(); + std::vector end_states; + GrammarFSMBuilderImpl builder(result_fsm); + if (expr.type == ExprType::kCharacterClassStar) { + builder.BuildCharacterClassStar(expr, start_state, &end_states); + } else { + builder.BuildCharacterClass(expr, start_state, &end_states); + } + return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); } FSMWithStartEnd GrammarFSMBuilderImpl::ByteString(const GrammarExpr& expr) { - return BuildSingleElement(expr); + FSM result_fsm; + int start_state = result_fsm.AddState(); + std::vector end_states; + GrammarFSMBuilderImpl builder(result_fsm); + builder.BuildByteString(expr, start_state, &end_states); + return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); } FSMWithStartEnd GrammarFSMBuilderImpl::Token(const GrammarExpr& expr) { - return BuildSingleElement(expr); + FSM result_fsm; + int start_state = result_fsm.AddState(); + std::vector end_states; + GrammarFSMBuilderImpl builder(result_fsm); + builder.BuildToken(expr, start_state, &end_states); + return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); } FSMWithStartEnd GrammarFSMBuilderImpl::ExcludeToken(const GrammarExpr& expr) { - return BuildSingleElement(expr); + FSM result_fsm; + int start_state = result_fsm.AddState(); + std::vector end_states; + GrammarFSMBuilderImpl builder(result_fsm); + builder.BuildExcludeToken(expr, start_state, &end_states); + return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); } std::optional GrammarFSMBuilderImpl::TokenTagDispatch( - const Grammar::Impl::TokenTagDispatch& ttd + const Grammar::Impl::TokenTagDispatch& token_tag_dispatch ) { - int num_triggers = static_cast(ttd.trigger_rule_pairs.size()); - bool loop = ttd.loop_after_dispatch; - int num_states = 1 + num_triggers + (loop ? 0 : 1); - FSM fsm(num_states); - std::vector ends; - int start = 0; - ends.push_back(start); - int end_state = -1; - if (!loop) { - end_state = num_states - 1; - ends.push_back(end_state); - } - std::vector self_loop_exclude; - for (const auto& [token_id, rule_id] : ttd.trigger_rule_pairs) { - self_loop_exclude.push_back(token_id); - } - for (auto excl_id : ttd.excludes) { - self_loop_exclude.push_back(excl_id); - } - std::sort(self_loop_exclude.begin(), self_loop_exclude.end()); - self_loop_exclude.erase( - std::unique(self_loop_exclude.begin(), self_loop_exclude.end()), self_loop_exclude.end() - ); - for (int i = 0; i < num_triggers; ++i) { - int dispatch_state = 1 + i; - auto [token_id, rule_id] = ttd.trigger_rule_pairs[i]; - fsm.AddTokenEdge(start, dispatch_state, {token_id}); - int target = loop ? start : end_state; - fsm.AddRuleEdge(dispatch_state, target, static_cast(rule_id)); - } - fsm.AddExcludeTokenEdge(start, start, self_loop_exclude); - return FSMWithStartEnd(fsm, start, ends); + FSM result_fsm; + int start_state = result_fsm.AddState(); + std::vector end_states; + GrammarFSMBuilderImpl builder(result_fsm); + builder.BuildTokenTagDispatch(token_tag_dispatch, start_state, &end_states); + return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); } -bool GrammarFSMBuilderImpl::BuildElement( - const GrammarExpr& element_expr, int start_state, std::vector* end_states +void GrammarFSMBuilderImpl::BuildExpression( + const GrammarExpr& expr, + const Grammar& grammar, + int start_state, + std::vector* end_states ) { XGRAMMAR_DCHECK(start_state >= 0 && start_state < target_fsm_.NumStates()); - end_states->clear(); - switch (element_expr.type) { - case (ExprType::kByteString): { - int current_state = start_state; - for (int32_t byte : element_expr) { - int next_state = target_fsm_.AddState(); - target_fsm_.AddEdge( - current_state, next_state, static_cast(byte), static_cast(byte) - ); - current_state = next_state; - } - end_states->push_back(current_state); - return true; - } - case (ExprType::kRuleRef): { - int end_state = target_fsm_.AddState(); - target_fsm_.AddRuleEdge(start_state, end_state, element_expr[0]); - end_states->push_back(end_state); - return true; - } - case (ExprType::kCharacterClass): - case (ExprType::kCharacterClassStar): { - BuildCharacterClass(element_expr, start_state, end_states); - return true; - } - case (ExprType::kRepeat): { - int end_state = target_fsm_.AddState(); - target_fsm_.AddRepeatEdge( - start_state, end_state, element_expr[0], element_expr[1], element_expr[2] + switch (expr.type) { + case ExprType::kEmptyStr: + return BuildEmptyString(start_state, end_states); + case ExprType::kByteString: + return BuildByteString(expr, start_state, end_states); + case ExprType::kCharacterClass: + return BuildCharacterClass(expr, start_state, end_states); + case ExprType::kCharacterClassStar: + return BuildCharacterClassStar(expr, start_state, end_states); + case ExprType::kRuleRef: + return BuildRuleRef(expr, start_state, end_states); + case ExprType::kRepeat: + return BuildRepeat(expr, start_state, end_states); + case ExprType::kToken: + return BuildToken(expr, start_state, end_states); + case ExprType::kExcludeToken: + return BuildExcludeToken(expr, start_state, end_states); + case ExprType::kSequence: + return BuildSequence(expr, grammar, start_state, end_states); + case ExprType::kChoices: + return BuildChoices(expr, grammar, start_state, end_states); + case ExprType::kRegex: + return BuildRegex( + grammar->GetRegexString(expr), + grammar->GetRegexIsJSONString(expr), + start_state, + end_states ); - end_states->push_back(end_state); - return true; - } - case (ExprType::kToken): - case (ExprType::kExcludeToken): { - int end_state = target_fsm_.AddState(); - std::vector token_ids(element_expr.begin(), element_expr.end()); - if (element_expr.type == ExprType::kToken) { - target_fsm_.AddTokenEdge(start_state, end_state, token_ids); - } else { - target_fsm_.AddExcludeTokenEdge(start_state, end_state, token_ids); - } - end_states->push_back(end_state); - return true; - } - default: { - return false; - } + case ExprType::kTagDispatch: + return BuildTagDispatch(grammar->GetTagDispatch(expr), start_state, end_states); + case ExprType::kTokenTagDispatch: + return BuildTokenTagDispatch(grammar->GetTokenTagDispatch(expr), start_state, end_states); } + XGRAMMAR_UNREACHABLE(); } -bool GrammarFSMBuilderImpl::BuildSequence( +void GrammarFSMBuilderImpl::BuildEmptyString(int start_state, std::vector* end_states) { + end_states->clear(); + end_states->push_back(start_state); +} + +void GrammarFSMBuilderImpl::BuildByteString( + const GrammarExpr& expr, int start_state, std::vector* end_states +) { + XGRAMMAR_DCHECK(expr.type == ExprType::kByteString); + end_states->clear(); + int current_state = start_state; + for (int32_t byte : expr) { + int next_state = target_fsm_.AddState(); + target_fsm_.AddEdge( + current_state, next_state, static_cast(byte), static_cast(byte) + ); + current_state = next_state; + } + end_states->push_back(current_state); +} + +void GrammarFSMBuilderImpl::BuildRuleRef( + const GrammarExpr& expr, int start_state, std::vector* end_states +) { + XGRAMMAR_DCHECK(expr.type == ExprType::kRuleRef); + end_states->clear(); + int end_state = target_fsm_.AddState(); + target_fsm_.AddRuleEdge(start_state, end_state, expr[0]); + end_states->push_back(end_state); +} + +void GrammarFSMBuilderImpl::BuildRepeat( + const GrammarExpr& expr, int start_state, std::vector* end_states +) { + XGRAMMAR_DCHECK(expr.type == ExprType::kRepeat); + end_states->clear(); + int end_state = target_fsm_.AddState(); + target_fsm_.AddRepeatEdge(start_state, end_state, expr[0], expr[1], expr[2]); + end_states->push_back(end_state); +} + +void GrammarFSMBuilderImpl::BuildToken( + const GrammarExpr& expr, int start_state, std::vector* end_states +) { + XGRAMMAR_DCHECK(expr.type == ExprType::kToken); + end_states->clear(); + int end_state = target_fsm_.AddState(); + target_fsm_.AddTokenEdge(start_state, end_state, std::vector(expr.begin(), expr.end())); + end_states->push_back(end_state); +} + +void GrammarFSMBuilderImpl::BuildExcludeToken( + const GrammarExpr& expr, int start_state, std::vector* end_states +) { + XGRAMMAR_DCHECK(expr.type == ExprType::kExcludeToken); + end_states->clear(); + int end_state = target_fsm_.AddState(); + target_fsm_.AddExcludeTokenEdge( + start_state, end_state, std::vector(expr.begin(), expr.end()) + ); + end_states->push_back(end_state); +} + +void GrammarFSMBuilderImpl::BuildSequence( const GrammarExpr& expr, const Grammar& grammar, int start_state, @@ -1618,26 +1695,20 @@ bool GrammarFSMBuilderImpl::BuildSequence( end_states->clear(); if (expr.size() == 0) { end_states->push_back(start_state); - return true; + return; } - if (!BuildElement(grammar->GetGrammarExpr(expr[0]), start_state, end_states)) { - return false; - } + BuildExpression(grammar->GetGrammarExpr(expr[0]), grammar, start_state, end_states); std::vector next_end_states; for (int index = 1; index < static_cast(expr.size()); ++index) { int element_start = target_fsm_.AddState(); - if (!BuildElement(grammar->GetGrammarExpr(expr[index]), element_start, &next_end_states)) { - return false; - } + BuildExpression(grammar->GetGrammarExpr(expr[index]), grammar, element_start, &next_end_states); for (int32_t previous_end_state : *end_states) { target_fsm_.AddEpsilonEdge(previous_end_state, element_start); } end_states->swap(next_end_states); } - - return true; } std::optional GrammarFSMBuilderImpl::Sequence( @@ -1647,13 +1718,11 @@ std::optional GrammarFSMBuilderImpl::Sequence( int start_state = result_fsm.AddState(); std::vector end_states; GrammarFSMBuilderImpl builder(result_fsm); - if (!builder.BuildSequence(expr, grammar, start_state, &end_states)) { - return std::nullopt; - } + builder.BuildSequence(expr, grammar, start_state, &end_states); return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); } -bool GrammarFSMBuilderImpl::BuildChoices( +void GrammarFSMBuilderImpl::BuildChoices( const GrammarExpr& expr, const Grammar& grammar, int start_state, @@ -1663,28 +1732,28 @@ bool GrammarFSMBuilderImpl::BuildChoices( XGRAMMAR_DCHECK(start_state >= 0 && start_state < target_fsm_.NumStates()); end_states->clear(); - int sequence_count = 0; + int non_empty_choice_count = 0; bool nullable = false; for (int32_t choice_id : expr) { const auto& choice_expr = grammar->GetGrammarExpr(choice_id); if (choice_expr.type == ExprType::kEmptyStr) { nullable = true; } else { - XGRAMMAR_DCHECK(choice_expr.type == ExprType::kSequence); - ++sequence_count; + ++non_empty_choice_count; } } - if (sequence_count == 0) { + if (non_empty_choice_count == 0) { end_states->push_back(start_state); - return true; + return; } - if (sequence_count == 1 && !nullable) { + if (non_empty_choice_count == 1 && !nullable) { for (int32_t choice_id : expr) { const auto& choice_expr = grammar->GetGrammarExpr(choice_id); if (choice_expr.type != ExprType::kEmptyStr) { - return BuildSequence(choice_expr, grammar, start_state, end_states); + BuildExpression(choice_expr, grammar, start_state, end_states); + return; } } XGRAMMAR_UNREACHABLE(); @@ -1697,9 +1766,7 @@ bool GrammarFSMBuilderImpl::BuildChoices( continue; } int branch_start_state = target_fsm_.AddState(); - if (!BuildSequence(choice_expr, grammar, branch_start_state, &branch_end_states)) { - return false; - } + BuildExpression(choice_expr, grammar, branch_start_state, &branch_end_states); target_fsm_.AddEpsilonEdge(start_state, branch_start_state); end_states->insert(end_states->end(), branch_end_states.begin(), branch_end_states.end()); } @@ -1709,26 +1776,112 @@ bool GrammarFSMBuilderImpl::BuildChoices( target_fsm_.AddEpsilonEdge(start_state, nullable_branch_state); end_states->push_back(nullable_branch_state); } - return true; } std::optional GrammarFSMBuilderImpl::Choices( const GrammarExpr& expr, const Grammar& grammar ) { - FSM result_fsm; - int start_state = result_fsm.AddState(); - std::vector end_states; - GrammarFSMBuilderImpl builder(result_fsm); - if (!builder.BuildChoices(expr, grammar, start_state, &end_states)) { - return std::nullopt; + return BuildExpressionFSM(expr, grammar); +} + +void GrammarFSMBuilderImpl::AppendFSM( + FSMWithStartEnd fsm, int start_state, std::vector* end_states +) { + const bool target_is_empty = target_fsm_.NumStates() == 1 && start_state == 0 && + target_fsm_.GetEdges(0).empty() && + target_fsm_.GetEdgeAuxData().empty() && fsm.GetStart() == 0; + if (target_is_empty) { + *end_states = fsm.GetEnds(); + target_fsm_ = std::move(fsm.GetFsm()); + return; } - FSMWithStartEnd result(result_fsm, start_state, std::move(end_states)); - result = result.SimplifyEpsilon(); - result = result.MergeEquivalentStates(); - return result; + + std::vector state_mapping; + target_fsm_.AddFSM(fsm.GetFsm(), &state_mapping); + target_fsm_.AddEpsilonEdge(start_state, state_mapping[fsm.GetStart()]); + end_states->clear(); + end_states->reserve(fsm.GetEnds().size()); + for (int end_state : fsm.GetEnds()) { + end_states->push_back(state_mapping[end_state]); + } +} + +void GrammarFSMBuilderImpl::BuildRegex( + const std::string& regex, bool json_string, int start_state, std::vector* end_states +) { + auto build_result = json_string ? RegexFSMBuilder::BuildWithForbiddenChars( + regex, GrammarFSMBuilder::JSONStringForbiddenChars() + ) + : RegexFSMBuilder::Build(regex); + if (build_result.IsErr()) { + auto error = std::move(build_result).UnwrapErr(); + if (rule_name_ != nullptr) { + XGRAMMAR_LOG(FATAL) << "Failed to build the automaton for rule " << *rule_name_ + << " with regex " << regex << ": " << error.what(); + } + XGRAMMAR_LOG(FATAL) << "Failed to build the automaton for regex " << regex << ": " + << error.what(); + } + AppendFSM(std::move(build_result).Unwrap(), start_state, end_states); +} + +void GrammarFSMBuilderImpl::BuildTagDispatch( + const Grammar::Impl::TagDispatch& tag_dispatch, + int start_state, + std::vector* end_states +) { + auto build_result = TagDispatch(tag_dispatch); + XGRAMMAR_CHECK(build_result.has_value()) << "Failed to build tag dispatch FSM"; + AppendFSM(std::move(*build_result), start_state, end_states); +} + +void GrammarFSMBuilderImpl::BuildTokenTagDispatch( + const Grammar::Impl::TokenTagDispatch& token_tag_dispatch, + int start_state, + std::vector* end_states +) { + int trigger_count = static_cast(token_tag_dispatch.trigger_rule_pairs.size()); + std::vector dispatch_states; + dispatch_states.reserve(trigger_count); + for (int index = 0; index < trigger_count; ++index) { + dispatch_states.push_back(target_fsm_.AddState()); + } + + int end_state = -1; + end_states->clear(); + end_states->push_back(start_state); + if (!token_tag_dispatch.loop_after_dispatch) { + end_state = target_fsm_.AddState(); + end_states->push_back(end_state); + } + + std::vector excluded_token_ids; + excluded_token_ids.reserve( + token_tag_dispatch.trigger_rule_pairs.size() + token_tag_dispatch.excludes.size() + ); + for (const auto& trigger_rule_pair : token_tag_dispatch.trigger_rule_pairs) { + excluded_token_ids.push_back(trigger_rule_pair.first); + } + excluded_token_ids.insert( + excluded_token_ids.end(), + token_tag_dispatch.excludes.begin(), + token_tag_dispatch.excludes.end() + ); + std::sort(excluded_token_ids.begin(), excluded_token_ids.end()); + excluded_token_ids.erase( + std::unique(excluded_token_ids.begin(), excluded_token_ids.end()), excluded_token_ids.end() + ); + + for (int index = 0; index < trigger_count; ++index) { + auto [token_id, rule_id] = token_tag_dispatch.trigger_rule_pairs[index]; + int dispatch_target = token_tag_dispatch.loop_after_dispatch ? start_state : end_state; + target_fsm_.AddTokenEdge(start_state, dispatch_states[index], {token_id}); + target_fsm_.AddRuleEdge(dispatch_states[index], dispatch_target, rule_id); + } + target_fsm_.AddExcludeTokenEdge(start_state, start_state, excluded_token_ids); } -std::optional GrammarFSMBuilderImpl::BuildTagDispatch( +std::optional GrammarFSMBuilderImpl::BuildTagDispatchFSM( const std::vector>& string_trigger_rules, bool loop_after_dispatch, const std::vector& excluded_strings @@ -1776,7 +1929,7 @@ std::optional GrammarFSMBuilderImpl::TagDispatch( tag_dispatch.tag_rule_pairs.begin(), tag_dispatch.tag_rule_pairs.end() ); - return BuildTagDispatch( + return BuildTagDispatchFSM( string_trigger_rules, tag_dispatch.loop_after_dispatch, tag_dispatch.excludes ); } diff --git a/cpp/grammar_impl.h b/cpp/grammar_impl.h index c9c919f55..a1c570662 100644 --- a/cpp/grammar_impl.h +++ b/cpp/grammar_impl.h @@ -233,7 +233,7 @@ class Grammar::Impl { }; /*! \brief Get the tag dispatch from the grammar expr. */ - TagDispatch GetTagDispatch(const GrammarExpr& grammar_expr) { + TagDispatch GetTagDispatch(const GrammarExpr& grammar_expr) const { XGRAMMAR_DCHECK(grammar_expr.type == GrammarExprType::kTagDispatch) << "GrammarExpr is not a tag dispatch"; @@ -265,7 +265,7 @@ class Grammar::Impl { } /*! \brief Get the tag dispatch from the grammar expr with the given id. */ - TagDispatch GetTagDispatch(int32_t grammar_expr_id) { + TagDispatch GetTagDispatch(int32_t grammar_expr_id) const { return GetTagDispatch(GetGrammarExpr(grammar_expr_id)); } @@ -277,7 +277,7 @@ class Grammar::Impl { }; /*! \brief Decode a kTokenTagDispatch expr into the TokenTagDispatch struct. */ - TokenTagDispatch GetTokenTagDispatch(const GrammarExpr& grammar_expr) { + TokenTagDispatch GetTokenTagDispatch(const GrammarExpr& grammar_expr) const { XGRAMMAR_DCHECK(grammar_expr.type == GrammarExprType::kTokenTagDispatch); TokenTagDispatch result; int pos = 0; @@ -297,7 +297,7 @@ class Grammar::Impl { } /*! \brief Get the token tag dispatch from the grammar expr with the given id. */ - TokenTagDispatch GetTokenTagDispatch(int32_t grammar_expr_id) { + TokenTagDispatch GetTokenTagDispatch(int32_t grammar_expr_id) const { return GetTokenTagDispatch(GetGrammarExpr(grammar_expr_id)); } diff --git a/cpp/tvm_ffi/tvm_ffi.cc b/cpp/tvm_ffi/tvm_ffi.cc index ed3a001bc..5f43e0c00 100644 --- a/cpp/tvm_ffi/tvm_ffi.cc +++ b/cpp/tvm_ffi/tvm_ffi.cc @@ -863,6 +863,18 @@ TVM_FFI_STATIC_INIT_BLOCK() { return ffi::String(_PrintGrammarFSMs(grammar_ref.as()->value)); } ) + .def( + "xgrammar.tvm_ffi_binding.testing._is_rule_fsm_accept_string", + [](O grammar_ref, int64_t rule_id, ffi::String input) { + const auto& grammar = grammar_ref.as()->value; + XGRAMMAR_CHECK(rule_id >= 0 && rule_id < grammar->NumRules()) + << "Rule id is out of range: " << rule_id; + const auto& rule_fsm = grammar->per_rule_fsms[rule_id]; + XGRAMMAR_CHECK(rule_fsm.has_value()) + << "Rule " << rule_id << " does not have a finite-state machine"; + return rule_fsm->GetFsm().AcceptString(input); + } + ) .def( "xgrammar.tvm_ffi_binding.testing.grammar_functor.structure_normalizer", [](O grammar_ref) { @@ -911,6 +923,14 @@ TVM_FFI_STATIC_INIT_BLOCK() { )); } ) + .def( + "xgrammar.tvm_ffi_binding.testing.grammar_functor.fsm_builder", + [](O grammar_ref) { + Grammar grammar = grammar_ref.as()->value; + GrammarFSMBuilder::Apply(&grammar); + return ffi::ObjectRef(ffi::make_object(std::move(grammar))); + } + ) .def( "xgrammar.tvm_ffi_binding.testing.grammar_functor.repetition_normalizer", [](O grammar_ref) { diff --git a/python/xgrammar/testing.py b/python/xgrammar/testing.py index c4c5bb7e0..e6e8d9ab8 100644 --- a/python/xgrammar/testing.py +++ b/python/xgrammar/testing.py @@ -192,6 +192,11 @@ def _is_grammar_accept_string( return grammar_matcher.is_terminated() +def _is_rule_fsm_accept_string(grammar: Grammar, rule_id: int, input_str: str) -> bool: + """Check whether a rule's already-built FSM accepts a string.""" + return bool(_core.testing._is_rule_fsm_accept_string(grammar._handle, rule_id, input_str)) + + def _get_masked_tokens_from_bitmask( bitmask: torch.Tensor, vocab_size: int, index: int = 0 ) -> List[int]: @@ -433,6 +438,13 @@ def grammar_optimizer(grammar: Grammar) -> Grammar: _core.testing.grammar_functor.grammar_optimizer(grammar._handle) ) + @staticmethod + def fsm_builder(grammar: Grammar) -> Grammar: + """Build rule FSMs without running other grammar passes.""" + return Grammar._create_from_handle( + _core.testing.grammar_functor.fsm_builder(grammar._handle) + ) + @staticmethod def repetition_normalizer(grammar: Grammar) -> Grammar: """Normalize the repetition expression.""" diff --git a/tests/python/test_fsm_sequence_build.py b/tests/python/test_fsm_sequence_build.py index 8d5891420..26ec3d10b 100644 --- a/tests/python/test_fsm_sequence_build.py +++ b/tests/python/test_fsm_sequence_build.py @@ -15,8 +15,12 @@ import xgrammar as xgr from xgrammar.testing import ( + GrammarFunctor, + _ebnf_to_grammar_no_normalization, _get_masked_tokens_from_bitmask, _get_matcher_from_grammar_and_tokenizer_info, + _is_rule_fsm_accept_string, + _print_grammar_fsms, ) @@ -78,6 +82,42 @@ def test_sequence_matches_reference_regex( assert checked > 1 +# --- Unnormalized nested expressions --- + + +def test_build_arbitrarily_nested_expressions_without_normalization(): + grammar = _ebnf_to_grammar_no_normalization( + """ +root ::= (("a" | "bb") ("c" | ("d" ("e" | "f")))) | (Regex("[0-9]+") "z") +""" + ) + grammar = GrammarFunctor.fsm_builder(grammar) + + for accepted in ("ac", "bbc", "ade", "adf", "bbde", "bbdf", "0z", "123z"): + assert _is_rule_fsm_accept_string(grammar, 0, accepted) + for rejected in ("", "a", "bc", "ad", "z", "12", "1zz", "ace"): + assert not _is_rule_fsm_accept_string(grammar, 0, rejected) + + +@pytest.mark.parametrize( + "grammar_str", + [ + """ +root ::= "a" TagDispatch(("tag", body), loop_after_dispatch=false) "z" +body ::= "b" +""", + """ +root ::= Token(2) TokenTagDispatch((3, body), excludes=(5,)) Token(6) +body ::= Token(4) +""", + ], +) +def test_build_nested_dispatch_without_normalization(grammar_str: str): + grammar = _ebnf_to_grammar_no_normalization(grammar_str) + grammar = GrammarFunctor.fsm_builder(grammar) + assert "None" not in _print_grammar_fsms(grammar) + + # --- UTF-8 multi-byte content in sequences ---