Conversation
Avoid rebuilding the complete graph after every merge round while preserving the existing state layout and small-FSM performance.
Lock down early-return behavior and exact layouts for duplicate, self-loop, and special-edge inputs.
Contributor
There was a problem hiding this comment.
Pull request overview
Note
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
This PR refactors FSMWithStartEnd::MergeEquivalentStates to incrementally merge equivalent FSM states using stable IDs, union-find classes, and a dirty-state worklist, rebuilding the physical FSM only once at the end. It also adds targeted correctness tests and a standalone benchmark to reproduce key workload shapes.
Changes:
- Replaced the previous “rebuild every round” merge approach with an incremental, worklist-driven merge algorithm.
- Added extensive unit tests covering leaf merges, deep prefixes/suffixes, cycles, special edges, and randomized language preservation.
- Added a benchmark program to measure MergeEquivalentStates performance across small/large and deep-merge shapes.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| tests/cpp/test_fsm.cc | Adds new tests that exercise many MergeEquivalentStates edge cases and invariants. |
| examples/benchmark/bench_merge_equivalent_states.cc | Adds a reproducible microbenchmark driver for measuring MergeEquivalentStates behavior. |
| cpp/fsm.cc | Implements the incremental equivalent-state merge algorithm with union-find + dirty worklist and single rebuild. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
1374
to
1392
| const auto& edges = result.GetFsm().GetEdges(source); | ||
| outgoing_row_sizes[source] = static_cast<int32_t>(edges.size()); | ||
| for (const auto& edge : edges) { | ||
| ++incoming_row_sizes[edge.target]; | ||
| for (int32_t edge_index = 0; edge_index < static_cast<int32_t>(edges.size()); ++edge_index) { | ||
| const auto& edge = edges[edge_index]; | ||
| ++original_incoming_row_sizes[edge.target]; | ||
| if ((edge.IsEpsilon() && edge.target == source) || | ||
| (edge_index > 0 && edge == edges[edge_index - 1])) { | ||
| input_edges_are_canonical = false; | ||
| } | ||
| if (edge_index > 0 && edges[edge_index - 1].min == edge.min && | ||
| edges[edge_index - 1].max == edge.max && edges[edge_index - 1].target != edge.target) { | ||
| int previous_target = edges[edge_index - 1].target; | ||
| bool merges_only_leaves = | ||
| result.GetFsm().GetEdges(previous_target).empty() && | ||
| result.GetFsm().GetEdges(edge.target).empty() && | ||
| result.IsEndState(previous_target) == result.IsEndState(edge.target); | ||
| may_have_non_leaf_merge |= !merges_only_leaves; | ||
| } | ||
| } | ||
| } |
Comment on lines
+1881
to
+1884
| std::vector<int> next_active_states; | ||
| next_active_states.reserve( | ||
| active_num_states - round_touched_states.size() + equivalent_classes.size() | ||
| ); |
Comment on lines
+1397
to
+1400
| for (const auto& edge : result.GetFsm().GetEdges(source)) { | ||
| original_incoming_edges.MutableRowAt(edge.target | ||
| )[original_incoming_write_positions[edge.target]++] = {source, edge.min, edge.max}; | ||
| } |
15 tasks
Ubospica
changed the title
…rges Drop the byte-identical-output constraint of MergeEquivalentStates and guarantee language equivalence instead. This removes the mechanisms that only existed to reproduce the old rebuild-per-round output byte by byte (has_rebuilt, current_ids and per-round renumbering, the round-local union-find, the leaf-only pre-loop, the 15-array arena), and canonicalizes class-level edges (dedup + epsilon self-loop removal) from round one. Also fix two over-merge bugs that exist on main and could change the accepted language: - Leaf merges chained with Case-1 merges through a shared state in the same round, combining a reaching-path equivalence with a continuation equivalence that do not compose. - Case 1 could pick the start state as a candidate, ignoring that the start state is also reachable by the empty string without any edge. Merged states are now numbered by each class's smallest original state id, so surviving states keep their original relative order. Exact-layout test expectations and FSM structure digests are updated accordingly; a 50k-random-FSM differential check against main shows zero language changes for the new implementation (main breaks 1257 of them).
4 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Changes
This PR splits the "merge equivalent states" optimization out of #727 and reworks
MergeEquivalentStates.The old implementation rebuilt and renumbered the whole FSM after every logical merge round. Deep shared prefixes or suffixes can take many rounds, so the same states and edges were copied and sorted over and over.
The new implementation uses a single incremental algorithm for all FSM sizes:
The result is guaranteed to accept the same language as the input; it is not guaranteed to be byte-identical to the old implementation's output. The earlier byte-identical goal did not hold anyway (see Correctness below), and dropping it removes the machinery that only existed to emulate the old rebuild-per-round output: the flag that switched edge canonicalization on after the first rebuild, the per-round renumbering used for tie-breaking, the round-local union-find, and the separate leaf-only fast path.
Correctness
This PR also fixes two over-merge bugs that exist on main and can change the accepted language:
Validation:
test_fsm_structure_stability.pyare regenerated. Note that the three digests that this PR previously failed in CI are exactly the large cases where the old rebuild-per-round emulation was not byte-identical; with the language-equivalence semantics the numbering is deterministic and the digests are stable again.12493908.Performance
Comparing main commit
c5717178against this change (850cf68e) on Linux x86-64, same optimized build, pinned to one core, three interleaved runs per build, median reported. Benchmark:examples/benchmark/bench_merge_equivalent_states.cc.Small FSMs:
Large FSMs:
Grammar construction via
Grammar::FromEBNF(long sequence of 10,000 element pairs, single choice with 8,000 alternatives, 8,000 chained rules) is within measurement noise of main (-5% to +4%).The only regressions are the synthetic single-round star graphs at 8,000 and 80,000 states, where one giant merge group is found immediately; the extra cost is the outgoing-edge index and the canonical edge collection. No grammar-construction input shows a regression.
Validation