Skip to main content

ferrox_models/grammar/
element.rs

1//! The compiled form of a GBNF grammar: elements, rules, and a cursor
2//! into them.
3//!
4//! Ported from `llama_gretype` / `llama_grammar_element` in llama.cpp's
5//! `src/llama-grammar.h`. The one representational change is the cursor:
6//! llama.cpp walks a grammar with `const llama_grammar_element *` raw
7//! pointers into the rule vectors, which is why `llama_grammar_clone_impl`
8//! has to rewrite every stack entry after a copy. [`RulePos`] is the same
9//! cursor expressed as `(rule, index)`, so a [`Grammar`](super::Grammar)
10//! clones for free.
11
12/// Element kinds, matching `enum llama_gretype` one for one.
13///
14/// The discriminants are llama.cpp's, so a serialized grammar could be
15/// compared against upstream without a translation table.
16#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
17#[repr(u8)]
18pub enum GreType {
19    /// End of a rule definition.
20    End = 0,
21    /// Start of an alternate definition for a rule.
22    Alt = 1,
23    /// Non-terminal: a reference to another rule, by rule id.
24    RuleRef = 2,
25    /// Terminal: a single Unicode code point.
26    Char = 3,
27    /// Inverted character set: `[^a]`, `[^a-b]`, `[^abc]`.
28    CharNot = 4,
29    /// Modifies a preceding `Char` or `CharAlt` into an inclusive range
30    /// (`[a-z]`).
31    CharRngUpper = 5,
32    /// Modifies a preceding `Char` or `CharRngUpper` by adding another
33    /// alternative to match (`[ab]`, `[a-zA]`).
34    CharAlt = 6,
35    /// Any character (`.`).
36    CharAny = 7,
37    /// Terminal: a token id (`<[42]>`).
38    Token = 8,
39    /// Inverted token (`!<[42]>`).
40    TokenNot = 9,
41}
42
43impl GreType {
44    /// True for the element kinds that consume a character.
45    ///
46    /// `llama_grammar_is_char_element`, used only by the printer and by
47    /// rule validation.
48    pub fn is_char_element(self) -> bool {
49        matches!(
50            self,
51            GreType::Char
52                | GreType::CharNot
53                | GreType::CharAlt
54                | GreType::CharRngUpper
55                | GreType::CharAny
56        )
57    }
58
59    /// True for the element kinds a parse stack may legally rest on.
60    ///
61    /// `llama_grammar_advance_stack` aborts if a stack top is anything
62    /// else; we return an error instead (see [`super::GrammarError`]).
63    pub fn is_stack_terminal(self) -> bool {
64        matches!(
65            self,
66            GreType::Char
67                | GreType::CharNot
68                | GreType::CharAny
69                | GreType::Token
70                | GreType::TokenNot
71        )
72    }
73}
74
75/// One grammar element: a kind plus a code point, rule id, or token id.
76#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
77pub struct GrammarElement {
78    pub gtype: GreType,
79    pub value: u32,
80}
81
82impl GrammarElement {
83    pub const fn new(gtype: GreType, value: u32) -> Self {
84        Self { gtype, value }
85    }
86
87    /// `llama_grammar_is_end_of_sequence`: true iff this position ends one
88    /// of the alternate definitions of a rule.
89    pub fn is_end_of_sequence(self) -> bool {
90        matches!(self.gtype, GreType::End | GreType::Alt)
91    }
92}
93
94/// One rule: a flat element sequence, alternates separated by [`GreType::Alt`],
95/// always terminated by [`GreType::End`].
96pub type GrammarRule = Vec<GrammarElement>;
97
98/// A cursor into the compiled rule table: `rules[rule][index]`.
99///
100/// Ordering is `(rule, index)` lexicographic, which stands in for
101/// llama.cpp's pointer comparison in `llama_grammar_advance_stack`'s
102/// `seen` set. Any total order does the job there; only distinctness
103/// matters.
104#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
105pub struct RulePos {
106    pub rule: u32,
107    pub index: u32,
108}
109
110impl RulePos {
111    pub const fn new(rule: u32, index: u32) -> Self {
112        Self { rule, index }
113    }
114
115    /// `pos + 1`: the next element in the same rule.
116    pub const fn next(self) -> Self {
117        Self {
118            rule: self.rule,
119            index: self.index + 1,
120        }
121    }
122
123    /// `pos + n`.
124    pub const fn advance(self, n: u32) -> Self {
125        Self {
126            rule: self.rule,
127            index: self.index + n,
128        }
129    }
130}
131
132/// A pushdown stack: cursors, innermost last. `stack.last()` is the
133/// position the next character has to satisfy.
134pub type GrammarStack = Vec<RulePos>;
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139
140    #[test]
141    fn end_and_alt_terminate_a_sequence() {
142        assert!(GrammarElement::new(GreType::End, 0).is_end_of_sequence());
143        assert!(GrammarElement::new(GreType::Alt, 0).is_end_of_sequence());
144        assert!(!GrammarElement::new(GreType::Char, b'a' as u32).is_end_of_sequence());
145        assert!(!GrammarElement::new(GreType::RuleRef, 3).is_end_of_sequence());
146    }
147
148    #[test]
149    fn char_range_modifiers_are_char_elements_but_not_stack_terminals() {
150        // llama.cpp's `advance_stack` aborts on a stack resting on
151        // CHAR_ALT or CHAR_RNG_UPPER; both are still "char elements" to
152        // the printer. The two predicates genuinely differ.
153        assert!(GreType::CharAlt.is_char_element());
154        assert!(GreType::CharRngUpper.is_char_element());
155        assert!(!GreType::CharAlt.is_stack_terminal());
156        assert!(!GreType::CharRngUpper.is_stack_terminal());
157        assert!(GreType::CharAny.is_char_element());
158        assert!(GreType::CharAny.is_stack_terminal());
159        assert!(!GreType::Token.is_char_element());
160        assert!(GreType::Token.is_stack_terminal());
161    }
162
163    #[test]
164    fn discriminants_match_llama_cpp() {
165        assert_eq!(GreType::End as u8, 0);
166        assert_eq!(GreType::Alt as u8, 1);
167        assert_eq!(GreType::RuleRef as u8, 2);
168        assert_eq!(GreType::Char as u8, 3);
169        assert_eq!(GreType::CharNot as u8, 4);
170        assert_eq!(GreType::CharRngUpper as u8, 5);
171        assert_eq!(GreType::CharAlt as u8, 6);
172        assert_eq!(GreType::CharAny as u8, 7);
173        assert_eq!(GreType::Token as u8, 8);
174        assert_eq!(GreType::TokenNot as u8, 9);
175    }
176
177    #[test]
178    fn cursor_orders_by_rule_then_index() {
179        assert!(RulePos::new(0, 5) < RulePos::new(1, 0));
180        assert!(RulePos::new(1, 0) < RulePos::new(1, 1));
181        assert_eq!(RulePos::new(2, 3).next(), RulePos::new(2, 4));
182        assert_eq!(RulePos::new(2, 3).advance(2), RulePos::new(2, 5));
183    }
184}