[Optim] Optim repetition expressions when the max repetition time is unbounded. (#403)

This PR optims the behavior of repetition expression when the max
repetition time is unbounded. It keeps the invariant of the repetition
states, instead of setting the upper as a large number. This can
significantly speed lower the uncertainty when the max repetition time
is unbounded, and the min repetition time is very small, like this
schema:
```
{
    "type": "string",
    "minLength": 1
}
```

Signed-off-by: Yuchuan
[[email protected]](mailto:[email protected])

Author: Seven-Streams

cpp/earley_parser.h

@@ -156,7 +156,7 @@ class StateEqualForParsing {
   bool operator()(const ParserState& lhs, const ParserState& rhs) const {
     return lhs.rule_id == rhs.rule_id && lhs.sequence_id == rhs.sequence_id &&
            lhs.element_id == rhs.element_id && lhs.rule_start_pos == rhs.rule_start_pos &&
-           lhs.sub_element_id == rhs.sub_element_id;
+           lhs.sub_element_id == rhs.sub_element_id && lhs.repeat_count == rhs.repeat_count;
   }
 };
 

cpp/grammar_parser.cc

@@ -8,6 +8,7 @@
 #include <picojson.h>
 
 #include <cstdint>
+#include <string>
 #include <variant>
 #include <vector>
 
@@ -758,12 +759,33 @@ int32_t EBNFParser::HandleQuestionQuantifier(int32_t grammar_expr_id) {
 int32_t EBNFParser::HandleRepetitionRange(
     const int32_t grammar_expr_id, int64_t lower, int64_t upper
 ) {
+  bool is_unbounded = false;
+  int32_t new_element;
   if (upper == -1) {
     // The repeation is unbounded, e.g. {2,}
-    upper = 0x7FFFFFFF;  // Use a large number to represent unbounded
+    is_unbounded = true;
+    const auto& rule_expr = builder_.GetGrammarExpr(grammar_expr_id);
+    if (rule_expr.type == GrammarBuilder::GrammarExprType::kCharacterClass) {
+      std::vector<GrammarBuilder::CharacterClassElement> character_ranges;
+      bool is_negative = rule_expr[0];
+      for (int i = 1; i < static_cast<int>(rule_expr.size()); i += 2) {
+        character_ranges.push_back({rule_expr[i], rule_expr[i + 1]});
+      }
+      new_element = builder_.AddCharacterClassStar(character_ranges, is_negative);
+    } else {
+      const auto& unbounded_rule_id =
+          builder_.AddEmptyRule(builder_.GetNewRuleName(cur_rule_name_ + "_repeat_inf"));
+      int recursion_sequence =
+          builder_.AddSequence({grammar_expr_id, builder_.AddRuleRef(unbounded_rule_id)});
+      int recursion_choice = builder_.AddChoices({builder_.AddEmptyStr(), recursion_sequence});
+      builder_.UpdateRuleBody(unbounded_rule_id, recursion_choice);
+      new_element = builder_.AddRuleRef(unbounded_rule_id);
+    }
+    upper = lower;
   }
-  const auto repeat_name = builder_.GetNewRuleName(cur_rule_name_) + "_xgrammar_repetition_context";
   std::vector<int32_t> elements;
+  const auto repeat_name = cur_rule_name_ + "_repeat_";
+  int cnt = 1;
   int splited_count = lower >= 4 ? 4 : lower;
   int nullable_splited_count = 0;
   if (splited_count != 4) {
@@ -772,9 +794,9 @@ int32_t EBNFParser::HandleRepetitionRange(
   }
   // The repetition sentence.
   if (upper != (splited_count + nullable_splited_count)) {
-    auto new_rule_name = builder_.GetNewRuleName(repeat_name);
     auto new_grammar_expr_id = builder_.AddChoices({builder_.AddSequence({grammar_expr_id})});
-    auto new_rule_id = builder_.AddRule(new_rule_name, new_grammar_expr_id);
+    auto new_rule_id =
+        builder_.AddRuleWithHint(repeat_name + std::to_string(cnt++), new_grammar_expr_id);
     elements.push_back(builder_.AddRepeat(
         new_rule_id, lower - splited_count, upper - splited_count - nullable_splited_count
     ));
@@ -783,20 +805,22 @@ int32_t EBNFParser::HandleRepetitionRange(
 
   // The nullable exprs.
   for (int i = 0; i < nullable_splited_count; i++) {
-    auto new_rule_name = builder_.GetNewRuleName(repeat_name);
     auto new_grammar_expr_id =
         builder_.AddChoices({builder_.AddEmptyStr(), builder_.AddSequence({grammar_expr_id})});
-    auto new_rule_id = builder_.AddRule(new_rule_name, new_grammar_expr_id);
+    auto new_rule_id =
+        builder_.AddRuleWithHint(repeat_name + std::to_string(cnt++), new_grammar_expr_id);
     elements.push_back(builder_.AddRuleRef(new_rule_id));
   }
 
   for (int i = 0; i < splited_count; i++) {
-    auto new_rule_name = builder_.GetNewRuleName(repeat_name);
     auto new_grammar_expr_id = builder_.AddChoices({builder_.AddSequence({grammar_expr_id})});
-    auto new_rule_id = builder_.AddRule(new_rule_name, new_grammar_expr_id);
+    auto new_rule_id =
+        builder_.AddRuleWithHint(repeat_name + std::to_string(cnt++), new_grammar_expr_id);
     elements.push_back(builder_.AddRuleRef(new_rule_id));
   }
-
+  if (is_unbounded) {
+    elements.push_back(new_element);
+  }
   // Add the lookahead elements
   std::vector<int32_t> lookahead_elements = elements;
   if (elements.empty()) {

tests/python/test_grammar_parser.py

@@ -146,10 +146,10 @@ def test_repetition_range_exact():
     """Test repetition range with exact count {n}."""
     before = """root ::= "a"{3}
 """
-    expected = """root ::= (((root_1_xgrammar_repetition_context root_1_xgrammar_repetition_context_1 root_1_xgrammar_repetition_context_2)))
-root_1_xgrammar_repetition_context ::= (("a")) (=(root_1_xgrammar_repetition_context_1 root_1_xgrammar_repetition_context_2))
-root_1_xgrammar_repetition_context_1 ::= (("a")) (=(root_1_xgrammar_repetition_context_2))
-root_1_xgrammar_repetition_context_2 ::= (("a"))
+    expected = """root ::= (((root_repeat_1 root_repeat_2 root_repeat_3)))
+root_repeat_1 ::= (("a")) (=(root_repeat_2 root_repeat_3))
+root_repeat_2 ::= (("a")) (=(root_repeat_3))
+root_repeat_3 ::= (("a"))
 """
     grammar = _ebnf_to_grammar_no_normalization(before)
     after = str(grammar)
@@ -160,11 +160,11 @@ def test_repetition_range_min_max():
     """Test repetition range with min and max {n,m}."""
     before = """root ::= "a"{2,4}
 """
-    expected = """root ::= (((root_1_xgrammar_repetition_context root_1_xgrammar_repetition_context_1 root_1_xgrammar_repetition_context_2 root_1_xgrammar_repetition_context_3)))
-root_1_xgrammar_repetition_context ::= ("" | ("a")) (=(root_1_xgrammar_repetition_context_1 root_1_xgrammar_repetition_context_2 root_1_xgrammar_repetition_context_3))
-root_1_xgrammar_repetition_context_1 ::= ("" | ("a")) (=(root_1_xgrammar_repetition_context_2 root_1_xgrammar_repetition_context_3))
-root_1_xgrammar_repetition_context_2 ::= (("a")) (=(root_1_xgrammar_repetition_context_3))
-root_1_xgrammar_repetition_context_3 ::= (("a"))
+    expected = """root ::= (((root_repeat_1 root_repeat_2 root_repeat_3 root_repeat_4)))
+root_repeat_1 ::= ("" | ("a")) (=(root_repeat_2 root_repeat_3 root_repeat_4))
+root_repeat_2 ::= ("" | ("a")) (=(root_repeat_3 root_repeat_4))
+root_repeat_3 ::= (("a")) (=(root_repeat_4))
+root_repeat_4 ::= (("a"))
 """
     grammar = _ebnf_to_grammar_no_normalization(before)
     after = str(grammar)
@@ -175,12 +175,10 @@ def test_repetition_range_min_only():
     """Test repetition range with only min {n,}."""
     before = """root ::= "a"{2,}
 """
-    expected = """root ::= (((root_1_xgrammar_repetition_context{0, 2147483643} root_1_xgrammar_repetition_context_1 root_1_xgrammar_repetition_context_2 root_1_xgrammar_repetition_context_3 root_1_xgrammar_repetition_context_4)))
-root_1_xgrammar_repetition_context ::= (("a")) (=(root_1_xgrammar_repetition_context_1 root_1_xgrammar_repetition_context_2 root_1_xgrammar_repetition_context_3 root_1_xgrammar_repetition_context_4))
-root_1_xgrammar_repetition_context_1 ::= ("" | ("a")) (=(root_1_xgrammar_repetition_context_2 root_1_xgrammar_repetition_context_3 root_1_xgrammar_repetition_context_4))
-root_1_xgrammar_repetition_context_2 ::= ("" | ("a")) (=(root_1_xgrammar_repetition_context_3 root_1_xgrammar_repetition_context_4))
-root_1_xgrammar_repetition_context_3 ::= (("a")) (=(root_1_xgrammar_repetition_context_4))
-root_1_xgrammar_repetition_context_4 ::= (("a"))
+    expected = """root ::= (((root_repeat_1 root_repeat_2 root_repeat_inf)))
+root_repeat_inf ::= ("" | ("a" root_repeat_inf))
+root_repeat_1 ::= (("a")) (=(root_repeat_2 root_repeat_inf))
+root_repeat_2 ::= (("a")) (=(root_repeat_inf))
 """
     grammar = _ebnf_to_grammar_no_normalization(before)
     after = str(grammar)
@@ -275,12 +273,12 @@ def test_combined_features():
 rule2 ::= [0-9]+ "." [0-9]*
 """
     expected = """root ::= (("start" root_1 "end"))
-rule1 ::= (((rule1_1_xgrammar_repetition_context rule1_1_xgrammar_repetition_context_1 rule1_1_xgrammar_repetition_context_2))) (=((":")))
+rule1 ::= (((rule1_repeat_1 rule1_repeat_2 rule1_repeat_3))) (=((":")))
 rule2 ::= ((rule2_1 "." [0-9]*))
 root_1 ::= ((((rule1) | (rule2)) root_1) | ((rule1) | (rule2)))
-rule1_1_xgrammar_repetition_context ::= ("" | ([a-z])) (=(rule1_1_xgrammar_repetition_context_1 rule1_1_xgrammar_repetition_context_2))
-rule1_1_xgrammar_repetition_context_1 ::= ("" | ([a-z])) (=(rule1_1_xgrammar_repetition_context_2))
-rule1_1_xgrammar_repetition_context_2 ::= (([a-z]))
+rule1_repeat_1 ::= ("" | ([a-z])) (=(rule1_repeat_2 rule1_repeat_3))
+rule1_repeat_2 ::= ("" | ([a-z])) (=(rule1_repeat_3))
+rule1_repeat_3 ::= (([a-z]))
 rule2_1 ::= (([0-9] rule2_1) | [0-9])
 """
     grammar = _ebnf_to_grammar_no_normalization(before)
@@ -354,35 +352,29 @@ def test_repetition_range():
 """
 
     expected = """root ::= ((a b c d e f g))
-a ::= ((a_1_xgrammar_repetition_context a_1_xgrammar_repetition_context_1))
-b ::= ((b_1_xgrammar_repetition_context{0, 1} b_1_xgrammar_repetition_context_1 b_1_xgrammar_repetition_context_2 b_1_xgrammar_repetition_context_3 b_1_xgrammar_repetition_context_4))
-c ::= ((c_1_xgrammar_repetition_context c_1_xgrammar_repetition_context_1))
-d ::= ((d_1_xgrammar_repetition_context{0, 2147483643} d_1_xgrammar_repetition_context_1 d_1_xgrammar_repetition_context_2 d_1_xgrammar_repetition_context_3 d_1_xgrammar_repetition_context_4))
-e ::= ((e_1_xgrammar_repetition_context{0, 2147483643} e_1_xgrammar_repetition_context_1 e_1_xgrammar_repetition_context_2 e_1_xgrammar_repetition_context_3 e_1_xgrammar_repetition_context_4))
-f ::= ((f_1_xgrammar_repetition_context f_1_xgrammar_repetition_context_1 f_1_xgrammar_repetition_context_2))
+a ::= ((a_repeat_1 a_repeat_2))
+b ::= ((b_repeat_1{0, 1} b_repeat_2 b_repeat_3 b_repeat_4 b_repeat_5))
+c ::= ((c_repeat_1 c_repeat_2))
+d ::= ((d_repeat_inf))
+e ::= ((e_repeat_1 e_repeat_2 e_repeat_inf))
+f ::= ((f_repeat_1 f_repeat_2 f_repeat_3))
 g ::= ("")
-a_1_xgrammar_repetition_context ::= ("" | ("a")) (=(a_1_xgrammar_repetition_context_1))
-a_1_xgrammar_repetition_context_1 ::= (("a"))
-b_1_xgrammar_repetition_context ::= ((a) | ("b")) (=(b_1_xgrammar_repetition_context_1 b_1_xgrammar_repetition_context_2 b_1_xgrammar_repetition_context_3 b_1_xgrammar_repetition_context_4))
-b_1_xgrammar_repetition_context_1 ::= ("" | (a) | ("b")) (=(b_1_xgrammar_repetition_context_2 b_1_xgrammar_repetition_context_3 b_1_xgrammar_repetition_context_4))
-b_1_xgrammar_repetition_context_2 ::= ("" | (a) | ("b")) (=(b_1_xgrammar_repetition_context_3 b_1_xgrammar_repetition_context_4))
-b_1_xgrammar_repetition_context_3 ::= ("" | (a) | ("b")) (=(b_1_xgrammar_repetition_context_4))
-b_1_xgrammar_repetition_context_4 ::= ((a) | ("b"))
-c_1_xgrammar_repetition_context ::= ("" | ("c")) (=(c_1_xgrammar_repetition_context_1))
-c_1_xgrammar_repetition_context_1 ::= ("" | ("c"))
-d_1_xgrammar_repetition_context ::= (("d")) (=(d_1_xgrammar_repetition_context_1 d_1_xgrammar_repetition_context_2 d_1_xgrammar_repetition_context_3 d_1_xgrammar_repetition_context_4))
-d_1_xgrammar_repetition_context_1 ::= ("" | ("d")) (=(d_1_xgrammar_repetition_context_2 d_1_xgrammar_repetition_context_3 d_1_xgrammar_repetition_context_4))
-d_1_xgrammar_repetition_context_2 ::= ("" | ("d")) (=(d_1_xgrammar_repetition_context_3 d_1_xgrammar_repetition_context_4))
-d_1_xgrammar_repetition_context_3 ::= ("" | ("d")) (=(d_1_xgrammar_repetition_context_4))
-d_1_xgrammar_repetition_context_4 ::= ("" | ("d"))
-e_1_xgrammar_repetition_context ::= (("e")) (=(e_1_xgrammar_repetition_context_1 e_1_xgrammar_repetition_context_2 e_1_xgrammar_repetition_context_3 e_1_xgrammar_repetition_context_4))
-e_1_xgrammar_repetition_context_1 ::= ("" | ("e")) (=(e_1_xgrammar_repetition_context_2 e_1_xgrammar_repetition_context_3 e_1_xgrammar_repetition_context_4))
-e_1_xgrammar_repetition_context_2 ::= ("" | ("e")) (=(e_1_xgrammar_repetition_context_3 e_1_xgrammar_repetition_context_4))
-e_1_xgrammar_repetition_context_3 ::= (("e")) (=(e_1_xgrammar_repetition_context_4))
-e_1_xgrammar_repetition_context_4 ::= (("e"))
-f_1_xgrammar_repetition_context ::= (("f")) (=(f_1_xgrammar_repetition_context_1 f_1_xgrammar_repetition_context_2))
-f_1_xgrammar_repetition_context_1 ::= (("f")) (=(f_1_xgrammar_repetition_context_2))
-f_1_xgrammar_repetition_context_2 ::= (("f"))
+a_repeat_1 ::= ("" | ("a")) (=(a_repeat_2))
+a_repeat_2 ::= (("a"))
+b_repeat_1 ::= ((a) | ("b")) (=(b_repeat_2 b_repeat_3 b_repeat_4 b_repeat_5))
+b_repeat_2 ::= ("" | (a) | ("b")) (=(b_repeat_3 b_repeat_4 b_repeat_5))
+b_repeat_3 ::= ("" | (a) | ("b")) (=(b_repeat_4 b_repeat_5))
+b_repeat_4 ::= ("" | (a) | ("b")) (=(b_repeat_5))
+b_repeat_5 ::= ((a) | ("b"))
+c_repeat_1 ::= ("" | ("c")) (=(c_repeat_2))
+c_repeat_2 ::= ("" | ("c"))
+d_repeat_inf ::= ("" | ("d" d_repeat_inf))
+e_repeat_inf ::= ("" | ("e" e_repeat_inf))
+e_repeat_1 ::= (("e")) (=(e_repeat_2 e_repeat_inf))
+e_repeat_2 ::= (("e")) (=(e_repeat_inf))
+f_repeat_1 ::= (("f")) (=(f_repeat_2 f_repeat_3))
+f_repeat_2 ::= (("f")) (=(f_repeat_3))
+f_repeat_3 ::= (("f"))
 """
 
     grammar = _ebnf_to_grammar_no_normalization(before)