metrale_grammar/grammar/
builder.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//
3// GrammarBuilder — the mutable construction API for the BNF AST.
4// Port of `GrammarBuilder` from xgrammar `cpp/grammar_builder.h`.
5//
6// The builder accumulates rules and expressions, writing expressions
7// into the CSR `expr_data`/`expr_indptr` vectors, and finally emits a
8// finished [`GrammarData`] via [`GrammarData::from_parts`].
9//
10// Layout reminder: each CSR slot is `[type, data_len, data0, data1, …]`.
11//
12// This file holds the core struct + expression-handling methods. Rule
13// handling and the error type live in `builder_rules.rs`; tests in
14// `builder_tests.rs` — split to keep each file under the 250-line cap.
15
16use std::cell::RefCell;
17use std::collections::HashMap;
18
19use super::data::{GrammarData, Rule};
20use super::expr::{GrammarExpr, GrammarExprType};
21
22#[path = "builder_rules.rs"]
23mod rules;
24#[cfg(test)]
25#[path = "builder_tests.rs"]
26mod tests;
27
28pub use rules::BuilderError;
29
30/// One element of a character class: an inclusive `[lower, upper]`
31/// codepoint range. Port of `GrammarBuilder::CharacterClassElement`.
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub struct CharacterClassElement {
34    /// Inclusive lower bound (unicode codepoint).
35    pub lower: i32,
36    /// Inclusive upper bound (unicode codepoint).
37    pub upper: i32,
38}
39
40impl CharacterClassElement {
41    /// A range `[lower, upper]`.
42    pub fn new(lower: i32, upper: i32) -> Self {
43        Self { lower, upper }
44    }
45}
46
47/// Decoded tag-dispatch description handed to [`GrammarBuilder::add_tag_dispatch`].
48/// Mirrors [`super::data::TagDispatch`] but is the *input* form: rule ids
49/// (not tag byte-string expr ids) are stored directly.
50#[derive(Debug, Clone, PartialEq, Eq, Default)]
51pub struct TagDispatchSpec {
52    /// `(tag, rule_id)` pairs.
53    pub tag_rule_pairs: Vec<(String, i32)>,
54    /// If true, EOS may stop the dispatch.
55    pub stop_eos: bool,
56    /// Strings that stop the dispatch.
57    pub stop_str: Vec<String>,
58    /// If true, the dispatch loops after dispatching.
59    pub loop_after_dispatch: bool,
60    /// Strings excluded by the tag dispatch.
61    pub excluded_str: Vec<String>,
62}
63
64/// Helper class to build a BNF grammar.
65///
66/// Equivalent to xgrammar's `GrammarBuilder`. Holds the in-progress
67/// rule list and CSR expression store, plus a name→id index.
68#[derive(Debug, Default)]
69pub struct GrammarBuilder {
70    rules: Vec<Rule>,
71    expr_data: Vec<i32>,
72    expr_indptr: Vec<i32>,
73    rule_name_to_id: HashMap<String, i32>,
74    /// Cache of the next suffix index to probe for each `name_hint`,
75    /// used by [`GrammarBuilder::get_new_rule_name`]. Port of upstream
76    /// `next_cnt_per_hint_` (xgrammar commit 96ae88b). Without this the
77    /// per-call probe restarts from `_1` every time, making repeated
78    /// calls with the same hint O(N) each (O(N^2) overall at high tool
79    /// counts). The cache is purely an optimization: every candidate is
80    /// still validated against `rule_name_to_id`, so the chosen name is
81    /// identical to the un-cached probe. `RefCell` keeps the public
82    /// `&self` signature stable.
83    next_cnt_per_hint: RefCell<HashMap<String, i32>>,
84}
85
86impl GrammarBuilder {
87    /// Create an empty builder.
88    pub fn new() -> Self {
89        Self::default()
90    }
91
92    /* ---- crate-internal field access for the `rules` submodule ---- */
93
94    pub(super) fn rules_vec(&self) -> &[Rule] {
95        &self.rules
96    }
97    pub(super) fn rules_vec_mut(&mut self) -> &mut Vec<Rule> {
98        &mut self.rules
99    }
100    pub(super) fn rule_name_to_id(&self) -> &HashMap<String, i32> {
101        &self.rule_name_to_id
102    }
103    pub(super) fn rule_name_to_id_mut(&mut self) -> &mut HashMap<String, i32> {
104        &mut self.rule_name_to_id
105    }
106    pub(super) fn next_cnt_per_hint(&self) -> &RefCell<HashMap<String, i32>> {
107        &self.next_cnt_per_hint
108    }
109
110    /// Finish building and return the [`GrammarData`], setting the root
111    /// rule to the rule with `root_rule_name`.
112    ///
113    /// Returns `Err` if the named rule is not present.
114    pub fn get(self, root_rule_name: &str) -> Result<GrammarData, BuilderError> {
115        let root_rule_id = self.get_rule_id(root_rule_name);
116        if root_rule_id == -1 {
117            return Err(BuilderError::RootRuleNotFound(root_rule_name.to_string()));
118        }
119        self.get_by_id(root_rule_id)
120    }
121
122    /// Finish building and return the [`GrammarData`], setting the root
123    /// rule to `root_rule_id`.
124    ///
125    /// Returns `Err` if `root_rule_id` is out of bounds.
126    pub fn get_by_id(self, root_rule_id: i32) -> Result<GrammarData, BuilderError> {
127        if root_rule_id < 0 || root_rule_id >= self.rules.len() as i32 {
128            return Err(BuilderError::RootRuleOutOfBounds(root_rule_id));
129        }
130        Ok(GrammarData::from_parts(
131            self.rules,
132            self.expr_data,
133            self.expr_indptr,
134            root_rule_id,
135        ))
136    }
137
138    /* ****************** GrammarExpr handling ****************** */
139
140    /// Append a raw `(type, data[])` expression to the CSR store and
141    /// return its expr id.
142    pub fn add_grammar_expr(&mut self, kind: GrammarExprType, data: &[i32]) -> i32 {
143        self.expr_indptr.push(self.expr_data.len() as i32);
144        self.expr_data.push(kind as i32);
145        self.expr_data.push(data.len() as i32);
146        self.expr_data.extend_from_slice(data);
147        self.expr_indptr.len() as i32 - 1
148    }
149
150    /// Add a `ByteString` expr from raw bytes (each `0..=255`).
151    pub fn add_byte_string_bytes(&mut self, bytes: &[i32]) -> i32 {
152        self.add_grammar_expr(GrammarExprType::ByteString, bytes)
153    }
154
155    /// Add a `ByteString` expr from a Rust string (UTF-8 bytes).
156    pub fn add_byte_string(&mut self, s: &str) -> i32 {
157        let bytes: Vec<i32> = s.bytes().map(|b| b as i32).collect();
158        self.add_grammar_expr(GrammarExprType::ByteString, &bytes)
159    }
160
161    /// Build the CSR payload of a character class:
162    /// `[is_negative, lower0, upper0, lower1, upper1, …]`.
163    fn char_class_data(elements: &[CharacterClassElement], is_negative: bool) -> Vec<i32> {
164        let mut data = Vec::with_capacity(1 + elements.len() * 2);
165        data.push(is_negative as i32);
166        for r in elements {
167            data.push(r.lower);
168            data.push(r.upper);
169        }
170        data
171    }
172
173    /// Add a `CharacterClass` expr, e.g. `[a-z]` / `[^a-z]`.
174    pub fn add_character_class(
175        &mut self,
176        elements: &[CharacterClassElement],
177        is_negative: bool,
178    ) -> i32 {
179        let data = Self::char_class_data(elements, is_negative);
180        self.add_grammar_expr(GrammarExprType::CharacterClass, &data)
181    }
182
183    /// Add a `CharacterClassStar` expr, e.g. `[a-z]*`.
184    pub fn add_character_class_star(
185        &mut self,
186        elements: &[CharacterClassElement],
187        is_negative: bool,
188    ) -> i32 {
189        let data = Self::char_class_data(elements, is_negative);
190        self.add_grammar_expr(GrammarExprType::CharacterClassStar, &data)
191    }
192
193    /// Add an `EmptyStr` expr (the empty string `""`).
194    pub fn add_empty_str(&mut self) -> i32 {
195        self.add_grammar_expr(GrammarExprType::EmptyStr, &[])
196    }
197
198    /// Add a `RuleRef` expr pointing at `rule_id`.
199    pub fn add_rule_ref(&mut self, rule_id: i32) -> i32 {
200        self.add_grammar_expr(GrammarExprType::RuleRef, &[rule_id])
201    }
202
203    /// Add a `Sequence` expr (concatenation of sub-expr ids).
204    pub fn add_sequence(&mut self, elements: &[i32]) -> i32 {
205        self.add_grammar_expr(GrammarExprType::Sequence, elements)
206    }
207
208    /// Add a `Choices` expr (alternation of sub-expr ids).
209    pub fn add_choices(&mut self, choices: &[i32]) -> i32 {
210        self.add_grammar_expr(GrammarExprType::Choices, choices)
211    }
212
213    /// Add a `Repeat` expr `[rule_id, min, max]`.
214    pub fn add_repeat(&mut self, ref_rule_id: i32, min_repeat: i32, max_repeat: i32) -> i32 {
215        self.add_grammar_expr(
216            GrammarExprType::Repeat,
217            &[ref_rule_id, min_repeat, max_repeat],
218        )
219    }
220
221    /// Add a `TagDispatch` expr. Tag strings and stop/exclude string
222    /// lists are interned as `ByteString`/`Choices` sub-exprs first,
223    /// matching the CSR layout decoded by [`GrammarData::tag_dispatch`].
224    pub fn add_tag_dispatch(&mut self, spec: &TagDispatchSpec) -> i32 {
225        let mut data: Vec<i32> = Vec::with_capacity(spec.tag_rule_pairs.len() * 2 + 4);
226        for (tag, rule_id) in &spec.tag_rule_pairs {
227            let tag_expr = self.add_byte_string(tag);
228            data.push(tag_expr);
229            data.push(*rule_id);
230        }
231        data.push(spec.stop_eos as i32);
232        let stop_ids: Vec<i32> = spec
233            .stop_str
234            .iter()
235            .map(|s| self.add_byte_string(s))
236            .collect();
237        data.push(self.add_choices(&stop_ids));
238        data.push(spec.loop_after_dispatch as i32);
239        let excl_ids: Vec<i32> = spec
240            .excluded_str
241            .iter()
242            .map(|s| self.add_byte_string(s))
243            .collect();
244        data.push(self.add_choices(&excl_ids));
245        self.add_grammar_expr(GrammarExprType::TagDispatch, &data)
246    }
247
248    /// Number of expressions stored so far.
249    pub fn num_grammar_exprs(&self) -> i32 {
250        self.expr_indptr.len() as i32
251    }
252
253    /// Borrowed view of expr `expr_id`. Panics if out of bounds — the
254    /// builder is internal and ids it returns are always valid.
255    pub fn get_grammar_expr(&self, expr_id: i32) -> GrammarExpr<'_> {
256        let start = self.expr_indptr[expr_id as usize] as usize;
257        let kind = GrammarExprType::from_i32(self.expr_data[start])
258            .expect("corrupt builder state: unknown GrammarExprType tag");
259        let data_len = self.expr_data[start + 1] as usize;
260        let data = &self.expr_data[start + 2..start + 2 + data_len];
261        GrammarExpr { kind, data }
262    }
263}
metrale_grammar/grammar/builder.rs

metrale_grammar/grammar/
builder.rs
