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}