metrale_grammar/grammar/
data.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//
3// GrammarData — the BNF grammar AST.
4// Port of `Grammar::Impl` from xgrammar `cpp/grammar_impl.h`.
5//
6// The AST is a set of rules; each rule's body is a `GrammarExpr`.
7// Expressions are stored CSR-style: all payloads packed end-to-end in
8// one `expr_data` vector, with `expr_indptr` recording each
9// expression's start offset. At offset `s` the layout is
10// `[type, data_len, data0, data1, …]`.
11
12use super::expr::{GrammarExpr, GrammarExprType};
13use crate::fsm::{CompactFsm, CompactFsmWithStartEnd};
14
15/// A named production rule. `rule_id` is this rule's index in
16/// [`GrammarData::rules`].
17#[derive(Debug, Clone, PartialEq, Eq)]
18pub struct Rule {
19    /// The rule name.
20    pub name: String,
21    /// The `GrammarExpr` id of the rule body.
22    pub body_expr_id: i32,
23    /// Id of the associated lookahead-assertion expr (a Sequence),
24    /// or `-1` when the rule has none.
25    pub lookahead_assertion_id: i32,
26    /// Whether the lookahead assertion is exact.
27    pub is_exact_lookahead: bool,
28}
29
30impl Rule {
31    /// A rule with no lookahead assertion.
32    pub fn new(name: impl Into<String>, body_expr_id: i32) -> Self {
33        Self {
34            name: name.into(),
35            body_expr_id,
36            lookahead_assertion_id: -1,
37            is_exact_lookahead: false,
38        }
39    }
40}
41
42/// Decoded form of a `TagDispatch` expression.
43/// Port of `Grammar::Impl::TagDispatch`.
44#[derive(Debug, Clone, PartialEq, Eq)]
45pub struct TagDispatch {
46    /// `(tag, rule_id)` pairs — matching a tag dispatches to its rule.
47    pub tag_rule_pairs: Vec<(String, i32)>,
48    /// If true, EOS may be generated and stops the dispatch.
49    pub stop_eos: bool,
50    /// Strings that stop the dispatch (only when `stop_eos` is false).
51    pub stop_str: Vec<String>,
52    /// If true, the dispatch loops after dispatching.
53    pub loop_after_dispatch: bool,
54    /// Strings excluded by the tag dispatch.
55    pub excluded_str: Vec<String>,
56}
57
58/// Number of trailing scalar parameters in a `TagDispatch` payload:
59/// `stop_eos`, `stop_str_expr_id`, `loop_after_dispatch`,
60/// `excluded_str_expr_id`.
61const TAG_DISPATCH_EXTRA: usize = 4;
62
63/// The BNF grammar AST. Equivalent to xgrammar's `Grammar::Impl`.
64///
65/// The FSM acceleration fields (`complete_fsm`, `per_rule_fsms`,
66/// `per_rule_fsm_hashes`, `per_rule_fsm_new_state_ids`) are populated
67/// by the W3 grammar-functor passes (`GrammarFSMBuilder` /
68/// `GrammarFSMHasher`); they are `Default::default()` until then.
69#[derive(Debug, Clone, Default)]
70pub struct GrammarData {
71    /// Rules, indexed by `rule_id`.
72    rules: Vec<Rule>,
73    /// CSR payload store for all expressions.
74    expr_data: Vec<i32>,
75    /// CSR offsets: `expr_indptr[id]` is expr `id`'s start in `expr_data`.
76    expr_indptr: Vec<i32>,
77    /// Root rule id, or `-1` if unset.
78    root_rule_id: i32,
79    /// Rule ids permitted to match the empty string.
80    pub allow_empty_rule_ids: Vec<i32>,
81    /// Whether grammar optimization passes have run.
82    pub optimized: bool,
83    /// The single FSM holding every rule's states/edges, built by
84    /// `GrammarFSMBuilder`. Empty until the FSM builder runs.
85    pub complete_fsm: CompactFsm,
86    /// Per-rule FSM views into `complete_fsm`. `None` for a rule whose
87    /// body could not be expressed as an FSM. Empty until built.
88    pub per_rule_fsms: Vec<Option<CompactFsmWithStartEnd>>,
89    /// Per-rule structural hash of the rule's FSM, or `None` when the
90    /// rule could not be hashed. Set by `GrammarFSMHasher`.
91    pub per_rule_fsm_hashes: Vec<Option<u64>>,
92    /// Per-rule `(original_state_id, new_state_id)` mapping produced
93    /// alongside the hash by `GrammarFSMHasher`.
94    pub per_rule_fsm_new_state_ids: Vec<Vec<(i32, i32)>>,
95}
96
97impl GrammarData {
98    /// An empty grammar with no root.
99    pub fn new() -> Self {
100        Self {
101            root_rule_id: -1,
102            ..Default::default()
103        }
104    }
105
106    /// Construct directly from CSR parts — used by the grammar builder
107    /// (W2) and the deserializer.
108    pub fn from_parts(
109        rules: Vec<Rule>,
110        expr_data: Vec<i32>,
111        expr_indptr: Vec<i32>,
112        root_rule_id: i32,
113    ) -> Self {
114        Self {
115            rules,
116            expr_data,
117            expr_indptr,
118            root_rule_id,
119            ..Default::default()
120        }
121    }
122
123    /// Mutable access to a rule's name — used by `RootRuleRenamer`.
124    pub fn set_rule_name(&mut self, rule_id: i32, name: impl Into<String>) {
125        self.rules[rule_id as usize].name = name.into();
126    }
127
128    /// Borrowed view of expr `expr_id` as `(type, &mut data[..])`. Used
129    /// by `RepetitionNormalizer` to rewrite a `Repeat` payload in place.
130    pub fn set_expr_data(&mut self, expr_id: i32, index: usize, value: i32) {
131        let start = self.expr_indptr[expr_id as usize] as usize;
132        self.expr_data[start + 2 + index] = value;
133    }
134
135    /// Number of rules.
136    pub fn num_rules(&self) -> i32 {
137        self.rules.len() as i32
138    }
139
140    /// The rule with the given id. Panics if out of bounds (matches
141    /// the C++ `XGRAMMAR_DCHECK` contract).
142    pub fn rule(&self, rule_id: i32) -> &Rule {
143        &self.rules[rule_id as usize]
144    }
145
146    /// Mutable access to a rule.
147    pub fn rule_mut(&mut self, rule_id: i32) -> &mut Rule {
148        &mut self.rules[rule_id as usize]
149    }
150
151    /// All rules.
152    pub fn rules(&self) -> &[Rule] {
153        &self.rules
154    }
155
156    /// The root rule id (`-1` if unset).
157    pub fn root_rule_id(&self) -> i32 {
158        self.root_rule_id
159    }
160
161    /// Set the root rule id.
162    pub fn set_root_rule_id(&mut self, id: i32) {
163        self.root_rule_id = id;
164    }
165
166    /// The root rule. Panics if the root is unset/out of bounds.
167    pub fn root_rule(&self) -> &Rule {
168        &self.rules[self.root_rule_id as usize]
169    }
170
171    /// Number of expressions.
172    pub fn num_exprs(&self) -> i32 {
173        self.expr_indptr.len() as i32
174    }
175
176    /// Append a raw `(type, data[])` expression to the CSR store and
177    /// return its new expr id. Used by `LookaheadAssertionAnalyzer`,
178    /// which adds Sequence exprs to an already-built grammar.
179    pub fn append_expr(&mut self, kind: GrammarExprType, data: &[i32]) -> i32 {
180        self.expr_indptr.push(self.expr_data.len() as i32);
181        self.expr_data.push(kind as i32);
182        self.expr_data.push(data.len() as i32);
183        self.expr_data.extend_from_slice(data);
184        self.expr_indptr.len() as i32 - 1
185    }
186
187    /// Borrowed view of expression `expr_id`. The CSR slot is
188    /// `[type, data_len, data0, …]`; this returns `(type, &data[..])`.
189    pub fn expr(&self, expr_id: i32) -> GrammarExpr<'_> {
190        let start = self.expr_indptr[expr_id as usize] as usize;
191        let kind = GrammarExprType::from_i32(self.expr_data[start])
192            .expect("corrupt grammar: unknown GrammarExprType tag");
193        let data_len = self.expr_data[start + 1] as usize;
194        let data = &self.expr_data[start + 2..start + 2 + data_len];
195        GrammarExpr { kind, data }
196    }
197
198    /// Decode a `ByteString` expression to its `String`.
199    pub fn byte_string(&self, expr_id: i32) -> String {
200        let e = self.expr(expr_id);
201        e.data.iter().map(|&b| (b as u8) as char).collect()
202    }
203
204    /// Decode a `TagDispatch` expression. Panics if `expr_id` is not a
205    /// tag dispatch (matches the C++ `XGRAMMAR_DCHECK`).
206    pub fn tag_dispatch(&self, expr_id: i32) -> TagDispatch {
207        let e = self.expr(expr_id);
208        assert_eq!(
209            e.kind,
210            GrammarExprType::TagDispatch,
211            "expr {expr_id} is not a TagDispatch"
212        );
213        let n = e.len();
214        assert!(n >= TAG_DISPATCH_EXTRA, "TagDispatch payload too short");
215        let pair_end = n - TAG_DISPATCH_EXTRA;
216
217        let mut tag_rule_pairs = Vec::with_capacity(pair_end / 2);
218        let mut i = 0;
219        while i < pair_end {
220            let tag = self.byte_string(e[i]);
221            let rule_id = e[i + 1];
222            tag_rule_pairs.push((tag, rule_id));
223            i += 2;
224        }
225
226        let stop_eos = e[pair_end] != 0;
227        let stop_str = self.choices_of_byte_strings(e[pair_end + 1]);
228        let loop_after_dispatch = e[pair_end + 2] != 0;
229        let excluded_str = self.choices_of_byte_strings(e[pair_end + 3]);
230
231        TagDispatch {
232            tag_rule_pairs,
233            stop_eos,
234            stop_str,
235            loop_after_dispatch,
236            excluded_str,
237        }
238    }
239
240    /// Decode a `Choices` expression whose branches are all
241    /// `ByteString`s into a `Vec<String>`.
242    fn choices_of_byte_strings(&self, expr_id: i32) -> Vec<String> {
243        let e = self.expr(expr_id);
244        assert_eq!(e.kind, GrammarExprType::Choices, "expected a Choices expr");
245        e.data.iter().map(|&id| self.byte_string(id)).collect()
246    }
247}
248
249#[cfg(test)]
250mod tests {
251    use super::*;
252
253    /// Build a tiny grammar by hand: one rule whose body is the byte
254    /// string "ab". CSR slot for a ByteString: `[type=0, len=2, 'a','b']`.
255    fn byte_string_grammar() -> GrammarData {
256        let expr_data = vec![
257            GrammarExprType::ByteString as i32,
258            2,
259            b'a' as i32,
260            b'b' as i32,
261        ];
262        let expr_indptr = vec![0];
263        let rules = vec![Rule::new("root", 0)];
264        GrammarData::from_parts(rules, expr_data, expr_indptr, 0)
265    }
266
267    #[test]
268    fn expr_view_decodes_csr() {
269        let g = byte_string_grammar();
270        assert_eq!(g.num_rules(), 1);
271        assert_eq!(g.num_exprs(), 1);
272        let e = g.expr(0);
273        assert_eq!(e.kind, GrammarExprType::ByteString);
274        assert_eq!(e.data, &[b'a' as i32, b'b' as i32]);
275    }
276
277    #[test]
278    fn byte_string_decodes() {
279        assert_eq!(byte_string_grammar().byte_string(0), "ab");
280    }
281
282    #[test]
283    fn root_rule_resolves() {
284        let g = byte_string_grammar();
285        assert_eq!(g.root_rule_id(), 0);
286        assert_eq!(g.root_rule().name, "root");
287        assert_eq!(g.root_rule().body_expr_id, 0);
288    }
289}
metrale_grammar/grammar/data.rs

metrale_grammar/grammar/
data.rs
