[Refac] Refactor the pipeline. (#428)

This is a rebased version of #383.
This PR refactors the pipeline. The functors in `grammar_functor` are
divided into three types:
- `grammar_normalizer`: When a grammar is constructed, they should be
called immediately.
- `grammar_optimizer`: When a grammar is going to be compiled, they will
be called.
- `grammar_constructor`: They are used to construct a new grammar, like
constructing the union of two grammars.

---------

Signed-off-by: Yuchuan <[email protected]>

Author: Seven-Streams

cpp/earley_parser.cc

@@ -241,6 +241,10 @@ EarleyParser::EarleyParser(
     const Grammar& grammar, const ParserState& init_state, const bool need_expand
 )
     : grammar_(grammar) {
+  if (!grammar->optimized) {
+    XGRAMMAR_LOG(FATAL) << "The grammar is not optimized. Please optimize the grammar before using "
+                           "the Earley parser.";
+  }
   // Check if the initial state is valid. If invalid, then we choose the root state as default.
   ParserState init = init_state;
   if (init_state.IsInvalid()) {

cpp/grammar_compiler.cc

@@ -645,16 +645,13 @@ class GrammarCompilerNoCache {
   std::unordered_map<int32_t, DynamicBitset> tag_dispatch_rule_id_to_second_slicing_bitset;
 };
 
-CompiledGrammar GrammarCompilerNoCache::MultiThreadCompileGrammar(Grammar grammar) {
+CompiledGrammar GrammarCompilerNoCache::MultiThreadCompileGrammar(Grammar grammar_unoptimized) {
   using GrammarExprType = Grammar::Impl::GrammarExprType;
 
   auto compiled_grammar_impl = std::make_shared<CompiledGrammar::Impl>();
 
-  compiled_grammar_impl->grammar = grammar;
+  compiled_grammar_impl->grammar = GrammarOptimizer::Apply(grammar_unoptimized);
   compiled_grammar_impl->tokenizer_info = tokenizer_info_;
-  grammar->allow_empty_rule_ids = AllowEmptyRuleAnalyzer::Apply(compiled_grammar_impl->grammar);
-  RepetitionNormalizer::Apply(&compiled_grammar_impl->grammar);
-  GrammarFSMBuilder::Apply(&compiled_grammar_impl->grammar);
   if (tokenizer_info_.GetVocabSize() == 0) {
     return CompiledGrammar(compiled_grammar_impl);
   }
@@ -711,15 +708,14 @@ CompiledGrammar GrammarCompilerNoCache::MultiThreadCompileGrammar(Grammar gramma
   // not need ThreadPool or std::mutex, which throws error in runtime in WebAssembly.
   std::optional<ThreadPool> thread_pool;
   std::optional<std::mutex> adaptive_token_mask_cache_mutex;
-
   if (max_threads_ > 1) {
     thread_pool.emplace(max_threads_);
     adaptive_token_mask_cache_mutex.emplace();
   }
 
   auto add_adaptive_token_mask = [&](const ParserState& state, bool is_root_rule) {
     auto grammar_matcher = GrammarMatcherForTokenMaskCache(
-        grammar, state, tag_dispatch_rule_id_to_second_slicing_bitset, false
+        compiled_grammar_impl->grammar, state, tag_dispatch_rule_id_to_second_slicing_bitset, false
     );
     auto cur_adaptive_token_mask_cache = grammar_matcher.GetAdaptiveTokenMask(
         tokenizer_info_.GetVocabSize(),
@@ -746,12 +742,13 @@ CompiledGrammar GrammarCompilerNoCache::MultiThreadCompileGrammar(Grammar gramma
     }
   };
 
-  auto root_rule_id = grammar->GetRootRuleId();
+  auto root_rule_id = compiled_grammar_impl->grammar->GetRootRuleId();
 
-  for (int32_t rule_id = 0; rule_id < static_cast<int>(grammar->NumRules()); ++rule_id) {
-    auto rule = grammar->GetRule(rule_id);
-    auto rule_body = grammar->GetGrammarExpr(rule.body_expr_id);
-    const auto& rule_fsm = grammar->per_rule_fsms[rule_id];
+  for (int32_t rule_id = 0; rule_id < static_cast<int>(compiled_grammar_impl->grammar->NumRules());
+       ++rule_id) {
+    auto rule = compiled_grammar_impl->grammar->GetRule(rule_id);
+    auto rule_body = compiled_grammar_impl->grammar->GetGrammarExpr(rule.body_expr_id);
+    const auto& rule_fsm = compiled_grammar_impl->grammar->per_rule_fsms[rule_id];
     if (rule_fsm.has_value()) {
       auto cur_stack_element =
           ParserState(rule_id, rule.body_expr_id, 0, ParserState::kNoPrevInputPos, 0);
@@ -768,15 +765,15 @@ CompiledGrammar GrammarCompilerNoCache::MultiThreadCompileGrammar(Grammar gramma
     }
     XGRAMMAR_DCHECK(rule_body.type == GrammarExprType::kChoices);
     for (auto sequence_id : rule_body) {
-      const auto& sequence = grammar->GetGrammarExpr(sequence_id);
+      const auto& sequence = compiled_grammar_impl->grammar->GetGrammarExpr(sequence_id);
       if (sequence.type == GrammarExprType::kEmptyStr) {
         continue;
       }
       XGRAMMAR_DCHECK(sequence.type == GrammarExprType::kSequence);
       auto state = ParserState(rule_id, sequence_id, 0, ParserState::kNoPrevInputPos, 0);
       for (int element_id = 0; element_id < sequence.size(); ++element_id) {
         state.element_id = element_id;
-        auto element = grammar->GetGrammarExpr(sequence[element_id]);
+        auto element = compiled_grammar_impl->grammar->GetGrammarExpr(sequence[element_id]);
         if (element.type == GrammarExprType::kRuleRef || element.type == GrammarExprType::kRepeat) {
           continue;
         }