[Fix] Do not merge the FSM start state with equivalent-prefix states

cpp/fsm.cc

@@ -1452,9 +1452,14 @@ FSMWithStartEnd FSMWithStartEnd::MergeEquivalentStates(int max_result_num_states
     }
 
     // Case 1: Like ab | ac | ad, then they can be merged into a(b | c | d).
+    // The start state must not be merged by this rule. The rule reasons that states with
+    // identical incoming edges are reached by identical prefixes, but the start state is
+    // additionally reached by the empty prefix. Merging it with a sibling would wrongly
+    // expose the sibling's suffixes at the very start, e.g. for (ab)*abc the FSM would
+    // start accepting "c".
     bool is_equiv_successor = false;
     for (int i = 0; i < n; i++) {
-      if (incoming_distinct_count[i] != 1 || union_find_set.Count(i)) {
+      if (i == result.GetStart() || incoming_distinct_count[i] != 1 || union_find_set.Count(i)) {
         continue;
       }
       int previous_state = single_incoming_source[i];
@@ -1469,7 +1474,7 @@ FSMWithStartEnd FSMWithStartEnd::MergeEquivalentStates(int max_result_num_states
         }
         auto edges_to_sibling = siblings.Slice(group_begin, group_end);
         group_begin = group_end;
-        if (sibling <= i || incoming_distinct_count[sibling] != 1 ||
+        if (sibling <= i || sibling == result.GetStart() || incoming_distinct_count[sibling] != 1 ||
             result.IsEndState(sibling) != result.IsEndState(i)) {
           continue;
         }

tests/cpp/test_fsm.cc

@@ -523,6 +523,55 @@ TEST(XGrammarFSMTest, MergeEquivalentStatesNoCrossRuleChaining) {
   EXPECT_FALSE(merged.AcceptString("ybn"));
 }
 
+TEST(XGrammarFSMTest, MergeEquivalentStatesKeepsStartStateUnmerged) {
+  // FSM of (ab)*abc after epsilon simplification. In the first round states 1 and 2 are
+  // merged as equivalent successors of state 0. After that, the start state 0 and state 3
+  // share the single predecessor {1,2} with identical incoming edges 'b'. Merging them
+  // would wrongly accept "c", because the start state is also reached by the empty prefix.
+  FSMWithStartEnd fsm_wse;
+  for (int i = 0; i < 5; ++i) {
+    fsm_wse.AddState();
+  }
+  fsm_wse.SetStartState(0);
+  fsm_wse.AddEndState(4);
+  fsm_wse.GetFsm().AddEdge(0, 1, 'a', 'a');
+  fsm_wse.GetFsm().AddEdge(1, 0, 'b', 'b');
+  fsm_wse.GetFsm().AddEdge(0, 2, 'a', 'a');
+  fsm_wse.GetFsm().AddEdge(2, 3, 'b', 'b');
+  fsm_wse.GetFsm().AddEdge(3, 4, 'c', 'c');
+
+  auto merged = fsm_wse.MergeEquivalentStates();
+
+  // States 1 and 2 are still merged.
+  EXPECT_EQ(merged.NumStates(), 4);
+  EXPECT_TRUE(merged.AcceptString("abc"));
+  EXPECT_TRUE(merged.AcceptString("ababc"));
+  EXPECT_FALSE(merged.AcceptString("c"));
+  EXPECT_FALSE(merged.AcceptString("ab"));
+  EXPECT_FALSE(merged.AcceptString(""));
+}
+
+TEST(XGrammarFSMTest, MergeEquivalentStatesStartStateWithSelfLoop) {
+  // [a]*abc: the character class star is a self-loop on the start state, so the start
+  // state has incoming edges identical to its successor's. Merging them would wrongly
+  // accept "bc".
+  auto fsm_wse = RegexFSMBuilder::Build("[a]*abc").Unwrap();
+  fsm_wse = fsm_wse.SimplifyEpsilon();
+  auto merged = fsm_wse.MergeEquivalentStates();
+  EXPECT_TRUE(merged.AcceptString("abc"));
+  EXPECT_TRUE(merged.AcceptString("aaabc"));
+  EXPECT_FALSE(merged.AcceptString("bc"));
+  EXPECT_FALSE(merged.AcceptString("c"));
+
+  fsm_wse = RegexFSMBuilder::Build("(ab)*abc").Unwrap();
+  fsm_wse = fsm_wse.SimplifyEpsilon();
+  merged = fsm_wse.MergeEquivalentStates();
+  EXPECT_TRUE(merged.AcceptString("abc"));
+  EXPECT_TRUE(merged.AcceptString("ababc"));
+  EXPECT_FALSE(merged.AcceptString("c"));
+  EXPECT_FALSE(merged.AcceptString("ab"));
+}
+
 TEST(XGrammarFSMTest, SimplifyEpsilonPreservesAcceptance) {
   // (ab)+ : the accepting state's only edge is an epsilon back to the non-accepting start.
   // Merging the two states would wrongly make the start state accepting.

tests/python/test_grammar_matcher_ebnf.py

@@ -29,7 +29,20 @@ def test_simple():
     grammar = xgr.Grammar.from_ebnf(grammar_str)
     assert _is_grammar_accept_string(grammar, "bab")
     assert not _is_grammar_accept_string(grammar, "abb")
-    assert _is_grammar_accept_string(grammar, "cab")
+
+
+def test_char_class_star_prefix_requires_first_char():
+    """Regression test: [a]* forms a self-loop on the start state of the rule FSM. The
+    FSM state merging used to merge the start state with the state after the first "a",
+    wrongly accepting "bc" without any leading "a"."""
+    grammar = xgr.Grammar.from_ebnf('root ::= [a]* "a" "b" "c"')
+    assert _is_grammar_accept_string(grammar, "abc")
+    assert _is_grammar_accept_string(grammar, "aabc")
+    assert _is_grammar_accept_string(grammar, "aaabc")
+    assert not _is_grammar_accept_string(grammar, "bc")
+    assert not _is_grammar_accept_string(grammar, "c")
+    assert not _is_grammar_accept_string(grammar, "ab")
+    assert not _is_grammar_accept_string(grammar, "cab")
 
 
 input_accepted_test_repetition = (