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}