metrale_grammar/fsm/
edge.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//
3// FSMEdge — the edge type of a finite-state machine.
4// Port of `struct FSMEdge` from xgrammar `cpp/fsm.h`.
5
6/// Sentinel edge-type values stored in `FSMEdge::min`.
7///
8/// When `min >= 0` the edge is a character range `[min, max]`.
9/// Negative values select a special edge type.
10pub mod edge_type {
11    /// `min >= 0` — a character range. This is the boundary value.
12    pub const CHAR_RANGE: i16 = 0;
13    /// Epsilon transition (no input consumed).
14    pub const EPSILON: i16 = -1;
15    /// Reference to another rule; `max` holds the rule id.
16    pub const RULE_REF: i16 = -2;
17    /// Accepts an end-of-sequence token.
18    pub const EOS: i16 = -3;
19    /// Repeated rule reference; `max` indexes into `edge_aux_data`.
20    pub const REPEAT_REF: i16 = -4;
21}
22
23/// The largest character value an edge range may carry.
24pub const MAX_CHAR: i16 = 255;
25
26/// An edge of a finite-state machine.
27///
28/// `min`/`max` encode the edge kind (see [`edge_type`]); `target` is the
29/// destination state id. The C++ struct is `alignas(8)` and packs
30/// `(i16, i16, i32)` — this layout is preserved so the compact
31/// representation matches the C++ memory layout byte-for-byte.
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
33pub struct FsmEdge {
34    /// Edge-kind discriminant / range lower bound. See [`edge_type`].
35    pub min: i16,
36    /// Range upper bound, rule id, or aux-data index depending on `min`.
37    pub max: i16,
38    /// Destination state id.
39    pub target: i32,
40}
41
42impl FsmEdge {
43    /// Construct an edge. Panics if a char-range edge has `min > max`,
44    /// matching the C++ `XGRAMMAR_DCHECK`.
45    pub fn new(min: i16, max: i16, target: i32) -> Self {
46        debug_assert!(
47            !(min >= edge_type::CHAR_RANGE && min > max),
48            "Invalid FsmEdge: min > max. min={min}, max={max}"
49        );
50        Self { min, max, target }
51    }
52
53    /// True if this is a character-range edge.
54    pub fn is_char_range(&self) -> bool {
55        self.min >= edge_type::CHAR_RANGE
56    }
57
58    /// True if this is an epsilon transition.
59    pub fn is_epsilon(&self) -> bool {
60        self.min == edge_type::EPSILON
61    }
62
63    /// True if this is a rule reference.
64    pub fn is_rule_ref(&self) -> bool {
65        self.min == edge_type::RULE_REF
66    }
67
68    /// True if this is an EOS transition.
69    pub fn is_eos(&self) -> bool {
70        self.min == edge_type::EOS
71    }
72
73    /// True if this is a repeat reference.
74    pub fn is_repeat_ref(&self) -> bool {
75        self.min == edge_type::REPEAT_REF
76    }
77
78    /// True if the edge uses `edge_aux_data` (currently only repeat refs).
79    pub fn is_aux_edge(&self) -> bool {
80        self.is_repeat_ref()
81    }
82
83    /// Referenced rule id, or `-1` if this is not a rule reference.
84    pub fn ref_rule_id(&self) -> i32 {
85        if self.is_rule_ref() {
86            self.max as i32
87        } else {
88            -1
89        }
90    }
91
92    /// Index into `edge_aux_data`, or `-1` if not a repeat reference.
93    pub fn aux_index(&self) -> i16 {
94        if self.is_repeat_ref() { self.max } else { -1 }
95    }
96}
97
98// FSMEdge's C++ `operator<` sorts by the tuple (min, max, target).
99impl PartialOrd for FsmEdge {
100    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
101        Some(self.cmp(other))
102    }
103}
104impl Ord for FsmEdge {
105    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
106        (self.min, self.max, self.target).cmp(&(other.min, other.max, other.target))
107    }
108}
109
110/// A view into `edge_aux_data` for a repeat edge.
111///
112/// Layout in the backing slice is `[rule_id, lower, upper]`.
113#[derive(Debug, Clone, Copy, PartialEq, Eq)]
114pub struct RepeatEdgeRef {
115    /// The rule being repeated.
116    pub rule_id: i16,
117    /// Minimum repeat count.
118    pub lower: i32,
119    /// Maximum repeat count (`-1` for unbounded).
120    pub upper: i32,
121}
122
123impl RepeatEdgeRef {
124    /// Build a [`RepeatEdgeRef`] from the 3-element aux slice starting at `idx`.
125    pub fn from_aux(data: &[i32], idx: usize) -> Self {
126        Self {
127            rule_id: data[idx] as i16,
128            lower: data[idx + 1],
129            upper: data[idx + 2],
130        }
131    }
132}
133
134/// Compare two edges by `(min, max)` only — used to deduplicate ranges.
135/// Mirrors the C++ `FSMEdgeRangeComparator`.
136pub fn cmp_edge_range(lhs: &FsmEdge, rhs: &FsmEdge) -> std::cmp::Ordering {
137    (lhs.min, lhs.max).cmp(&(rhs.min, rhs.max))
138}
139
140#[cfg(test)]
141mod tests {
142    use super::*;
143
144    #[test]
145    fn char_range_predicates() {
146        let e = FsmEdge::new(b'a' as i16, b'z' as i16, 3);
147        assert!(e.is_char_range());
148        assert!(!e.is_epsilon());
149        assert!(!e.is_rule_ref());
150        assert!(!e.is_eos());
151        assert!(!e.is_repeat_ref());
152        assert!(!e.is_aux_edge());
153        assert_eq!(e.ref_rule_id(), -1);
154        assert_eq!(e.aux_index(), -1);
155    }
156
157    #[test]
158    fn epsilon_edge() {
159        let e = FsmEdge::new(edge_type::EPSILON, 0, 7);
160        assert!(e.is_epsilon());
161        assert!(!e.is_char_range());
162    }
163
164    #[test]
165    fn rule_ref_edge() {
166        let e = FsmEdge::new(edge_type::RULE_REF, 5, 2);
167        assert!(e.is_rule_ref());
168        assert_eq!(e.ref_rule_id(), 5);
169    }
170
171    #[test]
172    fn eos_edge() {
173        let e = FsmEdge::new(edge_type::EOS, 0, 1);
174        assert!(e.is_eos());
175    }
176
177    #[test]
178    fn repeat_ref_edge() {
179        let e = FsmEdge::new(edge_type::REPEAT_REF, 9, 4);
180        assert!(e.is_repeat_ref());
181        assert!(e.is_aux_edge());
182        assert_eq!(e.aux_index(), 9);
183    }
184
185    #[test]
186    fn ordering_is_tuple_lexicographic() {
187        let a = FsmEdge::new(1, 2, 3);
188        let b = FsmEdge::new(1, 2, 4);
189        let c = FsmEdge::new(1, 3, 0);
190        assert!(a < b);
191        assert!(b < c);
192        let mut v = vec![c, b, a];
193        v.sort();
194        assert_eq!(v, vec![a, b, c]);
195    }
196
197    #[test]
198    fn range_comparator_ignores_target() {
199        let a = FsmEdge::new(1, 2, 99);
200        let b = FsmEdge::new(1, 2, 0);
201        assert_eq!(cmp_edge_range(&a, &b), std::cmp::Ordering::Equal);
202    }
203
204    #[test]
205    fn repeat_edge_ref_from_aux() {
206        let data = vec![0, 0, 0, 7, 2, 5];
207        let r = RepeatEdgeRef::from_aux(&data, 3);
208        assert_eq!(r.rule_id, 7);
209        assert_eq!(r.lower, 2);
210        assert_eq!(r.upper, 5);
211    }
212}
metrale_grammar/fsm/edge.rs

metrale_grammar/fsm/
edge.rs
