metrale_grammar/grammar/
builder_rules.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//
3// GrammarBuilder — rule-handling methods and the builder error type.
4// Split out of `builder.rs` to keep each file under the 250-line cap.
5// Port of the rule-handling half of `GrammarBuilder`
6// (xgrammar `cpp/grammar_builder.h`).
7
8use super::GrammarBuilder;
9use crate::grammar::data::Rule;
10
11impl GrammarBuilder {
12    /// Add a fully-formed rule and return its id.
13    ///
14    /// Returns `Err` if a rule of the same name already exists.
15    pub fn add_rule(&mut self, rule: Rule) -> Result<i32, BuilderError> {
16        if self.rule_name_to_id().contains_key(&rule.name) {
17            return Err(BuilderError::DuplicateRule(rule.name.clone()));
18        }
19        let id = self.rules_vec().len() as i32;
20        let name = rule.name.clone();
21        self.rules_vec_mut().push(rule);
22        self.rule_name_to_id_mut().insert(name, id);
23        Ok(id)
24    }
25
26    /// Add a rule with the given `name` and `body_expr_id`.
27    pub fn add_rule_named(
28        &mut self,
29        name: impl Into<String>,
30        body_expr_id: i32,
31    ) -> Result<i32, BuilderError> {
32        self.add_rule(Rule::new(name, body_expr_id))
33    }
34
35    /// Add a rule with a body, deriving a unique name from `name_hint`.
36    pub fn add_rule_with_hint(
37        &mut self,
38        name_hint: &str,
39        body_expr_id: i32,
40    ) -> Result<i32, BuilderError> {
41        let name = self.get_new_rule_name(name_hint);
42        self.add_rule(Rule::new(name, body_expr_id))
43    }
44
45    /// Number of rules.
46    pub fn num_rules(&self) -> i32 {
47        self.rules_vec().len() as i32
48    }
49
50    /// The rule with the given id.
51    pub fn get_rule(&self, rule_id: i32) -> &Rule {
52        &self.rules_vec()[rule_id as usize]
53    }
54
55    /// Add a body-less rule (`body_expr_id = -1`). The body must be set
56    /// later with [`GrammarBuilder::update_rule_body`]. Useful when the
57    /// rule id is needed to build the body (recursive rules).
58    pub fn add_empty_rule(&mut self, name: impl Into<String>) -> Result<i32, BuilderError> {
59        self.add_rule(Rule::new(name, -1))
60    }
61
62    /// Add a body-less rule with a unique name derived from `name_hint`.
63    pub fn add_empty_rule_with_hint(&mut self, name_hint: &str) -> Result<i32, BuilderError> {
64        let name = self.get_new_rule_name(name_hint);
65        self.add_rule(Rule::new(name, -1))
66    }
67
68    /// Set the body of rule `rule_id`.
69    ///
70    /// Returns `Err` if `rule_id` is out of range.
71    pub fn update_rule_body(
72        &mut self,
73        rule_id: i32,
74        body_expr_id: i32,
75    ) -> Result<(), BuilderError> {
76        if rule_id < 0 || rule_id >= self.rules_vec().len() as i32 {
77            return Err(BuilderError::RuleIdOutOfRange(rule_id));
78        }
79        self.rules_vec_mut()[rule_id as usize].body_expr_id = body_expr_id;
80        Ok(())
81    }
82
83    /// Set the body of the rule named `rule_name`.
84    pub fn update_rule_body_named(
85        &mut self,
86        rule_name: &str,
87        body_expr_id: i32,
88    ) -> Result<(), BuilderError> {
89        let rule_id = self.get_rule_id(rule_name);
90        if rule_id == -1 {
91            return Err(BuilderError::RuleNameNotFound(rule_name.to_string()));
92        }
93        self.update_rule_body(rule_id, body_expr_id)
94    }
95
96    /// Set the lookahead-assertion expr id of rule `rule_id`. An id of
97    /// `-1` means no assertion.
98    pub fn update_lookahead_assertion(
99        &mut self,
100        rule_id: i32,
101        lookahead_assertion_id: i32,
102    ) -> Result<(), BuilderError> {
103        if rule_id < 0 || rule_id >= self.rules_vec().len() as i32 {
104            return Err(BuilderError::RuleIdOutOfRange(rule_id));
105        }
106        self.rules_vec_mut()[rule_id as usize].lookahead_assertion_id = lookahead_assertion_id;
107        Ok(())
108    }
109
110    /// Set the lookahead-assertion expr id of the rule named `rule_name`.
111    pub fn update_lookahead_assertion_named(
112        &mut self,
113        rule_name: &str,
114        lookahead_assertion_id: i32,
115    ) -> Result<(), BuilderError> {
116        let rule_id = self.get_rule_id(rule_name);
117        if rule_id == -1 {
118            return Err(BuilderError::RuleNameNotFound(rule_name.to_string()));
119        }
120        self.update_lookahead_assertion(rule_id, lookahead_assertion_id)
121    }
122
123    /// Mark the lookahead assertion of rule `rule_id` as exact.
124    pub fn update_lookahead_exact(
125        &mut self,
126        rule_id: i32,
127        is_exact: bool,
128    ) -> Result<(), BuilderError> {
129        if rule_id < 0 || rule_id >= self.rules_vec().len() as i32 {
130            return Err(BuilderError::RuleIdOutOfRange(rule_id));
131        }
132        self.rules_vec_mut()[rule_id as usize].is_exact_lookahead = is_exact;
133        Ok(())
134    }
135
136    /// Find an unused rule name starting from `name_hint`, appending an
137    /// integer suffix (`_1`, `_2`, …) on collision.
138    ///
139    /// The first probed suffix is read from `next_cnt_per_hint` (a cache
140    /// of the last suffix reached for this hint) rather than restarting
141    /// from `1` every call. This makes repeated calls with the same hint
142    /// amortized O(1) instead of O(N) — see the field doc on
143    /// [`GrammarBuilder`]. Behavior is unchanged: every candidate is
144    /// still verified absent from `rule_name_to_id`, and the cache only
145    /// advances (never skips a still-free lower suffix, because suffixes
146    /// for a given hint are only ever consumed in increasing order).
147    pub fn get_new_rule_name(&self, name_hint: &str) -> String {
148        if !self.rule_name_to_id().contains_key(name_hint) {
149            return name_hint.to_string();
150        }
151        let mut cache = self.next_cnt_per_hint().borrow_mut();
152        let cnt = cache.entry(name_hint.to_string()).or_insert(1);
153        loop {
154            let candidate = format!("{name_hint}_{cnt}");
155            if !self.rule_name_to_id().contains_key(&candidate) {
156                return candidate;
157            }
158            *cnt += 1;
159        }
160    }
161
162    /// Rule id of the rule named `name`, or `-1` if absent.
163    pub fn get_rule_id(&self, name: &str) -> i32 {
164        self.rule_name_to_id().get(name).copied().unwrap_or(-1)
165    }
166}
167
168/// Errors emitted by [`GrammarBuilder`].
169#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
170pub enum BuilderError {
171    /// A rule with this name was added twice.
172    #[error("rule \"{0}\" is added multiple times")]
173    DuplicateRule(String),
174    /// The requested root rule name does not exist.
175    #[error("the root rule with name \"{0}\" is not found")]
176    RootRuleNotFound(String),
177    /// The requested root rule id is out of bounds.
178    #[error("the root rule id {0} is out of bound")]
179    RootRuleOutOfBounds(i32),
180    /// A rule id passed to an update method is out of range.
181    #[error("rule id {0} is out of range")]
182    RuleIdOutOfRange(i32),
183    /// A rule name passed to an update method is not found.
184    #[error("rule \"{0}\" is not found")]
185    RuleNameNotFound(String),
186}