diff --git a/cpp/grammar_functor.cc b/cpp/grammar_functor.cc index eae30fe71..20cdd1af6 100644 --- a/cpp/grammar_functor.cc +++ b/cpp/grammar_functor.cc @@ -1089,41 +1089,17 @@ class GrammarFSMBuilderImpl { const static uint32_t kMin4BytesUnicode = 0xF0808080; const static uint32_t kMax4BytesUnicode = 0xF7BFBFBF; - void Apply(Grammar* grammar) { + 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; 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) { @@ -1158,7 +1134,6 @@ 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 @@ -1167,21 +1142,75 @@ class GrammarFSMBuilderImpl { 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( + static std::optional BuildTagDispatchFSM( const std::vector>& string_trigger_rules, bool loop_after_dispatch, const std::vector& excluded_strings ); - static FSMWithStartEnd BuildNegativeCharacterClass(const GrammarExpr& expr); + + private: + static FSMWithStartEnd BuildRuleFSM(const Grammar& grammar, int rule_id); + static FSMWithStartEnd BuildExpressionFSM( + const GrammarExpr& expr, const Grammar& grammar, const std::string* rule_name = nullptr + ); + void BuildExpression( + const GrammarExpr& expr, + const Grammar& grammar, + int start_state, + std::vector* end_states + ); + 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 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 // 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), @@ -1197,7 +1226,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; } @@ -1205,7 +1234,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); @@ -1213,14 +1242,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]++; @@ -1228,15 +1257,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; } @@ -1244,7 +1273,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); @@ -1252,14 +1281,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]++; @@ -1267,12 +1296,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; } @@ -1280,7 +1309,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); @@ -1288,14 +1317,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]++; @@ -1303,16 +1332,14 @@ 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_CHECK(min <= max) << "Invalid character range: min (" << min << ") > max (" << max << ")"; // Ensure max and min are valid unicode value. @@ -1346,51 +1373,53 @@ 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 ); @@ -1409,18 +1438,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]) { @@ -1429,7 +1446,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), @@ -1438,225 +1455,433 @@ 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::AddCharacterClassTransitions( + const GrammarExpr& expr, int start_state, int end_state +) { + XGRAMMAR_DCHECK( + expr.type == ExprType::kCharacterClass || expr.type == ExprType::kCharacterClassStar + ); bool is_negative = expr[0]; - FSMWithStartEnd result_fsm; if (is_negative) { - result_fsm = BuildNegativeCharacterClass(expr); - return result_fsm; + BuildNegativeCharacterClass(expr, start_state, end_state); + } else { + 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); + } } +} + +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); +} + +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::BuildExpressionFSM( + const GrammarExpr& expr, const Grammar& grammar, const std::string* rule_name +) { + FSM 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; - } else { - end_state = result_fsm.AddState(); + std::vector 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(); } - 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; +} + +FSMWithStartEnd GrammarFSMBuilderImpl::RuleRef(const GrammarExpr& 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) { + 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 result_fsm; + return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); } -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::ByteString(const GrammarExpr& 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) { - 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}); + 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) { - 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}); + 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)); } -std::optional GrammarFSMBuilderImpl::Sequence( - const GrammarExpr& expr, const Grammar& grammar +void GrammarFSMBuilderImpl::BuildExpression( + const GrammarExpr& expr, + const Grammar& grammar, + int start_state, + std::vector* end_states ) { - 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) { - case (ExprType::kByteString): { - fsm_lists.push_back(ByteString(sequence_expr)); - break; - } - case (ExprType::kRuleRef): { - fsm_lists.push_back(RuleRef(sequence_expr)); - break; - } - case (ExprType::kCharacterClass): - case (ExprType::kCharacterClassStar): { - fsm_lists.push_back(CharacterClass(sequence_expr)); - break; - } - case (ExprType::kRepeat): { - fsm_lists.push_back(Repeat(sequence_expr)); - break; - } - case (ExprType::kToken): { - fsm_lists.push_back(Token(sequence_expr)); - break; - } - case (ExprType::kExcludeToken): { - fsm_lists.push_back(ExcludeToken(sequence_expr)); - break; - } - default: { - return std::nullopt; - } - } - } - - // Check if the sequence is empty. - if (fsm_lists.empty()) { - FSMWithStartEnd empty_fsm; - empty_fsm.AddState(); - empty_fsm.SetStartState(0); - empty_fsm.AddEndState(0); - return empty_fsm; + XGRAMMAR_DCHECK(start_state >= 0 && start_state < target_fsm_.NumStates()); + 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 + ); + case ExprType::kTagDispatch: + return BuildTagDispatch(grammar->GetTagDispatch(expr), start_state, end_states); + case ExprType::kTokenTagDispatch: + return BuildTokenTagDispatch(grammar->GetTokenTagDispatch(expr), start_state, end_states); } - - return FSMWithStartEnd::Concat(fsm_lists); + XGRAMMAR_UNREACHABLE(); } -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; +void GrammarFSMBuilderImpl::BuildEmptyString(int start_state, std::vector* end_states) { + end_states->clear(); + end_states->push_back(start_state); } -FSMWithStartEnd GrammarFSMBuilderImpl::ByteString(const GrammarExpr& expr) { +void GrammarFSMBuilderImpl::BuildByteString( + const GrammarExpr& expr, int start_state, std::vector* end_states +) { 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( + 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; } - result_fsm.AddEndState(current_state); - return result_fsm; + end_states->push_back(current_state); } -std::optional GrammarFSMBuilderImpl::Choices( +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, + std::vector* end_states +) { + XGRAMMAR_DCHECK(start_state >= 0 && start_state < target_fsm_.NumStates()); + end_states->clear(); + if (expr.size() == 0) { + end_states->push_back(start_state); + return; + } + + 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(); + 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); + } +} + +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); + builder.BuildSequence(expr, grammar, start_state, &end_states); + return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); +} + +void 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(start_state >= 0 && start_state < target_fsm_.NumStates()); + end_states->clear(); + + int non_empty_choice_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 { + ++non_empty_choice_count; } - XGRAMMAR_DCHECK(choice_expr.type == ExprType::kSequence); - auto fsm_result = Sequence(choice_expr, grammar); - if (!fsm_result.has_value()) { - return std::nullopt; + } + + if (non_empty_choice_count == 0) { + end_states->push_back(start_state); + return; + } + + 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) { + BuildExpression(choice_expr, grammar, start_state, end_states); + return; + } } - 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(); + 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()); } + 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); } +} - auto result = FSMWithStartEnd::Union(fsm_list); - result = result.SimplifyEpsilon(); - result = result.MergeEquivalentStates(); - return result; +std::optional GrammarFSMBuilderImpl::Choices( + const GrammarExpr& expr, const Grammar& grammar +) { + 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; + } + + 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 @@ -1704,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 ); } @@ -2852,7 +3077,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/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 new file mode 100644 index 000000000..26ec3d10b --- /dev/null +++ b/tests/python/test_fsm_sequence_build.py @@ -0,0 +1,284 @@ +"""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 ( + 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, +) + + +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 + + +# --- 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 --- + + +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) diff --git a/tests/python/test_fsm_structure_stability.py b/tests/python/test_fsm_structure_stability.py new file mode 100644 index 000000000..cfcf591f6 --- /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, + }, + "a7a376ee74bc9e59e36da46fd7dddc9c322555c96110120260779bd07a9d14c5", + ), + ( + "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