[Fix] Fix slow FSM optimization methods SimplifyEpsilon and MergeEquivalentSuccessors (#407)

This PR refactors part of the fsm functions. In general, this PR:
- Fix the `MergeEquivalentSuccessors` and `SimplifyEpsilon` to speed up.
- Set proper limitations of number of states for the fsm functions.

---------

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

Author: Seven-Streams

cpp/fsm.cc

@@ -16,11 +16,11 @@
 #include <cstdint>
 #include <functional>
 #include <string>
-#include <unordered_map>
 #include <unordered_set>
 #include <vector>
 
 #include "support/compact_2d_array.h"
+#include "support/logging.h"
 #include "support/reflection.h"
 #include "support/utils.h"
 #include "xgrammar/exception.h"
@@ -302,9 +302,10 @@ class FSM {
    * \brief Add a whole FSM to the current FSM.
    * \param fsm The FSM to be added.
    * \param state_mapping The mapping from the state ids of the added FSM to the new ids in the
-   * current FSM. The result is cleared at the beginning.
+   * current FSM. The result is cleared at the beginning. If the fsm's state id starts from 0, use
+   * it for efficiency.
    */
-  void AddFSM(const FSM& fsm, std::unordered_map<int, int>* state_mapping = nullptr);
+  void AddFSM(const FSM& fsm, std::vector<int>* state_mapping = nullptr);
 
   /****************** FSM Construction Methods ******************/
 
@@ -319,7 +320,7 @@ class FSM {
    * \param new_num_states The new number of states.
    * \return The rebuilt FSM.
    */
-  FSM RebuildWithMapping(std::unordered_map<int, int>& state_mapping, int new_num_states) const;
+  FSM RebuildWithMapping(const std::vector<int>& state_mapping, int new_num_states) const;
 
   /*!
    * \brief Sort the edges of the FSM by their min, max and target.
@@ -483,20 +484,6 @@ class FSMWithStartEndBase {
   // For serialization only
   FSMWithStartEndBase() = default;
 
-  /*! \brief Constructs an FSMWithStartEnd with a given FSM, start state, and end states. */
-  FSMWithStartEndBase(
-      const FSMType& fsm, int start, const std::unordered_set<int>& ends, bool is_dfa = false
-  )
-      : fsm_(fsm), start_(start), is_dfa_(is_dfa) {
-    ends_.resize(fsm.NumStates(), false);
-    for (const auto& end : ends) {
-      XGRAMMAR_DCHECK(end < fsm.NumStates())
-          << "End state " << end << " is out of bounds for FSM with " << fsm.NumStates()
-          << " states.";
-      ends_[end] = true;
-    }
-  }
-
   FSMWithStartEndBase(
       const FSMType& fsm, int start, const std::vector<bool>& ends, bool is_dfa = false
   )
@@ -653,9 +640,8 @@ class FSMWithStartEnd : public FSMWithStartEndBase<FSM> {
    * \param state_mapping The mapping from old state ids to new state ids.
    * \param new_num_states The new number of states.
    */
-  FSMWithStartEnd RebuildWithMapping(
-      std::unordered_map<int, int>& state_mapping, int new_num_states
-  );
+  FSMWithStartEnd RebuildWithMapping(const std::vector<int>& state_mapping, int new_num_states)
+      const;
 
   /*!
    * \brief Add the underlying FSM to another complete FSM that could contain multiple FSMs.
@@ -666,7 +652,7 @@ class FSMWithStartEnd : public FSMWithStartEndBase<FSM> {
    * cleared at the beginning. Should not be nullptr.
    * \return The FSMWithStartEnd that points to the complete FSM.
    */
-  FSMWithStartEnd AddToCompleteFSM(FSM* complete_fsm, std::unordered_map<int, int>* state_mapping);
+  FSMWithStartEnd AddToCompleteFSM(FSM* complete_fsm, std::vector<int>* state_mapping);
 
   /*!
    * \brief Transform the FSMWithStartEnd to a CompactFSMWithStartEnd.
@@ -735,29 +721,29 @@ class FSMWithStartEnd : public FSMWithStartEndBase<FSM> {
    * \details If a --\epsilon--> b, and either 1) b doesn't have any other inward edges, or
    * 2) a doesn't have any other outward edges, we can merge a and b.
    */
-  FSMWithStartEnd SimplifyEpsilon() const;
+  FSMWithStartEnd SimplifyEpsilon(int max_num_states = 1e8) const;
 
   /*!
    * \brief Merge equivalent states in the FSM.
    * \details If two states are 1) pointed to by edges with the same label from the same state, and
    * 2) they are not pointed to by other edges, then we can merge them.
    * \example n0 --(c)--> n1, n0 --(c)--> n2, then we can merge n1 and n2.
    */
-  FSMWithStartEnd MergeEquivalentSuccessors() const;
+  FSMWithStartEnd MergeEquivalentSuccessors(int max_num_states = 1e5) const;
 
   /*!
    * \brief Transform the FSM to a DFA.
    * \param max_result_num_states The maximum number of states in the DFA.
    * \return The DFA.
    */
-  Result<FSMWithStartEnd> ToDFA(int max_result_num_states = 1e6) const;
+  Result<FSMWithStartEnd> ToDFA(int max_num_states = 1e3) const;
 
   /*!
    * \brief Minimize the DFA.
    * \param max_result_num_states The maximum number of states in the DFA.
    * \return The minimized DFA.
    */
-  Result<FSMWithStartEnd> MinimizeDFA(int max_result_num_states = 1e6) const;
+  Result<FSMWithStartEnd> MinimizeDFA(int max_num_states = 1e3) const;
 };
 
 /*!

cpp/fsm.h

@@ -192,7 +192,7 @@ Result<std::pair<int, int>> RegexIR::CheckRepeat(const std::string& regex, int&
 Result<FSMWithStartEnd> RegexIR::Build() const {
   if (states.empty()) {
     FSM empty_fsm(1);
-    FSMWithStartEnd result(empty_fsm, 0, std::unordered_set<int>{0}, false);
+    FSMWithStartEnd result(empty_fsm, 0, {true}, false);
     return ResultOk(std::move(result));
   }
   std::vector<FSMWithStartEnd> fsm_list;
@@ -336,7 +336,7 @@ Result<FSMWithStartEnd> RegexIR::visit(const RegexIR::Repeat& state) const {
 
 FSMWithStartEnd RegexIR::BuildLeafFSMFromRegex(const std::string& regex) {
   FSM empty_fsm(0);
-  FSMWithStartEnd result(empty_fsm, 0, std::unordered_set<int>{}, true);
+  FSMWithStartEnd result(empty_fsm, 0, {}, true);
   // Handle the regex string.
   if (!(regex[0] == '[' && regex[regex.size() - 1] == ']')) {
     result.AddState();
@@ -495,7 +495,9 @@ FSMWithStartEnd RegexIR::BuildLeafFSMFromRegex(const std::string& regex) {
         new_fsm.AddEdge(0, 1, last, 0xFF);
       }
     }
-    result = FSMWithStartEnd(new_fsm, 0, std::unordered_set<int>{1}, false);
+    std::vector<bool> ends(new_fsm.NumStates(), false);
+    ends[1] = true;
+    result = FSMWithStartEnd(new_fsm, 0, ends, false);
   } else {
     // TODO: The support for rules.
     XGRAMMAR_LOG(WARNING) << "rule is not supported yet.";
@@ -830,7 +832,13 @@ std::optional<FSMWithStartEnd> TrieFSMBuilderImpl::Build(
   if (add_back_edges) {
     AddBackEdges(&fsm, start, ends);
   }
-  return FSMWithStartEnd(fsm, start, ends);
+
+  std::vector<bool> is_end_state(fsm.NumStates(), false);
+  for (const auto& end : ends) {
+    is_end_state[end] = true;
+  }
+
+  return FSMWithStartEnd(fsm, start, is_end_state);
 }
 
 void TrieFSMBuilderImpl::AddBackEdges(FSM* fsm, int start, const std::unordered_set<int>& ends) {

cpp/fsm_builder.cc

@@ -1070,10 +1070,18 @@ class StructuralTagGrammarCreatorImpl : public GrammarMutator {
 
 class GrammarFSMBuilderImpl {
  public:
+  const static uint32_t kMax1ByteUnicode = 0x7F;
+  const static uint32_t kMin2BytesUnicode = 0xC080;
+  const static uint32_t kMax2BytesUnicode = 0xDFBF;
+  const static uint32_t kMin3BytesUnicode = 0xE08080;
+  const static uint32_t kMax3BytesUnicode = 0xEFBFBF;
+  const static uint32_t kMin4BytesUnicode = 0xF0808080;
+  const static uint32_t kMax4BytesUnicode = 0xF7BFBFBF;
+
   void Apply(Grammar* grammar) {
     FSM complete_fsm;
     std::vector<std::optional<FSMWithStartEnd>> per_rule_fsms((*grammar)->NumRules());
-    std::unordered_map<int, int> state_mapping;
+    std::vector<int> state_mapping;
 
     for (int i = 0; i < (*grammar)->NumRules(); ++i) {
       auto rule = (*grammar)->GetRule(i);
@@ -1108,13 +1116,13 @@ class GrammarFSMBuilderImpl {
   }
 
   /* Basic Building functions.*/
-  static std::optional<FSMWithStartEnd> RuleRef(const GrammarExpr& expr);
-  static std::optional<FSMWithStartEnd> CharacterClass(const GrammarExpr& expr);
-  static std::optional<FSMWithStartEnd> ByteString(const GrammarExpr& expr);
+  static FSMWithStartEnd RuleRef(const GrammarExpr& expr);
+  static FSMWithStartEnd CharacterClass(const GrammarExpr& expr);
+  static FSMWithStartEnd ByteString(const GrammarExpr& expr);
   static std::optional<FSMWithStartEnd> Sequence(const GrammarExpr& expr, const Grammar& grammar);
   static std::optional<FSMWithStartEnd> Choices(const GrammarExpr& expr, const Grammar& grammar);
   static std::optional<FSMWithStartEnd> TagDispatch(const Grammar::Impl::TagDispatch& tag_dispatch);
-
+  static void AddCharacterRange(FSMWithStartEnd& fsm, int from, int to, uint32_t min, uint32_t max);
   /* Building tool funtions.*/
   static std::optional<FSMWithStartEnd> BuildTagDispatchWithEOSStop(
       const std::vector<std::pair<std::string, int>>& tag_dispatch_rules, bool loop_after_dispatch
@@ -1124,17 +1132,9 @@ class GrammarFSMBuilderImpl {
       const std::vector<std::string>& stop_strings,
       bool loop_after_dispatch
   );
-  static std::optional<FSMWithStartEnd> BuildNegativeCharacterClass(const GrammarExpr& expr);
+  static FSMWithStartEnd BuildNegativeCharacterClass(const GrammarExpr& expr);
 };
 
-const static uint32_t kMax1ByteUnicode = 0x7F;
-const static uint32_t kMin2BytesUnicode = 0xC080;
-const static uint32_t kMax2BytesUnicode = 0xDFBF;
-const static uint32_t kMin3BytesUnicode = 0xE08080;
-const static uint32_t kMax3BytesUnicode = 0xEFBFBF;
-const static uint32_t kMin4BytesUnicode = 0xF0808080;
-const static uint32_t kMax4BytesUnicode = 0xF7BFBFBF;
-
 // This function will add a range [min, max] of characters to the FSM, and the length
 // of the characters are the same.
 void AddSameLengthCharacterRange(
@@ -1268,7 +1268,9 @@ void AddSameLengthCharacterRange(
 }
 
 // This function will add a range [min, max] of unicode characters to the FSM.
-void AddCharacterRange(FSMWithStartEnd& fsm, int from, int to, uint32_t min, uint32_t max) {
+void GrammarFSMBuilderImpl::AddCharacterRange(
+    FSMWithStartEnd& fsm, 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,9 +1348,7 @@ void AddCharacterRange(FSMWithStartEnd& fsm, int from, int to, uint32_t min, uin
   return;
 }
 
-std::optional<FSMWithStartEnd> GrammarFSMBuilderImpl::BuildNegativeCharacterClass(
-    const GrammarExpr& expr
-) {
+FSMWithStartEnd GrammarFSMBuilderImpl::BuildNegativeCharacterClass(const GrammarExpr& expr) {
   XGRAMMAR_DCHECK(
       expr.type == ExprType::kCharacterClass || expr.type == ExprType::kCharacterClassStar
   );
@@ -1400,15 +1400,12 @@ std::optional<FSMWithStartEnd> GrammarFSMBuilderImpl::BuildNegativeCharacterClas
   return result_fsm;
 }
 
-std::optional<FSMWithStartEnd> GrammarFSMBuilderImpl::CharacterClass(const GrammarExpr& expr) {
+FSMWithStartEnd GrammarFSMBuilderImpl::CharacterClass(const GrammarExpr& expr) {
   bool is_negative = expr[0];
   FSMWithStartEnd result_fsm;
   if (is_negative) {
-    auto optional_fsm = BuildNegativeCharacterClass(expr);
-    if (!optional_fsm.has_value()) {
-      return std::nullopt;
-    }
-    return result_fsm = std::move(optional_fsm.value());
+    result_fsm = BuildNegativeCharacterClass(expr);
+    return result_fsm;
   }
   int start_state = result_fsm.AddState();
   result_fsm.SetStartState(start_state);
@@ -1438,28 +1435,16 @@ std::optional<FSMWithStartEnd> GrammarFSMBuilderImpl::Sequence(
     const auto& sequence_expr = grammar->GetGrammarExpr(sequence_id);
     switch (sequence_expr.type) {
       case (ExprType::kByteString): {
-        auto fsm = ByteString(sequence_expr);
-        if (!fsm.has_value()) {
-          return std::nullopt;
-        }
-        fsm_lists.push_back(std::move(fsm.value()));
+        fsm_lists.push_back(ByteString(sequence_expr));
         break;
       }
       case (ExprType::kRuleRef): {
-        auto fsm = RuleRef(sequence_expr);
-        if (!fsm.has_value()) {
-          return std::nullopt;
-        }
-        fsm_lists.push_back(std::move(fsm.value()));
+        fsm_lists.push_back(RuleRef(sequence_expr));
         break;
       }
       case (ExprType::kCharacterClass):
       case (ExprType::kCharacterClassStar): {
-        auto fsm = CharacterClass(sequence_expr);
-        if (!fsm.has_value()) {
-          return std::nullopt;
-        }
-        fsm_lists.push_back(std::move(fsm.value()));
+        fsm_lists.push_back(CharacterClass(sequence_expr));
         break;
       }
       default: {
@@ -1480,7 +1465,7 @@ std::optional<FSMWithStartEnd> GrammarFSMBuilderImpl::Sequence(
   return FSMWithStartEnd::Concat(fsm_lists);
 }
 
-std::optional<FSMWithStartEnd> GrammarFSMBuilderImpl::RuleRef(const GrammarExpr& expr) {
+FSMWithStartEnd GrammarFSMBuilderImpl::RuleRef(const GrammarExpr& expr) {
   FSMWithStartEnd result_fsm;
   result_fsm.AddState();
   result_fsm.AddState();
@@ -1490,7 +1475,7 @@ std::optional<FSMWithStartEnd> GrammarFSMBuilderImpl::RuleRef(const GrammarExpr&
   return result_fsm;
 }
 
-std::optional<FSMWithStartEnd> GrammarFSMBuilderImpl::ByteString(const GrammarExpr& expr) {
+FSMWithStartEnd GrammarFSMBuilderImpl::ByteString(const GrammarExpr& expr) {
   XGRAMMAR_DCHECK(expr.type == ExprType::kByteString);
   FSMWithStartEnd result_fsm;
   int current_state = result_fsm.AddState();
@@ -1580,13 +1565,13 @@ std::optional<FSMWithStartEnd> GrammarFSMBuilderImpl::BuildTagDispatchWithStopSt
       old_ends.insert(end);
     }
   }
-  std::unordered_set<int> ends;
+  std::vector<bool> ends(trie_fsm.NumStates(), false);
 
   // The final end states are the end of each stop string.
   for (int i = static_cast<int>(tag_dispatch_rules.size());
        i < static_cast<int>(trie_end_states.size());
        i++) {
-    ends.insert(trie_end_states[i]);
+    ends[trie_end_states[i]] = true;
   }
 
   if (loop_after_dispatch) {
@@ -1611,11 +1596,12 @@ std::optional<FSMWithStartEnd> GrammarFSMBuilderImpl::BuildTagDispatchWithStopSt
       }
     }
 
-    std::unordered_map<int, int> stop_trie_to_trie_map;
+    std::vector<int> stop_trie_to_trie_map;
     trie_fsm.AddFSM(stop_trie_fsm, &stop_trie_to_trie_map);
+    ends.resize(trie_fsm.NumStates(), false);
     int start_of_stop_trie = stop_trie_to_trie_map[stop_trie_start];
     for (auto state : stop_trie_ends) {
-      ends.insert(stop_trie_to_trie_map[state]);
+      ends[stop_trie_to_trie_map[state]] = true;
     }
 
     for (int i = 0; i < static_cast<int>(tag_dispatch_rules.size()); i++) {
@@ -1642,7 +1628,7 @@ std::optional<FSMWithStartEnd> GrammarFSMBuilderImpl::BuildTagDispatchWithEOSSto
   auto trie_fsm = trie_result->GetFsm();
   auto start = trie_result->GetStart();
   std::unordered_set<int> old_ends;
-  std::unordered_set<int> ends;
+  std::vector<bool> ends(trie_fsm.NumStates(), false);
   for (int end = 0; end < trie_result->NumStates(); end++) {
     if (trie_result->IsEndState(end)) {
       old_ends.insert(end);
@@ -1652,7 +1638,7 @@ std::optional<FSMWithStartEnd> GrammarFSMBuilderImpl::BuildTagDispatchWithEOSSto
   // The final end states are all but old_ends.
   for (int i = 0; i < trie_fsm.NumStates(); i++) {
     if (old_ends.count(i) == 0) {
-      ends.insert(i);
+      ends[i] = true;
     }
   }
 
@@ -1663,7 +1649,7 @@ std::optional<FSMWithStartEnd> GrammarFSMBuilderImpl::BuildTagDispatchWithEOSSto
       next_state = start;
     } else {
       next_state = trie_fsm.AddState();
-      ends.insert(next_state);
+      ends.push_back(true);
     }
     trie_fsm.AddRuleEdge(end_states[i], next_state, tag_dispatch_rules[i].second);
   }
@@ -1758,15 +1744,15 @@ void GrammarFSMBuilder::Apply(Grammar* grammar) { GrammarFSMBuilderImpl().Apply(
 
 void RepetitionNormalizer::Apply(Grammar* grammar) { RepetitionNormalizerImpl().Apply(grammar); }
 
-std::optional<FSMWithStartEnd> GrammarFSMBuilder::RuleRef(const GrammarExpr& expr) {
+FSMWithStartEnd GrammarFSMBuilder::RuleRef(const GrammarExpr& expr) {
   return GrammarFSMBuilderImpl::RuleRef(expr);
 }
 
-std::optional<FSMWithStartEnd> GrammarFSMBuilder::CharacterClass(const GrammarExpr& expr) {
+FSMWithStartEnd GrammarFSMBuilder::CharacterClass(const GrammarExpr& expr) {
   return GrammarFSMBuilderImpl::CharacterClass(expr);
 }
 
-std::optional<FSMWithStartEnd> GrammarFSMBuilder::ByteString(const GrammarExpr& expr) {
+FSMWithStartEnd GrammarFSMBuilder::ByteString(const GrammarExpr& expr) {
   return GrammarFSMBuilderImpl::ByteString(expr);
 }
 

cpp/grammar_functor.cc

@@ -347,9 +347,9 @@ class GrammarFSMBuilder {
 
  public:
   static void Apply(Grammar* grammar);
-  static std::optional<FSMWithStartEnd> RuleRef(const GrammarExpr& expr);
-  static std::optional<FSMWithStartEnd> CharacterClass(const GrammarExpr& expr);
-  static std::optional<FSMWithStartEnd> ByteString(const GrammarExpr& expr);
+  static FSMWithStartEnd RuleRef(const GrammarExpr& expr);
+  static FSMWithStartEnd CharacterClass(const GrammarExpr& expr);
+  static FSMWithStartEnd ByteString(const GrammarExpr& expr);
   static std::optional<FSMWithStartEnd> Sequence(const GrammarExpr& expr, const Grammar& grammar);
   static std::optional<FSMWithStartEnd> Choices(const GrammarExpr& expr, const Grammar& grammar);
   static std::optional<FSMWithStartEnd> TagDispatch(const Grammar::Impl::TagDispatch& tag_dispatch);