metrale_grammar/grammar/expr.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2//
3// GrammarExpr — the node type of the BNF grammar AST.
4// Port of the `GrammarExprType` enum + `GrammarExpr` view from
5// xgrammar `cpp/grammar_impl.h`.
6
7/// The type of a grammar expression node.
8///
9/// Each `GrammarExpr` is a `(type, data[])` pair; the `data` layout
10/// depends on the type — see each variant's doc.
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
12#[repr(i32)]
13pub enum GrammarExprType {
14 /// `[byte0, byte1, …]` — a string of bytes (0-255), UTF-8 capable.
15 ByteString = 0,
16 /// `[is_negative, lower0, upper0, lower1, upper1, …]` — a set of
17 /// unicode codepoint ranges, e.g. `[a-z]`, negatable `[^a-z]`.
18 CharacterClass = 1,
19 /// A star quantifier over a character class, e.g. `[a-z]*`.
20 CharacterClassStar = 2,
21 /// `[]` — the empty string `""`.
22 EmptyStr = 3,
23 /// `[rule_id]` — a reference to another rule.
24 RuleRef = 4,
25 /// `[expr_id0, expr_id1, …]` — concatenation of sub-expressions.
26 Sequence = 5,
27 /// `[expr_id0, expr_id1, …]` — alternation; any branch may match.
28 Choices = 6,
29 /// `[tag_expr0, rule_id0, …, stop_eos, stop_str_expr_id,
30 /// loop_after_dispatch, excluded_str_expr_id]` — tag dispatch.
31 TagDispatch = 7,
32 /// `[rule_id, min_repeat_count, max_repeat_count]`.
33 Repeat = 8,
34}
35
36impl GrammarExprType {
37 /// Reconstruct from the raw `i32` stored in the CSR data vector.
38 pub fn from_i32(v: i32) -> Option<Self> {
39 Some(match v {
40 0 => Self::ByteString,
41 1 => Self::CharacterClass,
42 2 => Self::CharacterClassStar,
43 3 => Self::EmptyStr,
44 4 => Self::RuleRef,
45 5 => Self::Sequence,
46 6 => Self::Choices,
47 7 => Self::TagDispatch,
48 8 => Self::Repeat,
49 _ => return None,
50 })
51 }
52}
53
54/// A borrowed view of one grammar expression: its type and a slice of
55/// its `i32` payload. Equivalent to xgrammar's `GrammarExpr` struct,
56/// but lifetime-bound to the owning [`super::data::GrammarData`]
57/// instead of holding a raw pointer.
58#[derive(Debug, Clone, Copy)]
59pub struct GrammarExpr<'a> {
60 pub kind: GrammarExprType,
61 pub data: &'a [i32],
62}
63
64impl<'a> GrammarExpr<'a> {
65 /// Number of elements in the payload.
66 pub fn len(&self) -> usize {
67 self.data.len()
68 }
69
70 /// Whether the payload is empty.
71 pub fn is_empty(&self) -> bool {
72 self.data.is_empty()
73 }
74}
75
76impl std::ops::Index<usize> for GrammarExpr<'_> {
77 type Output = i32;
78 fn index(&self, i: usize) -> &i32 {
79 &self.data[i]
80 }
81}
82
83#[cfg(test)]
84mod tests {
85 use super::*;
86
87 #[test]
88 fn type_roundtrips() {
89 for v in 0..=8 {
90 let t = GrammarExprType::from_i32(v).unwrap();
91 assert_eq!(t as i32, v);
92 }
93 assert!(GrammarExprType::from_i32(9).is_none());
94 assert!(GrammarExprType::from_i32(-1).is_none());
95 }
96}