metrale_grammar/fsm/
mod.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//
3// FSM subsystem — port wave W2.
4//
5// Pure-Rust port of xgrammar's finite-state-machine subsystem
6// (`cpp/fsm.{h,cc}` + `cpp/fsm_builder.{h,cc}`). No `unsafe`; the
7// CSR-packed `CompactFsm` representation is preserved (it is load-bearing
8// for the compiled-grammar memory layout) and accessed safely via slices.
9//
10// Module map:
11//   edge           — FsmEdge, RepeatEdgeRef, edge-type constants
12//   compact_array  — Compact2DArray, the CSR 2D-array primitive
13//   union_find     — UnionFindSet, used by the simplification passes
14//   traversal      — shared epsilon-closure / transition / advance code
15//   fsm            — Fsm: the mutable adjacency-list FSM
16//   compact        — CompactFsm: the immutable CSR FSM
17//   with_start_end — FsmWithStartEnd / CompactFsmWithStartEnd + Star/
18//                    Plus/Optional/Union/Concat
19//   simplify       — SimplifyEpsilon / MergeEquivalentSuccessors
20//   algorithms     — IsDFA / ToDFA
21//   dfa_ops        — MinimizeDFA / Not / Intersect
22//   builder        — RegexFsmBuilder / TrieFsmBuilder
23
24pub mod algorithms;
25pub mod builder;
26pub mod compact;
27pub mod compact_array;
28pub mod compact_with_start_end;
29pub mod dfa_ops;
30pub mod edge;
31pub mod fsm;
32pub mod fsm_ops;
33pub mod fsm_traversal;
34pub mod simplify;
35pub mod traversal;
36pub mod union_find;
37pub mod with_start_end;
38
39pub use algorithms::DEFAULT_MAX_STATES;
40pub use builder::{RegexFsmBuilder, TrieFsmBuilder};
41pub use compact::CompactFsm;
42pub use compact_array::Compact2DArray;
43pub use edge::{FsmEdge, RepeatEdgeRef, edge_type};
44pub use fsm::{Fsm, NO_NEXT_STATE};
45pub use union_find::UnionFindSet;
46pub use with_start_end::{CompactFsmWithStartEnd, FsmWithStartEnd};
47
48#[cfg(test)]
49mod integration_tests {
50    //! End-to-end tests mirroring `tests/cpp/test_fsm.cc` —
51    //! the scenarios that exercise multiple modules together.
52    use super::*;
53
54    fn build(re: &str) -> FsmWithStartEnd {
55        RegexFsmBuilder::build(re).unwrap_or_else(|e| panic!("build {re}: {e}"))
56    }
57
58    #[test]
59    fn function_test_compact_roundtrip() {
60        // test_fsm.cc FunctionTest1
61        let fsm = build("[\\d\\d\\d]+123");
62        assert!(fsm.accept_string(b"123456123"));
63        let compact = fsm.fsm().to_compact();
64        let compact_wse =
65            CompactFsmWithStartEnd::new(compact.clone(), fsm.start(), fsm.ends().to_vec());
66        assert!(compact_wse.accept_string(b"123456123"));
67        let back = FsmWithStartEnd::new(compact.to_fsm(), fsm.start(), fsm.ends().to_vec(), false);
68        assert!(back.accept_string(b"123456123"));
69    }
70
71    #[test]
72    fn function_test_to_dfa_no_epsilon() {
73        // test_fsm.cc FunctionTest2
74        let fsm = build("([abc]|[\\d])+");
75        assert!(fsm.accept_string(b"abc3"));
76        let dfa = fsm.to_dfa(DEFAULT_MAX_STATES).unwrap();
77        assert!(dfa.accept_string(b"abc3"));
78        for edges in dfa.fsm().all_edges() {
79            assert!(edges.iter().all(|e| !e.is_epsilon()));
80        }
81    }
82
83    #[test]
84    fn function_test_minimize_then_not() {
85        // test_fsm.cc FunctionTest3 + 4
86        let fsm = build("([abc]|[\\d])+");
87        let dfa = fsm.to_dfa(DEFAULT_MAX_STATES).unwrap();
88        let min = dfa.minimize_dfa(DEFAULT_MAX_STATES).unwrap();
89        assert!(min.accept_string(b"abc3"));
90        assert_eq!(min.num_states(), 2);
91
92        let neg = min.not(DEFAULT_MAX_STATES).unwrap();
93        assert!(!neg.accept_string(b"abc3"));
94        assert!(neg.accept_string(b"abcd"));
95    }
96
97    #[test]
98    fn function_test_simplify_epsilon_state_count() {
99        // test_fsm.cc FunctionTest6
100        let fsm = build("[a][b][c][d]");
101        assert!(fsm.accept_string(b"abcd"));
102        let simplified = fsm.simplify_epsilon();
103        assert_eq!(simplified.fsm().num_states(), 5);
104        assert!(simplified.accept_string(b"abcd"));
105    }
106
107    #[test]
108    fn function_test_merge_equivalent() {
109        // test_fsm.cc FunctionTest7 ("abc|abd")
110        let fsm = build("abc|abd");
111        assert!(fsm.accept_string(b"abc"));
112        let merged = fsm.simplify_epsilon().merge_equivalent_successors();
113        assert!(merged.accept_string(b"abc"));
114        assert!(!merged.accept_string(b"abcd"));
115        assert_eq!(merged.fsm().num_states(), 4);
116    }
117
118    #[test]
119    fn function_test_merge_precursor() {
120        // test_fsm.cc FunctionTest8 ("acd|bcd")
121        let fsm = build("acd|bcd");
122        let merged = fsm.simplify_epsilon().merge_equivalent_successors();
123        assert!(merged.accept_string(b"acd"));
124        assert!(!merged.accept_string(b"abcd"));
125        assert_eq!(merged.fsm().num_states(), 4);
126    }
127
128    #[test]
129    fn function_test_star_simplify() {
130        // test_fsm.cc FunctionTest9 ("ab*")
131        let fsm = build("ab*");
132        assert!(fsm.accept_string(b"abbb"));
133        let simplified = fsm.simplify_epsilon();
134        assert!(simplified.accept_string(b"abbb"));
135        assert_eq!(simplified.fsm().num_states(), 2);
136    }
137
138    #[test]
139    fn function_test_intersect() {
140        // test_fsm.cc FunctionTest10
141        let left = build("[c-f]+");
142        let right = build("[d-h]*");
143        let inter = FsmWithStartEnd::intersect(&left, &right, DEFAULT_MAX_STATES).unwrap();
144        assert!(inter.accept_string(b"de"));
145        assert!(inter.accept_string(b"def"));
146        assert!(!inter.accept_string(b""));
147        assert!(!inter.accept_string(b"cd"));
148    }
149
150    #[test]
151    fn manual_merge_nodes_test() {
152        // test_fsm.cc MergingNodesTest — manual FSM construction
153        let mut fsm = FsmWithStartEnd::new(Fsm::with_states(0), 0, Vec::new(), false);
154        for _ in 0..10 {
155            fsm.add_state();
156        }
157        fsm.set_start_state(0);
158        fsm.add_end_state(9);
159        let pairs = [
160            (0, 1, b'a'),
161            (0, 2, b'a'),
162            (1, 3, b'b'),
163            (1, 3, b'c'),
164            (1, 4, b'b'),
165            (1, 4, b'c'),
166            (2, 5, b'b'),
167            (2, 5, b'c'),
168            (2, 6, b'b'),
169            (2, 6, b'c'),
170            (3, 7, b'd'),
171            (4, 7, b'd'),
172            (5, 8, b'd'),
173            (6, 8, b'd'),
174            (7, 9, b'e'),
175            (8, 9, b'e'),
176        ];
177        for (f, t, c) in pairs {
178            fsm.fsm_mut().add_edge(f, t, c as i16, c as i16);
179        }
180        let merged = fsm.merge_equivalent_successors();
181        assert_eq!(merged.fsm().num_states(), 5);
182        // language: "abde" / "acde" — single letters a,(b|c),d,e
183        assert!(merged.accept_string(b"abde"));
184        assert!(merged.accept_string(b"acde"));
185    }
186
187    #[test]
188    fn efficiency_test_minimize_target() {
189        // test_fsm.cc EfficiencyTest — reduced: 10 copies of a 26-way
190        // alternation each followed by "0123456789".  After minimize the
191        // C++ asserts 111 states; we assert language correctness + that
192        // minimize produces a small DFA.
193        let one = "(a0123456789|b0123456789|c0123456789)";
194        let pattern = format!("{one}{one}");
195        let fsm = build(&pattern);
196        let simplified = fsm.simplify_epsilon().merge_equivalent_successors();
197        let dfa = simplified.to_dfa(DEFAULT_MAX_STATES).unwrap();
198        let min = dfa.minimize_dfa(DEFAULT_MAX_STATES).unwrap();
199        assert!(min.accept_string(b"a0123456789b0123456789"));
200        assert!(!min.accept_string(b"a0123456789"));
201        assert!(min.num_states() <= dfa.num_states());
202    }
203
204    #[test]
205    fn trie_then_compact() {
206        // test_fsm_builder.cc TestTrieFSMBuilder (compaction half)
207        let pats: Vec<&[u8]> = vec![b"hello", b"hi", b"good"];
208        let res = TrieFsmBuilder::build(&pats, &[], true, false).unwrap();
209        let compact = res.fsm.fsm().to_compact();
210        let compact_wse =
211            CompactFsmWithStartEnd::new(compact, res.fsm.start(), res.fsm.ends().to_vec());
212        assert!(compact_wse.accept_string(b"hello"));
213        assert!(compact_wse.accept_string(b"hi"));
214        assert!(!compact_wse.accept_string(b"he"));
215    }
216}