metrale_grammar/fsm/
fsm.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//
3// Fsm — a mutable finite-state machine (NFA or DFA).
4// Port of `class FSM` (+ `FSM::Impl`) from xgrammar `cpp/fsm.{h,cc}`.
5//
6// Representation: an adjacency list `edges[state] = Vec<FsmEdge>`, plus a
7// flat `edge_aux_data: Vec<i32>` holding `[rule_id, lower, upper]` triples
8// for repeat-reference edges. The C++ pimpl/shared_ptr indirection is
9// dropped — Rust callers `clone()` explicitly when they need a copy.
10
11use super::edge::{FsmEdge, RepeatEdgeRef, edge_type};
12
13/// Returned by [`Fsm::next_state`] when no transition exists.
14pub const NO_NEXT_STATE: i32 = -1;
15
16/// A mutable finite-state machine.
17#[derive(Debug, Clone, Default, PartialEq, Eq)]
18pub struct Fsm {
19    /// `edges[s]` is the list of outgoing edges from state `s`.
20    pub(crate) edges: Vec<Vec<FsmEdge>>,
21    /// Flat auxiliary data for repeat-reference edges.
22    pub(crate) edge_aux_data: Vec<i32>,
23}
24
25impl Fsm {
26    /// Create an FSM with `num_states` states and no edges.
27    pub fn with_states(num_states: usize) -> Self {
28        Self {
29            edges: vec![Vec::new(); num_states],
30            edge_aux_data: Vec::new(),
31        }
32    }
33
34    /// Create an FSM from a ready-made adjacency list.
35    pub fn from_edges(edges: Vec<Vec<FsmEdge>>, edge_aux_data: Vec<i32>) -> Self {
36        Self {
37            edges,
38            edge_aux_data,
39        }
40    }
41
42    /* ----------------------- visitors ----------------------- */
43
44    /// Number of states.
45    pub fn num_states(&self) -> usize {
46        self.edges.len()
47    }
48
49    /// Total number of edges across all states.
50    pub fn num_edges(&self) -> usize {
51        self.edges.iter().map(Vec::len).sum()
52    }
53
54    /// All adjacency rows.
55    pub fn all_edges(&self) -> &[Vec<FsmEdge>] {
56        &self.edges
57    }
58
59    /// Outgoing edges of `state`.
60    pub fn edges(&self, state: usize) -> &[FsmEdge] {
61        &self.edges[state]
62    }
63
64    /// Mutable outgoing edges of `state`.
65    pub fn edges_mut(&mut self, state: usize) -> &mut Vec<FsmEdge> {
66        &mut self.edges[state]
67    }
68
69    /// The repeat-edge auxiliary data buffer.
70    pub fn edge_aux_data(&self) -> &[i32] {
71        &self.edge_aux_data
72    }
73
74    /// Replace the repeat-edge auxiliary data buffer.
75    pub fn set_edge_aux_data(&mut self, data: Vec<i32>) {
76        self.edge_aux_data = data;
77    }
78
79    /// Decode the repeat-edge info stored at aux index `idx`.
80    pub fn repeat_edge_info(&self, idx: i16) -> RepeatEdgeRef {
81        RepeatEdgeRef::from_aux(&self.edge_aux_data, idx as usize)
82    }
83
84    /* ----------------------- mutators ----------------------- */
85
86    /// Append a fresh state; returns its id.
87    pub fn add_state(&mut self) -> usize {
88        self.edges.push(Vec::new());
89        self.edges.len() - 1
90    }
91
92    /// Add a character-range edge `from -[min,max]-> to`.
93    pub fn add_edge(&mut self, from: usize, to: usize, min: i16, max: i16) {
94        debug_assert!(from < self.edges.len());
95        self.edges[from].push(FsmEdge::new(min, max, to as i32));
96    }
97
98    /// Add a raw typed edge (`type` is one of [`edge_type`], `value`
99    /// goes into `max`).
100    pub fn add_typed_edge(&mut self, from: usize, to: usize, edge_type: i16, value: i16) {
101        self.add_edge(from, to, edge_type, value);
102    }
103
104    /// Add an epsilon transition.
105    pub fn add_epsilon_edge(&mut self, from: usize, to: usize) {
106        self.add_typed_edge(from, to, edge_type::EPSILON, 0);
107    }
108
109    /// Add a rule-reference edge for `rule_id`.
110    pub fn add_rule_edge(&mut self, from: usize, to: usize, rule_id: i16) {
111        self.add_typed_edge(from, to, edge_type::RULE_REF, rule_id);
112    }
113
114    /// Add an EOS edge.
115    pub fn add_eos_edge(&mut self, from: usize, to: usize) {
116        self.add_typed_edge(from, to, edge_type::EOS, 0);
117    }
118
119    /// Add a repeat-reference edge, allocating `[rule_id, lower, upper]`
120    /// into `edge_aux_data`. The source state must have no other edges.
121    pub fn add_repeat_edge(
122        &mut self,
123        from: usize,
124        to: usize,
125        rule_id: i32,
126        lower: i32,
127        upper: i32,
128    ) {
129        debug_assert!(
130            self.edges[from].is_empty(),
131            "A state with a kRepeatRef edge must have no other outgoing edges."
132        );
133        let aux_index = self.edge_aux_data.len() as i16;
134        self.edge_aux_data.push(rule_id);
135        self.edge_aux_data.push(lower);
136        self.edge_aux_data.push(upper);
137        self.add_typed_edge(from, to, edge_type::REPEAT_REF, aux_index);
138    }
139
140    /// Splice `other` into `self`, offsetting all of `other`'s state ids.
141    /// Returns the state mapping `other_id -> new_id`.
142    pub fn add_fsm(&mut self, other: &Fsm) -> Vec<usize> {
143        let old_num_states = self.num_states();
144        let aux_offset = self.edge_aux_data.len() as i16;
145        self.edge_aux_data.extend_from_slice(&other.edge_aux_data);
146
147        let state_mapping: Vec<usize> = (0..other.num_states())
148            .map(|i| i + old_num_states)
149            .collect();
150
151        self.edges
152            .resize(self.edges.len() + other.num_states(), Vec::new());
153
154        for (i, row) in other.edges.iter().enumerate() {
155            for edge in row {
156                let max_val = if edge.is_aux_edge() && aux_offset > 0 {
157                    edge.max + aux_offset
158                } else {
159                    edge.max
160                };
161                self.add_edge(
162                    i + old_num_states,
163                    edge.target as usize + old_num_states,
164                    edge.min,
165                    max_val,
166                );
167            }
168        }
169        state_mapping
170    }
171
172    /* --------------------- construction ---------------------- */
173
174    /// Sort every state's edges by `(min, max, target)`.
175    pub fn sort_edges(&mut self) {
176        for row in &mut self.edges {
177            row.sort_unstable();
178        }
179    }
180
181    /// Rebuild with remapped state ids. Epsilon self-loops created by the
182    /// mapping are dropped; duplicate edges per row are removed.
183    pub fn rebuild_with_mapping(&self, state_mapping: &[usize], new_num_states: usize) -> Fsm {
184        let mut new_edges: Vec<Vec<FsmEdge>> = vec![Vec::new(); new_num_states];
185        for (i, row) in self.edges.iter().enumerate() {
186            for edge in row {
187                if edge.is_epsilon() && state_mapping[i] == state_mapping[edge.target as usize] {
188                    continue; // skip epsilon self-loops
189                }
190                new_edges[state_mapping[i]].push(FsmEdge::new(
191                    edge.min,
192                    edge.max,
193                    state_mapping[edge.target as usize] as i32,
194                ));
195            }
196        }
197        for row in &mut new_edges {
198            row.sort_unstable();
199            row.dedup();
200        }
201        Fsm::from_edges(new_edges, self.edge_aux_data.clone())
202    }
203}
204
205// Traversal methods (next_state / epsilon_closure / advance /
206// possible_rules / reachable_states) live in `fsm_traversal.rs`.
207
208#[cfg(test)]
209#[path = "fsm_tests.rs"]
210mod tests;
metrale_grammar/fsm/fsm.rs

metrale_grammar/fsm/
fsm.rs
