metrale_grammar/fsm/
simplify.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//
3// FSM simplification passes — `SimplifyEpsilon` and
4// `MergeEquivalentSuccessors`. Port of the corresponding methods of
5// `FSMWithStartEnd` from `cpp/fsm.cc`.
6
7use super::union_find::UnionFindSet;
8use super::with_start_end::FsmWithStartEnd;
9
10/// A compact edge endpoint used for the incoming/outgoing CSR rows of
11/// [`FsmWithStartEnd::merge_equivalent_successors`]. Port of upstream
12/// `EndpointEdge` (xgrammar commit 96ae88b).
13///
14/// `peer` is the *source* state in incoming rows and the *target* state
15/// in outgoing rows. Outgoing rows are sorted by `(peer, min, max)` so
16/// successors of one source form contiguous groups.
17#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
18struct EndpointEdge {
19    peer: i32,
20    min: i16,
21    max: i16,
22}
23
24/// A CSR (compressed-sparse-row) array of [`EndpointEdge`]: `data`
25/// holds all rows contiguously, `indptr[i]..indptr[i + 1]` delimits
26/// row `i`. Replaces the per-state `HashMap`s that upstream commit
27/// 96ae88b removed: row sizing + filling is two flat passes with no
28/// per-edge allocation, and the buffers are reused across iterations.
29#[derive(Debug, Default)]
30struct EdgeCsr {
31    data: Vec<EndpointEdge>,
32    indptr: Vec<i32>,
33}
34
35impl EdgeCsr {
36    /// Resize the CSR so row `i` has `row_sizes[i]` default slots,
37    /// reusing existing allocations. Port of `ResetWithRowSizes`.
38    fn reset_with_row_sizes(&mut self, row_sizes: &[i32]) {
39        self.indptr.clear();
40        self.indptr.reserve(row_sizes.len() + 1);
41        self.indptr.push(0);
42        let mut acc = 0i32;
43        for &sz in row_sizes {
44            debug_assert!(sz >= 0);
45            acc += sz;
46            self.indptr.push(acc);
47        }
48        self.data.clear();
49        self.data.resize(
50            acc as usize,
51            EndpointEdge {
52                peer: 0,
53                min: 0,
54                max: 0,
55            },
56        );
57    }
58
59    /// Immutable view of row `i`.
60    fn row(&self, i: usize) -> &[EndpointEdge] {
61        let start = self.indptr[i] as usize;
62        let end = self.indptr[i + 1] as usize;
63        &self.data[start..end]
64    }
65
66    /// Mutable view of row `i`.
67    fn row_mut(&mut self, i: usize) -> &mut [EndpointEdge] {
68        let start = self.indptr[i] as usize;
69        let end = self.indptr[i + 1] as usize;
70        &mut self.data[start..end]
71    }
72}
73
74impl FsmWithStartEnd {
75    /// Merge states linked by removable epsilon transitions.
76    ///
77    /// `a --eps--> b` is collapsible when either (1) `a` has no other
78    /// outgoing edge, or (2) `b` has no other incoming edge.
79    pub fn simplify_epsilon(&self) -> FsmWithStartEnd {
80        if self.is_dfa {
81            return self.clone();
82        }
83        let n = self.num_states();
84
85        let mut uf = UnionFindSet::new();
86        let mut in_degree = vec![0i32; n];
87        let mut epsilon_edges: Vec<(usize, usize)> = Vec::new();
88
89        for i in 0..n {
90            let edges = self.fsm.edges(i);
91            for edge in edges {
92                in_degree[edge.target as usize] += 1;
93                if edge.is_epsilon() {
94                    if edges.len() == 1 {
95                        // case 1: `a` has only this outgoing edge
96                        uf.add(i as i32);
97                        uf.add(edge.target);
98                        uf.union(i as i32, edge.target);
99                        in_degree[edge.target as usize] -= 1;
100                    } else {
101                        epsilon_edges.push((i, edge.target as usize));
102                    }
103                }
104            }
105        }
106
107        // Build the equivalence representative per node.
108        let mut equiv_node = vec![0usize; n];
109        for i in 0..n {
110            if uf.contains(i as i32) {
111                let rep = uf.find(i as i32) as usize;
112                equiv_node[i] = rep;
113                if rep != i {
114                    in_degree[rep] += in_degree[i];
115                }
116            } else {
117                equiv_node[i] = i;
118            }
119        }
120
121        // case 2: `a --eps--> b`, `b` has no other incoming edge.
122        for &(from_raw, to_raw) in &epsilon_edges {
123            let from = equiv_node[from_raw];
124            let to = equiv_node[to_raw];
125            if in_degree[to] == 1 && equiv_node[self.start] != to {
126                uf.add(from as i32);
127                uf.add(to as i32);
128                uf.union(from as i32, to as i32);
129            }
130        }
131
132        let eq_classes = uf.all_sets();
133        if eq_classes.is_empty() {
134            return self.clone();
135        }
136
137        let mut new_to_old = vec![-1i64; n];
138        for (i, class) in eq_classes.iter().enumerate() {
139            for &state in class {
140                new_to_old[state as usize] = i as i64;
141            }
142        }
143        let mut cnt = eq_classes.len();
144        for slot in new_to_old.iter_mut() {
145            if *slot == -1 {
146                *slot = cnt as i64;
147                cnt += 1;
148            }
149        }
150        let mapping: Vec<usize> = new_to_old.iter().map(|&v| v as usize).collect();
151        self.rebuild_with_mapping(&mapping, cnt)
152    }
153
154    /// Merge states with identical incoming or outgoing transition
155    /// structure (`ab | ac | ad` -> `a(b|c|d)`, and the mirror case).
156    ///
157    /// Upstream renamed this `MergeEquivalentStates` (commit 8d22ba0);
158    /// the port keeps the original name.
159    ///
160    /// The incoming/outgoing edge relations are stored in two reused
161    /// CSR arrays (`EdgeCsr`) instead of per-state hash maps — port
162    /// of upstream commit 96ae88b (#616), which found `build_maps`
163    /// dominated the pass through hash-map and small-vector churn.
164    pub fn merge_equivalent_successors(&self) -> FsmWithStartEnd {
165        // No merge is possible with fewer than 4 states: a Case 1 merge
166        // needs >=2 sinks sharing a source, a Case 2 merge needs >=2
167        // sources sharing a sink (upstream commit 96ae88b, #616).
168        if self.num_states() < 4 {
169            return self.copy();
170        }
171        let mut result = self.copy();
172        result.fsm_mut().sort_edges();
173        let mut uf = UnionFindSet::new();
174        let mut changed = true;
175
176        // Scratch buffers hoisted out of the loop so their capacity is
177        // reused across iterations (upstream commit 96ae88b).
178        let mut incoming = EdgeCsr::default();
179        let mut outgoing = EdgeCsr::default();
180        let mut incoming_row_sizes: Vec<i32> = Vec::new();
181        let mut outgoing_row_sizes: Vec<i32> = Vec::new();
182        let mut incoming_write_pos: Vec<i32> = Vec::new();
183        let mut outgoing_write_pos: Vec<i32> = Vec::new();
184        // For each state: distinct-peer count and (when count == 1) the
185        // single peer; -1 otherwise.
186        let mut incoming_distinct: Vec<i32> = Vec::new();
187        let mut outgoing_distinct: Vec<i32> = Vec::new();
188        let mut single_incoming: Vec<i32> = Vec::new();
189        let mut single_outgoing: Vec<i32> = Vec::new();
190        let mut no_succ_end: Vec<usize> = Vec::new();
191        let mut no_succ_non_end: Vec<usize> = Vec::new();
192
193        while changed {
194            uf.clear();
195            let n = result.num_states();
196
197            // First pass: count incoming/outgoing edges per state.
198            incoming_row_sizes.clear();
199            incoming_row_sizes.resize(n, 0);
200            outgoing_row_sizes.clear();
201            outgoing_row_sizes.resize(n, 0);
202            for source in 0..n {
203                let edges = result.fsm().edges(source);
204                outgoing_row_sizes[source] = edges.len() as i32;
205                for edge in edges {
206                    incoming_row_sizes[edge.target as usize] += 1;
207                }
208            }
209
210            // Allocate CSR rows (reusing backing storage).
211            incoming.reset_with_row_sizes(&incoming_row_sizes);
212            outgoing.reset_with_row_sizes(&outgoing_row_sizes);
213            incoming_write_pos.clear();
214            incoming_write_pos.resize(n, 0);
215            outgoing_write_pos.clear();
216            outgoing_write_pos.resize(n, 0);
217
218            // Second pass: fill incoming/outgoing rows. Incoming rows
219            // are naturally grouped by source (sources scanned in
220            // order); outgoing rows are sorted by `(peer, min, max)`.
221            for source in 0..n {
222                for edge in result.fsm().edges(source) {
223                    let t = edge.target as usize;
224                    let ip = incoming_write_pos[t];
225                    incoming.row_mut(t)[ip as usize] = EndpointEdge {
226                        peer: source as i32,
227                        min: edge.min,
228                        max: edge.max,
229                    };
230                    incoming_write_pos[t] = ip + 1;
231                    let op = outgoing_write_pos[source];
232                    outgoing.row_mut(source)[op as usize] = EndpointEdge {
233                        peer: edge.target,
234                        min: edge.min,
235                        max: edge.max,
236                    };
237                    outgoing_write_pos[source] = op + 1;
238                }
239                outgoing.row_mut(source).sort_unstable();
240            }
241
242            // Distinct-peer counts. A peer count of 1 with its identity
243            // mirrors the old `map.len() == 1` + `keys().next()` usage.
244            distinct_peers(&incoming, n, &mut incoming_distinct, &mut single_incoming);
245            distinct_peers(&outgoing, n, &mut outgoing_distinct, &mut single_outgoing);
246
247            let mut equiv_successor = false;
248            // Case 1: ab|ac|ad -> a(b|c|d)
249            for i in 0..n {
250                if incoming_distinct[i] != 1 || uf.contains(i as i32) {
251                    continue;
252                }
253                let prev_state = single_incoming[i] as usize;
254                let edges_to_i = incoming.row(i);
255                let siblings = outgoing.row(prev_state);
256                let mut group_begin = 0usize;
257                while group_begin < siblings.len() {
258                    let sibling = siblings[group_begin].peer as usize;
259                    let mut group_end = group_begin + 1;
260                    while group_end < siblings.len() && siblings[group_end].peer as usize == sibling
261                    {
262                        group_end += 1;
263                    }
264                    let edges_to_sibling = &siblings[group_begin..group_end];
265                    group_begin = group_end;
266                    if sibling <= i
267                        || incoming_distinct[sibling] != 1
268                        || result.is_end_state(sibling) != result.is_end_state(i)
269                    {
270                        continue;
271                    }
272                    if edges_to_i.len() != edges_to_sibling.len() {
273                        continue;
274                    }
275                    let same = edges_to_i
276                        .iter()
277                        .zip(edges_to_sibling)
278                        .all(|(a, b)| a.min == b.min && a.max == b.max);
279                    if same {
280                        uf.add(i as i32);
281                        uf.add(sibling as i32);
282                        uf.union(i as i32, sibling as i32);
283                        equiv_successor = true;
284                    }
285                }
286            }
287
288            // Case 2: ba|ca|da -> (b|c|d)a, plus dead-end merges.
289            let mut equiv_precursor = false;
290            no_succ_end.clear();
291            no_succ_non_end.clear();
292            for i in 0..n {
293                if outgoing_distinct[i] == 0 {
294                    if result.is_end_state(i) {
295                        no_succ_end.push(i);
296                    } else {
297                        no_succ_non_end.push(i);
298                    }
299                    continue;
300                }
301                if outgoing_distinct[i] != 1 || uf.contains(i as i32) {
302                    continue;
303                }
304                let next_state = single_outgoing[i] as usize;
305                let node_edges = outgoing.row(i);
306                let siblings = incoming.row(next_state);
307                let mut group_begin = 0usize;
308                while group_begin < siblings.len() {
309                    let sibling = siblings[group_begin].peer as usize;
310                    while group_begin < siblings.len()
311                        && siblings[group_begin].peer as usize == sibling
312                    {
313                        group_begin += 1;
314                    }
315                    // Skip a sibling already merged earlier this iteration
316                    // (typically by Case 1): chaining a Case 2 merge onto it
317                    // can over-merge via transitive closure (upstream 8d22ba0,
318                    // #632). The `equiv_precursor` flag below was also fixed
319                    // from a wrongly-set `equiv_successor` in the same commit.
320                    if sibling <= i
321                        || uf.contains(sibling as i32)
322                        || outgoing_distinct[sibling] != 1
323                        || result.is_end_state(i) != result.is_end_state(sibling)
324                    {
325                        continue;
326                    }
327                    let sibling_edges = outgoing.row(sibling);
328                    if sibling_edges.len() != node_edges.len() {
329                        continue;
330                    }
331                    let same = sibling_edges
332                        .iter()
333                        .zip(node_edges)
334                        .all(|(a, b)| a.min == b.min && a.max == b.max);
335                    if same {
336                        uf.add(i as i32);
337                        uf.add(sibling as i32);
338                        uf.union(i as i32, sibling as i32);
339                        equiv_precursor = true;
340                    }
341                }
342            }
343
344            if no_succ_end.len() > 1 {
345                for &s in &no_succ_end[1..] {
346                    uf.add(no_succ_end[0] as i32);
347                    uf.add(s as i32);
348                    uf.union(no_succ_end[0] as i32, s as i32);
349                    equiv_precursor = true;
350                }
351            }
352            if no_succ_non_end.len() > 1 {
353                for &s in &no_succ_non_end[1..] {
354                    uf.add(no_succ_non_end[0] as i32);
355                    uf.add(s as i32);
356                    uf.union(no_succ_non_end[0] as i32, s as i32);
357                    equiv_precursor = true;
358                }
359            }
360
361            changed = equiv_successor || equiv_precursor;
362            if changed {
363                let eq_classes = uf.all_sets();
364                let mut old_to_new = vec![-1i64; n];
365                for (idx, class) in eq_classes.iter().enumerate() {
366                    for &state in class {
367                        old_to_new[state as usize] = idx as i64;
368                    }
369                }
370                let mut cnt = eq_classes.len();
371                for slot in old_to_new.iter_mut() {
372                    if *slot == -1 {
373                        *slot = cnt as i64;
374                        cnt += 1;
375                    }
376                }
377                let mapping: Vec<usize> = old_to_new.iter().map(|&v| v as usize).collect();
378                result = result.rebuild_with_mapping(&mapping, cnt);
379                result.fsm_mut().sort_edges();
380            }
381        }
382        result
383    }
384}
385
386/// Compute, for each of the first `n` rows of `csr`, the number of
387/// distinct `peer` values and (when exactly one) that peer's id.
388///
389/// `distinct[i]` and `single[i]` are resized and overwritten. A row
390/// with one distinct peer mirrors the old `map.len() == 1` check; the
391/// single peer mirrors `map.keys().next()`. Rows are pre-sorted by
392/// `peer`, so distinct peers are counted by scanning for changes.
393fn distinct_peers(csr: &EdgeCsr, n: usize, distinct: &mut Vec<i32>, single: &mut Vec<i32>) {
394    distinct.clear();
395    distinct.resize(n, 0);
396    single.clear();
397    single.resize(n, -1);
398    for (state, slot) in distinct.iter_mut().enumerate().take(n) {
399        let row = csr.row(state);
400        if row.is_empty() {
401            continue;
402        }
403        let mut count = 1;
404        let mut only = row[0].peer;
405        for w in row.windows(2) {
406            if w[0].peer != w[1].peer {
407                count += 1;
408                only = -1;
409            }
410        }
411        *slot = count;
412        single[state] = only;
413    }
414}
415
416#[cfg(test)]
417#[path = "simplify_tests.rs"]
418mod tests;
metrale_grammar/fsm/simplify.rs

metrale_grammar/fsm/
simplify.rs
