metrale_grammar/compiler/
compiler.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//
3// GrammarCompiler — port of `class GrammarCompiler` / `GrammarCompiler
4// ::Impl` from `cpp/grammar_compiler.cc`.
5//
6// A tokenizer-bound cache that produces `CompiledGrammar`s, avoiding
7// redundant preprocessing of grammars / schemas seen before.
8
9use crate::grammar::functor::{GrammarNormalizer, GrammarOptimizer};
10use crate::grammar::{GrammarData, parse_ebnf};
11use crate::schema::{SchemaConverterOptions, builtin_json_grammar_ebnf, json_schema_to_ebnf};
12use crate::structural_tag::structural_tag_to_grammar;
13use crate::tokenizer::TokenizerInfo;
14
15use super::compile::compile_optimized_grammar;
16use super::compiled_grammar::CompiledGrammar;
17use super::grammar_cache::{GrammarCache, UNLIMITED_BYTES};
18use super::mask_snapshot::{self, SnapshotError, SnapshotIdentity};
19use super::rule_cache::{RuleLevelCache, UNLIMITED_SIZE};
20
21/// An error produced while compiling a grammar, schema or tag.
22#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
23pub enum CompileError {
24    /// The EBNF source failed to parse.
25    #[error("failed to parse grammar: {0}")]
26    Grammar(String),
27    /// The JSON schema failed to convert or parse.
28    #[error("failed to compile JSON schema: {0}")]
29    Schema(String),
30    /// A structural tag failed to compile.
31    #[error("failed to compile structural tag: {0}")]
32    StructuralTag(String),
33}
34
35/// The cache key — exactly the request parameters that determine the
36/// compiled output. Port of the C++ `GrammarCompilerCacheKeys` union.
37#[derive(Debug, Clone, PartialEq, Eq, Hash)]
38enum CacheKey {
39    /// `compile_builtin_json_grammar`.
40    BuiltinJson,
41    /// `compile_grammar_from_ebnf(ebnf, root_rule)`.
42    Ebnf { ebnf: String, root_rule: String },
43    /// `compile_json_schema(schema, options...)`.
44    Schema {
45        schema: String,
46        any_whitespace: bool,
47        indent: Option<i32>,
48        separators: Option<(String, String)>,
49        strict_mode: bool,
50        max_whitespace_cnt: Option<i32>,
51    },
52    /// `compile_structural_tag(json)`.
53    StructuralTag(String),
54}
55
56/// A tokenizer-bound grammar compiler with an internal cache.
57///
58/// Each compiler instance is tied to one [`TokenizerInfo`]; use a
59/// separate compiler per vocabulary. Compilation results are cached
60/// (when `cache_enabled`) keyed by the request parameters.
61pub struct GrammarCompiler {
62    tokenizer_info: TokenizerInfo,
63    max_threads: usize,
64    cache_enabled: bool,
65    cache_limit_bytes: i64,
66    /// Tier-1 compiled-grammar cache, LRU-bounded by `cache_limit_bytes`.
67    ///
68    /// Was a bare `DashMap` with no eviction: `CacheKey::Schema` is keyed by
69    /// the full tool-schema text, so agentic traffic minted a permanent entry
70    /// per request (~100 MB/min under BFCL; issue #368).
71    cache: GrammarCache<CacheKey>,
72    /// Cross-grammar adaptive-token-mask cache (Tier 2, upstream commit
73    /// `bfb2a79`). `Some` only when `cache_enabled` — a rule
74    /// structurally identical to one compiled for a previous request
75    /// reuses its masks. Shared (cloned `Arc` inside) into every
76    /// `CompiledGrammar` this compiler produces. Port of the C++
77    /// `GrammarCompiler::Impl::rule_level_cache_`.
78    rule_cache: Option<RuleLevelCache>,
79}
80
81impl GrammarCompiler {
82    /// Construct a compiler bound to `tokenizer_info`.
83    ///
84    /// * `max_threads` — bounds explicit top-k mask prewarming. Lazy
85    ///   matcher lookups still compute only the requested mask. Must be >= 1.
86    /// * `cache_enabled` — whether to cache compiled grammars.
87    /// * `cache_limit_bytes` — memory budget, ENFORCED by LRU eviction on
88    ///   both tiers. `-1` means unlimited, which is now an explicit opt-in
89    ///   rather than the serve's default: unbounded was the shape of the leak
90    ///   in issue #368. The budget splits as upstream does — ~1/3 to the
91    ///   rule-level cache, the remaining ~2/3 to this grammar-level one.
92    ///
93    /// Port of the C++ `GrammarCompiler` constructor.
94    pub fn new(
95        tokenizer_info: TokenizerInfo,
96        max_threads: usize,
97        cache_enabled: bool,
98        cache_limit_bytes: i64,
99    ) -> Self {
100        assert!(max_threads >= 1, "max_threads must be >= 1");
101        assert!(
102            cache_limit_bytes >= -1,
103            "cache_limit_bytes must be -1 (unlimited) or non-negative"
104        );
105        // The cross-grammar `RuleLevelCache` exists only when caching is
106        // enabled. Its memory budget mirrors the C++ split: with an
107        // unlimited (`-1`) overall limit it is unbounded; otherwise it
108        // gets `limit - limit/3*2` bytes (~1/3 of the budget — the other
109        // ~2/3 goes to the grammar-level cache).
110        let rule_cache = if cache_enabled {
111            let budget = if cache_limit_bytes == -1 {
112                UNLIMITED_SIZE
113            } else {
114                (cache_limit_bytes - cache_limit_bytes / 3 * 2) as usize
115            };
116            Some(RuleLevelCache::new(budget))
117        } else {
118            None
119        };
120        Self {
121            tokenizer_info,
122            max_threads,
123            cache_enabled,
124            cache_limit_bytes,
125            cache: GrammarCache::new(if cache_limit_bytes == -1 {
126                UNLIMITED_BYTES
127            } else {
128                // The ~2/3 complement of the rule cache's `limit - limit/3*2`.
129                (cache_limit_bytes / 3 * 2) as usize
130            }),
131            rule_cache,
132        }
133    }
134
135    /// Run the full W3 pipeline (normalize then optimize) on a parsed
136    /// grammar, then compile it against this compiler's tokenizer.
137    fn compile_normalized(&self, grammar: GrammarData) -> CompiledGrammar {
138        let normalized = GrammarNormalizer::apply(grammar);
139        let optimized = GrammarOptimizer::apply(normalized);
140        compile_optimized_grammar(
141            optimized,
142            &self.tokenizer_info,
143            self.max_threads,
144            self.rule_cache.clone(),
145        )
146    }
147
148    /// Fetch from cache or compute via `f`, honoring `cache_enabled`.
149    fn get_or_compute<F>(&self, key: CacheKey, f: F) -> CompiledGrammar
150    where
151        F: FnOnce() -> CompiledGrammar,
152    {
153        if !self.cache_enabled {
154            return f();
155        }
156        if let Some(hit) = self.cache.get(&key) {
157            return hit;
158        }
159        let value = f();
160        self.cache.insert(key, value.clone());
161        value
162    }
163
164    /// Compile a grammar from an EBNF string. Port of
165    /// `GrammarCompiler::CompileGrammar(ebnf, root_rule_name)`.
166    pub fn compile_grammar_from_ebnf(
167        &self,
168        ebnf: &str,
169        root_rule_name: &str,
170    ) -> Result<CompiledGrammar, CompileError> {
171        // Parse eagerly so a malformed grammar is a typed error rather
172        // than a cached panic.
173        let grammar =
174            parse_ebnf(ebnf, root_rule_name).map_err(|e| CompileError::Grammar(e.to_string()))?;
175        let key = CacheKey::Ebnf {
176            ebnf: ebnf.to_string(),
177            root_rule: root_rule_name.to_string(),
178        };
179        Ok(self.get_or_compute(key, || self.compile_normalized(grammar)))
180    }
181
182    /// Compile an already-parsed [`GrammarData`]. Port of
183    /// `GrammarCompiler::CompileGrammar(const Grammar&)` — keyed by the
184    /// grammar's printed EBNF, as the C++ does.
185    pub fn compile_grammar(&self, grammar: GrammarData) -> CompiledGrammar {
186        let ebnf = crate::grammar::print_grammar(&grammar);
187        let root_rule = grammar.root_rule().name.clone();
188        let key = CacheKey::Ebnf { ebnf, root_rule };
189        self.get_or_compute(key, || self.compile_normalized(grammar))
190    }
191
192    /// Compile the builtin "any JSON value" grammar. Port of
193    /// `GrammarCompiler::CompileBuiltinJSONGrammar`.
194    pub fn compile_builtin_json_grammar(&self) -> Result<CompiledGrammar, CompileError> {
195        let ebnf = builtin_json_grammar_ebnf();
196        let grammar = parse_ebnf(&ebnf, "root").map_err(|e| CompileError::Schema(e.to_string()))?;
197        Ok(self.get_or_compute(CacheKey::BuiltinJson, || self.compile_normalized(grammar)))
198    }
199
200    /// Compile a JSON schema. Port of
201    /// `GrammarCompiler::CompileJSONSchema`.
202    #[allow(clippy::too_many_arguments)]
203    pub fn compile_json_schema(
204        &self,
205        schema: &str,
206        any_whitespace: bool,
207        indent: Option<i32>,
208        separators: Option<(String, String)>,
209        strict_mode: bool,
210        max_whitespace_cnt: Option<i32>,
211    ) -> Result<CompiledGrammar, CompileError> {
212        let options = SchemaConverterOptions {
213            any_whitespace,
214            indent,
215            separators: separators.clone(),
216            strict_mode,
217            max_whitespace_cnt,
218            ..SchemaConverterOptions::default()
219        };
220        let ebnf = json_schema_to_ebnf(schema, &options)
221            .map_err(|e| CompileError::Schema(e.to_string()))?;
222        let grammar = parse_ebnf(&ebnf, "root")
223            .map_err(|e| CompileError::Schema(format!("generated EBNF failed to parse: {e}")))?;
224        let key = CacheKey::Schema {
225            schema: schema.to_string(),
226            any_whitespace,
227            indent,
228            separators,
229            strict_mode,
230            max_whitespace_cnt,
231        };
232        Ok(self.get_or_compute(key, || self.compile_normalized(grammar)))
233    }
234
235    /// Compile a structural tag. Port of
236    /// `GrammarCompiler::CompileStructuralTag`.
237    ///
238    /// Delegates the structural-tag-JSON -> `GrammarData` conversion to
239    /// the W5 `structural_tag` module (`structural_tag_to_grammar`).
240    pub fn compile_structural_tag(
241        &self,
242        structural_tag_json: &str,
243    ) -> Result<CompiledGrammar, CompileError> {
244        let grammar = structural_tag_to_grammar(structural_tag_json)
245            .map_err(|e| CompileError::StructuralTag(e.to_string()))?;
246        let key = CacheKey::StructuralTag(structural_tag_json.to_string());
247        Ok(self.get_or_compute(key, || self.compile_normalized(grammar)))
248    }
249
250    /// Clear the internal compiled-grammar cache *and* the cross-grammar
251    /// rule-level mask cache. Port of `GrammarCompiler::ClearCache`.
252    pub fn clear_cache(&self) {
253        self.cache.clear();
254        if let Some(rule_cache) = &self.rule_cache {
255            rule_cache.clear();
256        }
257    }
258
259    /// Approximate bytes held by the cache — the grammar-level compiled
260    /// grammars plus the cross-grammar rule-level mask cache. Port of
261    /// `GrammarCompiler::GetCacheSizeBytes`.
262    pub fn cache_size_bytes(&self) -> i64 {
263        let grammar_bytes = self.cache.resident_bytes() as i64;
264        let rule_bytes = self
265            .rule_cache
266            .as_ref()
267            .map_or(0, |c| c.memory_size() as i64);
268        grammar_bytes + rule_bytes
269    }
270
271    /// The configured cache memory limit; `-1` means unlimited. Port of
272    /// `GrammarCompiler::CacheLimitBytes`.
273    pub fn cache_limit_bytes(&self) -> i64 {
274        self.cache_limit_bytes
275    }
276
277    /// The tokenizer this compiler is bound to.
278    pub fn tokenizer_info(&self) -> &TokenizerInfo {
279        &self.tokenizer_info
280    }
281
282    /// The identity an on-disk mask snapshot must match to be usable
283    /// with this compiler (#918).
284    pub fn snapshot_identity(&self) -> SnapshotIdentity {
285        SnapshotIdentity {
286            tokenizer_fingerprint: self.tokenizer_info.fingerprint(),
287            vocab_size: self.tokenizer_info.vocab_size(),
288        }
289    }
290
291    /// Seed the cross-grammar [`RuleLevelCache`] from a snapshot written
292    /// by an earlier process (#918).
293    ///
294    /// Returns the number of masks imported. `0` means a miss — no file,
295    /// a snapshot from a different tokenizer or build, or a corrupt one —
296    /// and the compiler simply computes as before. A compiler constructed
297    /// with `cache_enabled == false` has no rule cache and always
298    /// returns `0`.
299    pub fn load_mask_snapshot(&self, path: &std::path::Path) -> Result<usize, SnapshotError> {
300        let Some(cache) = &self.rule_cache else {
301            return Ok(0);
302        };
303        mask_snapshot::load_from_file(cache, self.snapshot_identity(), path)
304    }
305
306    /// Persist the cross-grammar [`RuleLevelCache`] so the NEXT process
307    /// starts warm (#918). Returns the number of masks written; at most
308    /// `max_entries`, most-recently-used first.
309    pub fn save_mask_snapshot(
310        &self,
311        path: &std::path::Path,
312        max_entries: usize,
313    ) -> Result<usize, SnapshotError> {
314        let Some(cache) = &self.rule_cache else {
315            return Ok(0);
316        };
317        mask_snapshot::save_to_file(cache, self.snapshot_identity(), path, max_entries)
318    }
319
320    /// The cross-grammar rule-level cache itself, when caching is on.
321    /// Exposed for the #918 snapshot tests and for callers that want to
322    /// share one cache across compilers.
323    pub fn rule_cache(&self) -> Option<&RuleLevelCache> {
324        self.rule_cache.as_ref()
325    }
326
327    /// Entries currently held by the cross-grammar rule-level cache —
328    /// `0` when caching is disabled. Used by the #918 tests and by the
329    /// serve path's warm-up logging.
330    pub fn rule_cache_len(&self) -> usize {
331        self.rule_cache.as_ref().map_or(0, |c| c.len())
332    }
333}
metrale_grammar/compiler/compiler.rs

metrale_grammar/compiler/
compiler.rs
