Skip to main content

maskforge_core/runtime/
mod.rs

1//! Provides compiled artifacts and stateful matchers.
2
3use std::sync::atomic::{AtomicUsize, Ordering};
4use std::sync::{Arc, OnceLock};
5
6use crate::automaton::RefEngine;
7use crate::error::{CompileError, ErrorCode, MatcherError, Provenance, Stage};
8use crate::index::{
9    build_delta, BindMode, BoundByteTrie, CompiledIndex, VocabFingerprint, VocabTrie,
10    VocabularyHandle,
11};
12use crate::mask::Bitmask;
13use crate::primitives::{StateId, TokenId};
14use crate::vocab::{CanonicalVocabulary, Vocabulary};
15
16mod artifact_cache;
17mod executable_cache;
18pub use artifact_cache::{ArtifactCache, COMPILER_SEMANTICS_VERSION};
19pub use executable_cache::{
20    global_executable_cache, CachedExecutable, ExecutableCache, ExecutableCacheStats,
21    ExecutableSchema, EXECUTABLE_SEMANTICS_VERSION,
22};
23
24/// Total mask-buffer cap in `u32` words: `rows * words_per_row` above this is rejected before
25/// allocation, so a large row index cannot force a multi-gigabyte buffer. `1 << 26` words = 256 MiB.
26const MAX_MASK_WORDS: usize = 1 << 26;
27
28/// Byte budget for the per-state mask cache. Once this many bytes are retained, further states are
29/// computed but not cached, so a large DFA plus a wide vocabulary cannot retain unbounded memory.
30const MAX_MASK_CACHE_BYTES: usize = 1 << 26;
31
32/// Conservative per-entry overhead charged on top of a cached mask's words: the `Arc<Bitmask>`
33/// control block, the `Vec` header, and the hash-map slot. Keeps the byte budget from being
34const MASK_CACHE_ENTRY_OVERHEAD: usize = 64;
35
36/// Budget for the INCREMENTAL packed-build peak of the `TrieJointBytePacked` mode: the payload
37/// table, a fixed header allowance, the dense walk scratch, and the DFS stack's exact peak (see
38const MAX_PACKED_MASK_BYTES: usize = 1 << 28;
39
40/// Measured packed row-fill cost per `u32` word: the delta between packed and forced-lazy
41/// `ir_to_index` time divided by `state_count * words_per_row`, calibrated on gpt2 (wpr=1571,
42const PACKED_ROW_FILL_NS_PER_WORD: usize = 2;
43
44/// Target ceiling on the packed row-fill contribution to compile latency - a policy choice, not a
45/// measured fact: 20ms keeps an eager packed compile from becoming visibly slow relative to a
46const PACKED_COMPILE_LATENCY_BUDGET_US: usize = 20_000;
47
48/// Fixed allowance added to the incremental packed-build peak estimate for the boxed-slice header, the
49/// artifact struct growth, and allocator rounding, so the declared budget is not silently exceeded by
50const PACKED_HEADER_ALLOWANCE: usize = 1 << 16;
51
52/// One contiguous packed mask matrix: `rows` holds `state_count * words_per_row` little-endian `u32`
53/// words, row `s` at `[s*wpr .. (s+1)*wpr]`. A `Box<[u32]>` (NOT `Arc<[u32]>`): `CompiledArtifact` is
54#[derive(Debug)]
55struct PackedMaskTable {
56    words_per_row: usize,
57    rows: Box<[u32]>,
58}
59
60impl PackedMaskTable {
61    /// The packed words for `s`, or `None` for an out-of-range (forged) state. Borrowed, no copy.
62    fn row(&self, s: StateId) -> Option<&[u32]> {
63        let wpr = self.words_per_row;
64        let start = (s.get() as usize).checked_mul(wpr)?;
65        let end = start.checked_add(wpr)?;
66        self.rows.get(start..end)
67    }
68
69    /// Retained payload bytes (the table words; metadata is a fixed small struct).
70    fn retained_bytes(&self) -> usize {
71        self.rows.len().saturating_mul(4)
72    }
73}
74
75/// A slot's outcome once decided. `Uncached` is permanent: production budgets never grow after
76/// construction, so a state rejected once stays rejected for the artifact's lifetime.
77#[derive(Clone, Debug)]
78enum SlotValue {
79    Cached(Result<Arc<Bitmask>, Arc<CompileError>>),
80    Uncached,
81}
82
83type MaskSlot = OnceLock<SlotValue>;
84
85/// Cap on the `LazyRows` slot array itself (not the masks it later retains): `1 << 28` = 256 MiB,
86/// mirroring the packed budget.
87const MAX_LAZY_ROWS_BYTES: usize = 1 << 28;
88
89/// One `OnceLock` per state: a query for `s` never contends with a query for a different state.
90#[derive(Debug)]
91struct LazyRows {
92    rows: Box<[MaskSlot]>,
93    retained_bytes: AtomicUsize,
94    budget_bytes: AtomicUsize,
95}
96
97impl LazyRows {
98    fn try_new(state_count: usize, budget_bytes: usize) -> Result<Self, CompileError> {
99        state_count
100            .checked_mul(std::mem::size_of::<MaskSlot>())
101            .filter(|&b| b <= MAX_LAZY_ROWS_BYTES)
102            .ok_or_else(|| {
103                CompileError::new(
104                    ErrorCode::InternalLimitExceeded,
105                    Stage::L4Bind,
106                    "lazy row slot array exceeds the size cap",
107                )
108            })?;
109        let mut rows = Vec::new();
110        rows.try_reserve_exact(state_count).map_err(|_| {
111            CompileError::new(
112                ErrorCode::InternalLimitExceeded,
113                Stage::L4Bind,
114                "lazy row slot array allocation failed",
115            )
116        })?;
117        rows.resize_with(state_count, OnceLock::new);
118        Ok(Self {
119            rows: rows.into_boxed_slice(),
120            retained_bytes: AtomicUsize::new(0),
121            budget_bytes: AtomicUsize::new(budget_bytes),
122        })
123    }
124
125    fn try_reserve(&self, bytes: usize) -> bool {
126        let budget = self.budget_bytes.load(Ordering::Acquire);
127        self.retained_bytes
128            .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| {
129                current.checked_add(bytes).filter(|&total| total <= budget)
130            })
131            .is_ok()
132    }
133
134    /// `state` must be `< state_count` (checked by the caller). Only the `OnceLock`-winning
135    /// closure ever reserves; losing racers read its published result and never reserve.
136    fn get_or_compute(
137        &self,
138        state: usize,
139        entry_bytes: usize,
140        compute: impl FnOnce() -> Result<Bitmask, CompileError>,
141    ) -> Result<Arc<Bitmask>, CompileError> {
142        let mut compute = Some(compute);
143        let value = match self.rows[state].get() {
144            Some(v) => v,
145            None => self.rows[state].get_or_init(|| {
146                if !self.try_reserve(entry_bytes) {
147                    return SlotValue::Uncached;
148                }
149                let result = compute.take().expect("reserved path runs once")().map(Arc::new);
150                if result.is_err() {
151                    self.retained_bytes.fetch_sub(entry_bytes, Ordering::AcqRel);
152                }
153                SlotValue::Cached(result.map_err(Arc::new))
154            }),
155        };
156        match value {
157            SlotValue::Cached(Ok(mask)) => Ok(mask.clone()),
158            SlotValue::Cached(Err(e)) => Err((**e).clone()),
159            SlotValue::Uncached => {
160                compute.take().expect("untouched on the uncached path")().map(Arc::new)
161            }
162        }
163    }
164
165    /// Conservative upper bound on what this structure could ever retain: the slot array itself
166    /// plus its configured mask-byte budget.
167    fn max_charge_bytes(&self) -> usize {
168        self.rows
169            .len()
170            .saturating_mul(std::mem::size_of::<MaskSlot>())
171            .saturating_add(self.budget_bytes.load(Ordering::Acquire))
172    }
173
174    #[cfg(test)]
175    fn filled_count(&self) -> usize {
176        self.rows
177            .iter()
178            .filter(|r| matches!(r.get(), Some(SlotValue::Cached(_))))
179            .count()
180    }
181
182    #[cfg(test)]
183    fn retained_bytes(&self) -> usize {
184        self.retained_bytes.load(Ordering::Acquire)
185    }
186
187    #[cfg(test)]
188    fn set_budget_for_test(&self, budget: usize) {
189        self.budget_bytes.store(budget, Ordering::Release);
190    }
191}
192
193/// Mask storage backing `cached_mask`. Packed mode needs none: its table is already fully
194/// materialized, so allocating per-state slots would be pure waste.
195#[derive(Debug)]
196enum MaskRows {
197    Disabled,
198    Lazy(LazyRows),
199}
200
201/// A byte engine bound to a vocabulary, shared by [`Arc`] across sequences and the FFI. Only
202/// `mask_cache` mutates after construction.
203#[derive(Debug)]
204pub struct CompiledArtifact {
205    engine: Arc<RefEngine>,
206    vocab: Arc<CanonicalVocabulary>,
207    provenance: Provenance,
208    start_is_dead: bool,
209    bind_mode: BindMode,
210    index: Option<Arc<CompiledIndex>>,
211    lazy_trie: Option<Arc<VocabTrie>>,
212    packed: Option<PackedMaskTable>,
213    mask_cache: MaskRows,
214}
215
216impl CompiledArtifact {
217    /// Binds `engine` to `vocab` under `provenance`, validating the result. Returns a structured
218    /// error (never a partial artifact) if provenance, start state, or token-id bounds are invalid.
219    pub fn new(
220        engine: RefEngine,
221        vocab: Arc<Vocabulary>,
222        provenance: Provenance,
223    ) -> Result<Self, CompileError> {
224        let canonical = Arc::new(CanonicalVocabulary::from_vocabulary(&vocab)?);
225        Self::new_with_canonical(engine, canonical, provenance)
226    }
227
228    /// The safe public bind (naive mode): takes the canonical vocabulary from ONE `VocabularyHandle`,
229    /// so the mask width and the id lookup always agree with each other by construction.
230    pub fn new_from_handle(
231        handle: &VocabularyHandle,
232        engine: RefEngine,
233        provenance: Provenance,
234    ) -> Result<Self, CompileError> {
235        Self::new_with_canonical(engine, handle.canonical().clone(), provenance)
236    }
237
238    /// Naive binding from a vocabulary-neutral shared engine.
239    pub fn new_from_handle_shared(
240        handle: &VocabularyHandle,
241        engine: Arc<RefEngine>,
242        provenance: Provenance,
243    ) -> Result<Self, CompileError> {
244        Self::new_full(
245            engine,
246            handle.canonical().clone(),
247            provenance,
248            BindMode::Naive,
249            None,
250            None,
251            None,
252        )
253    }
254
255    /// Binds using an already-built canonical vocabulary instead of rebuilding it (the warm-bind fast
256    /// path). Crate-internal: `vocab` is trusted to genuinely be canonical, so only a caller that
257    pub(crate) fn new_with_canonical(
258        engine: RefEngine,
259        vocab: Arc<CanonicalVocabulary>,
260        provenance: Provenance,
261    ) -> Result<Self, CompileError> {
262        Self::new_full(
263            Arc::new(engine),
264            vocab,
265            provenance,
266            BindMode::Naive,
267            None,
268            None,
269            None,
270        )
271    }
272
273    /// Builds the artifact with its FINAL bind mode decided up front - no field is mutated after.
274    /// Packed gets no `LazyRows` (its table is already complete); every other mode gets one.
275    fn new_full(
276        engine: Arc<RefEngine>,
277        vocab: Arc<CanonicalVocabulary>,
278        provenance: Provenance,
279        bind_mode: BindMode,
280        index: Option<Arc<CompiledIndex>>,
281        lazy_trie: Option<Arc<VocabTrie>>,
282        packed: Option<PackedMaskTable>,
283    ) -> Result<Self, CompileError> {
284        let start_is_dead = engine.is_dead(engine.start());
285        let mask_cache = if bind_mode == BindMode::TrieJointBytePacked {
286            MaskRows::Disabled
287        } else {
288            MaskRows::Lazy(LazyRows::try_new(
289                engine.state_count(),
290                MAX_MASK_CACHE_BYTES,
291            )?)
292        };
293        let artifact = Self {
294            engine,
295            vocab,
296            provenance,
297            start_is_dead,
298            bind_mode,
299            index,
300            lazy_trie,
301            packed,
302            mask_cache,
303        };
304        artifact.validate()?;
305        Ok(artifact)
306    }
307
308    /// Binds `engine` to `vocab` from a borrowed `trie`: the eager delta walks
309    /// (`TrieJointByte`/`TrieJointClass`) and the packed table (`TrieJointBytePacked`).
310    pub fn new_trie(
311        engine: RefEngine,
312        vocab: Arc<Vocabulary>,
313        provenance: Provenance,
314        mode: BindMode,
315        trie: &VocabTrie,
316    ) -> Result<Self, CompileError> {
317        let canonical = Arc::new(CanonicalVocabulary::from_vocabulary(&vocab)?);
318        let fingerprint = canonical.fingerprint();
319        match mode {
320            BindMode::TrieJointByte | BindMode::TrieJointClass => {
321                Self::new_trie_prevalidated(engine, canonical, provenance, mode, trie, fingerprint)
322            }
323            BindMode::TrieJointBytePacked => {
324                Self::new_packed_from_trie(engine, canonical, provenance, trie, fingerprint)
325            }
326            BindMode::TrieJointByteLazy | BindMode::Naive => Err(CompileError::new(
327                ErrorCode::ArtifactOutOfBounds,
328                Stage::L4Bind,
329                "this bind mode is not constructible from a borrowed trie; use new_from_bound_trie \
330                 (byte_trie_lazy) or new/new_from_handle (naive)",
331            )),
332        }
333    }
334
335    /// Same check as `new_trie`, but skips rebuilding `vocab`. Not public: a caller-suppliable
336    /// fingerprint here is never re-verified against `vocab`'s real content, so only
337    pub(crate) fn new_trie_prevalidated(
338        engine: RefEngine,
339        vocab: Arc<CanonicalVocabulary>,
340        provenance: Provenance,
341        mode: BindMode,
342        trie: &VocabTrie,
343        vocab_fingerprint: VocabFingerprint,
344    ) -> Result<Self, CompileError> {
345        Self::new_trie_prevalidated_shared(
346            Arc::new(engine),
347            vocab,
348            provenance,
349            mode,
350            trie,
351            vocab_fingerprint,
352        )
353    }
354
355    fn new_trie_prevalidated_shared(
356        engine: Arc<RefEngine>,
357        vocab: Arc<CanonicalVocabulary>,
358        provenance: Provenance,
359        mode: BindMode,
360        trie: &VocabTrie,
361        vocab_fingerprint: VocabFingerprint,
362    ) -> Result<Self, CompileError> {
363        if trie.vocab_fingerprint() != vocab_fingerprint {
364            return Err(CompileError::new(
365                ErrorCode::ArtifactMismatch,
366                Stage::L4Bind,
367                "trie was built from a different vocabulary than the one being bound",
368            ));
369        }
370        let index = build_delta(&engine, trie, mode, provenance)?;
371        Self::new_full(
372            engine,
373            vocab,
374            provenance,
375            mode,
376            Some(Arc::new(index)),
377            None,
378            None,
379        )
380    }
381
382    /// Binds `engine` to `handle`'s vocabulary using `bound`, a trie the trie cache already proved
383    /// paired with SOME `VocabularyHandle`. Rejects `bound` if it was not proven paired with THIS
384    pub fn new_from_bound_trie(
385        handle: &VocabularyHandle,
386        engine: RefEngine,
387        provenance: Provenance,
388        mode: BindMode,
389        bound: &BoundByteTrie,
390    ) -> Result<Self, CompileError> {
391        Self::new_from_bound_trie_shared(handle, Arc::new(engine), provenance, mode, bound)
392    }
393
394    /// Binds a vocabulary-neutral shared engine without cloning its states or transition tables.
395    pub fn new_from_bound_trie_shared(
396        handle: &VocabularyHandle,
397        engine: Arc<RefEngine>,
398        provenance: Provenance,
399        mode: BindMode,
400        bound: &BoundByteTrie,
401    ) -> Result<Self, CompileError> {
402        if bound.fingerprint() != handle.fingerprint() {
403            return Err(CompileError::new(
404                ErrorCode::ArtifactMismatch,
405                Stage::L4Bind,
406                "bound trie does not match this vocabulary handle",
407            ));
408        }
409        if mode == BindMode::TrieJointByteLazy {
410            return Self::new_lazy_trie_shared(
411                engine,
412                handle.canonical().clone(),
413                provenance,
414                bound.trie_arc(),
415                handle.fingerprint(),
416            );
417        }
418        if mode == BindMode::TrieJointBytePacked {
419            return Self::new_packed_from_trie_shared(
420                engine,
421                handle.canonical().clone(),
422                provenance,
423                bound.trie(),
424                handle.fingerprint(),
425            );
426        }
427        Self::new_trie_prevalidated_shared(
428            engine,
429            handle.canonical().clone(),
430            provenance,
431            mode,
432            bound.trie(),
433            handle.fingerprint(),
434        )
435    }
436
437    /// Recommends `requested` as-is unless it is `TrieJointBytePacked` and either the packed byte
438    /// budget (a safety limit) or the packed compile-latency budget (a policy target, see
439    #[must_use]
440    pub fn recommended_bind_mode(
441        state_count: usize,
442        mask_vocab_size: usize,
443        bound: &BoundByteTrie,
444        requested: BindMode,
445    ) -> BindMode {
446        if requested != BindMode::TrieJointBytePacked {
447            return requested;
448        }
449        let wpr = mask_vocab_size.div_ceil(32);
450        let over_budget = !matches!(
451            Self::packed_build_peak_bytes(state_count, wpr, bound.trie()),
452            Some(peak) if peak <= MAX_PACKED_MASK_BYTES
453        );
454        let over_latency = state_count
455            .checked_mul(wpr)
456            .and_then(|words| words.checked_mul(PACKED_ROW_FILL_NS_PER_WORD))
457            .is_none_or(|ns| ns / 1000 > PACKED_COMPILE_LATENCY_BUDGET_US);
458        if over_budget || over_latency {
459            BindMode::TrieJointByteLazy
460        } else {
461            requested
462        }
463    }
464
465    /// Binds `engine` to `vocab` in packed mode. Packed needs NO `Delta`: the packed payload is
466    /// `state_count * words_per_row * 4`, exact and known before any walk, so each state's row is
467    fn new_packed_from_trie(
468        engine: RefEngine,
469        vocab: Arc<CanonicalVocabulary>,
470        provenance: Provenance,
471        trie: &VocabTrie,
472        fingerprint: VocabFingerprint,
473    ) -> Result<Self, CompileError> {
474        Self::new_packed_from_trie_shared(Arc::new(engine), vocab, provenance, trie, fingerprint)
475    }
476
477    fn new_packed_from_trie_shared(
478        engine: Arc<RefEngine>,
479        vocab: Arc<CanonicalVocabulary>,
480        provenance: Provenance,
481        trie: &VocabTrie,
482        fingerprint: VocabFingerprint,
483    ) -> Result<Self, CompileError> {
484        Self::new_packed_from_trie_budgeted_shared(
485            engine,
486            vocab,
487            provenance,
488            trie,
489            fingerprint,
490            MAX_PACKED_MASK_BYTES,
491        )
492    }
493
494    /// `new_packed_from_trie` with an explicit `budget` (production passes `MAX_PACKED_MASK_BYTES`;
495    /// tests pass a small budget). Admission is EXACT: the only heap allocation the packed build makes
496    #[cfg(test)]
497    fn new_packed_from_trie_budgeted(
498        engine: RefEngine,
499        vocab: Arc<CanonicalVocabulary>,
500        provenance: Provenance,
501        trie: &VocabTrie,
502        fingerprint: VocabFingerprint,
503        budget: usize,
504    ) -> Result<Self, CompileError> {
505        Self::new_packed_from_trie_budgeted_shared(
506            Arc::new(engine),
507            vocab,
508            provenance,
509            trie,
510            fingerprint,
511            budget,
512        )
513    }
514
515    fn new_packed_from_trie_budgeted_shared(
516        engine: Arc<RefEngine>,
517        vocab: Arc<CanonicalVocabulary>,
518        provenance: Provenance,
519        trie: &VocabTrie,
520        fingerprint: VocabFingerprint,
521        budget: usize,
522    ) -> Result<Self, CompileError> {
523        if trie.vocab_fingerprint() != fingerprint {
524            return Err(CompileError::new(
525                ErrorCode::ArtifactMismatch,
526                Stage::L4Bind,
527                "trie was built from a different vocabulary than the one being bound",
528            ));
529        }
530        let n = engine.state_count();
531        let wpr = vocab.mask_width_words();
532        let payload = Self::packed_payload_bytes(n, wpr).ok_or_else(Self::packed_overflow)?;
533        let peak = Self::packed_build_peak_bytes(n, wpr, trie).ok_or_else(Self::packed_overflow)?;
534        if peak > budget {
535            return Err(CompileError::new(
536                ErrorCode::InternalLimitExceeded,
537                Stage::L4Bind,
538                "packed mask table plus walk scratch exceeds the memory budget; select byte_trie_lazy",
539            ));
540        }
541        let _permit = crate::mem_gate::acquire_vocab_build_bytes(peak)?;
542        let stack_fuel = trie.max_dfs_frontier();
543        let word_count = payload / 4;
544        let mut rows: Vec<u32> = Vec::new();
545        rows.try_reserve_exact(word_count)
546            .map_err(|_| Self::packed_alloc_failed())?;
547        rows.resize(word_count, 0);
548        crate::index::fill_packed_rows(
549            trie,
550            &engine,
551            &mut rows,
552            wpr,
553            vocab.mask_vocab_size(),
554            stack_fuel,
555        )?;
556        let packed = PackedMaskTable {
557            words_per_row: wpr,
558            rows: rows.into_boxed_slice(),
559        };
560        Self::new_full(
561            engine,
562            vocab,
563            provenance,
564            BindMode::TrieJointBytePacked,
565            None,
566            None,
567            Some(packed),
568        )
569    }
570
571    /// Retained packed payload bytes for `state_count` states of `wpr` words, checked (`None` on
572    /// overflow). The exact admission size (the packed table is the packed build's only allocation).
573    #[must_use]
574    fn packed_payload_bytes(state_count: usize, wpr: usize) -> Option<usize> {
575        state_count.checked_mul(wpr)?.checked_mul(4)
576    }
577
578    /// The exact incremental packed-build peak: payload table, fixed header allowance, dense walk
579    /// scratch, and the DFS stack's precomputed exact peak - the same total `new_packed_from_trie_
580    fn packed_build_peak_bytes(state_count: usize, wpr: usize, trie: &VocabTrie) -> Option<usize> {
581        let payload = Self::packed_payload_bytes(state_count, wpr)?;
582        let scratch_fixed = crate::index::packed_scratch_fixed_bytes()?;
583        let stack_bytes = crate::index::packed_stack_bytes(trie.max_dfs_frontier())?;
584        payload
585            .checked_add(PACKED_HEADER_ALLOWANCE)?
586            .checked_add(scratch_fixed)?
587            .checked_add(stack_bytes)
588    }
589
590    fn packed_overflow() -> CompileError {
591        CompileError::new(
592            ErrorCode::InternalLimitExceeded,
593            Stage::L4Bind,
594            "packed mask matrix size overflows",
595        )
596    }
597
598    fn packed_alloc_failed() -> CompileError {
599        CompileError::new(
600            ErrorCode::InternalLimitExceeded,
601            Stage::L4Bind,
602            "packed mask matrix allocation failed",
603        )
604    }
605
606    /// The borrowed packed row for `s` in packed mode (no allocation), else `None`. The native
607    /// zero-copy serving view; a Python view must hold an `Arc` to keep the artifact alive.
608    #[must_use]
609    pub fn mask_row_words(&self, s: StateId) -> Option<&[u32]> {
610        self.packed.as_ref()?.row(s)
611    }
612
613    /// The WHOLE packed mask matrix, borrowed flat as `state_count() * words_per_row()` little-endian
614    /// `u32` words (row `s` at `[s*words_per_row .. (s+1)*words_per_row]`), or `None` outside packed
615    #[must_use]
616    pub fn packed_mask_table_words(&self) -> Option<&[u32]> {
617        self.packed.as_ref().map(|p| &p.rows[..])
618    }
619
620    /// Retained bytes of the packed mask matrix (0 when not in packed mode), for memory reporting.
621    #[must_use]
622    pub fn packed_retained_bytes(&self) -> usize {
623        self.packed
624            .as_ref()
625            .map_or(0, PackedMaskTable::retained_bytes)
626    }
627
628    /// Retained bytes of the eager delta index (0 outside `TrieJointByte`/`TrieJointClass` modes),
629    /// for the compiled-artifact cache's byte-budget accounting.
630    #[must_use]
631    pub fn index_retained_bytes(&self) -> usize {
632        self.index.as_ref().map_or(0, |i| i.heap_bytes())
633    }
634
635    /// Retained bytes of the byte engine itself (states, transitions, class table) - present in
636    /// every bind mode, so the compiled-artifact cache charges it alongside the packed/index bytes
637    #[must_use]
638    pub fn engine_retained_bytes(&self) -> usize {
639        self.engine.heap_bytes()
640    }
641
642    /// Conservative upper bound on total retained bytes, valid for ANY bind mode and stable for
643    /// the artifact's whole lifetime - safe to charge once at cache-admission time.
644    #[must_use]
645    pub fn artifact_cache_charge_upper_bound(&self) -> usize {
646        let lazy_bytes = match &self.mask_cache {
647            MaskRows::Disabled => 0,
648            MaskRows::Lazy(rows) => rows.max_charge_bytes(),
649        };
650        let trie_bytes = self.lazy_trie.as_ref().map_or(0, |t| t.heap_bytes());
651        self.engine_retained_bytes()
652            .saturating_add(self.packed_retained_bytes())
653            .saturating_add(self.index_retained_bytes())
654            .saturating_add(lazy_bytes)
655            .saturating_add(trie_bytes)
656    }
657
658    /// Binds `engine` to `vocab` in lazy byte-trie mode: no eager `Delta`; each state's row is
659    /// walked on first query and cached. Rejects a trie built from a different vocabulary.
660    #[cfg(test)]
661    pub(crate) fn new_lazy_trie(
662        engine: RefEngine,
663        vocab: Arc<CanonicalVocabulary>,
664        provenance: Provenance,
665        trie: Arc<VocabTrie>,
666        vocab_fingerprint: VocabFingerprint,
667    ) -> Result<Self, CompileError> {
668        Self::new_lazy_trie_shared(Arc::new(engine), vocab, provenance, trie, vocab_fingerprint)
669    }
670
671    fn new_lazy_trie_shared(
672        engine: Arc<RefEngine>,
673        vocab: Arc<CanonicalVocabulary>,
674        provenance: Provenance,
675        trie: Arc<VocabTrie>,
676        vocab_fingerprint: VocabFingerprint,
677    ) -> Result<Self, CompileError> {
678        if trie.vocab_fingerprint() != vocab_fingerprint {
679            return Err(CompileError::new(
680                ErrorCode::ArtifactMismatch,
681                Stage::L4Bind,
682                "trie was built from a different vocabulary than the one being bound",
683            ));
684        }
685        Self::new_full(
686            engine,
687            vocab,
688            provenance,
689            BindMode::TrieJointByteLazy,
690            None,
691            Some(trie),
692            None,
693        )
694    }
695
696    /// Validates the bindings this layer owns: provenance, the start state in range, every bound
697    /// token id inside the mask, and (unless the language is empty) at least one accepting state.
698    pub fn validate(&self) -> Result<(), CompileError> {
699        self.provenance.validate()?;
700        let out_of_bounds = |msg| {
701            Err(CompileError::new(
702                ErrorCode::ArtifactOutOfBounds,
703                Stage::ArtifactValidate,
704                msg,
705            ))
706        };
707        if (self.engine.start().get() as usize) >= self.engine.state_count() {
708            return out_of_bounds("artifact start state is out of range");
709        }
710        if !self.start_is_dead
711            && !(0..self.engine.state_count())
712                .filter_map(|raw| StateId::try_from(raw).ok())
713                .any(|s| self.engine.is_accepting(s))
714        {
715            return Err(CompileError::new(
716                ErrorCode::Malformed,
717                Stage::ArtifactValidate,
718                "a non-empty-language artifact has no accepting state",
719            ));
720        }
721        Ok(())
722    }
723
724    /// The anchored start state.
725    #[must_use]
726    pub fn start(&self) -> StateId {
727        self.engine.start()
728    }
729
730    /// The canonical DEAD state.
731    #[must_use]
732    pub fn dead(&self) -> StateId {
733        self.engine.dead()
734    }
735
736    /// Whether the start state is DEAD (the language accepts nothing). Diagnostic only.
737    #[must_use]
738    pub fn start_is_dead(&self) -> bool {
739        self.start_is_dead
740    }
741
742    /// The artifact's provenance record.
743    #[must_use]
744    pub fn provenance(&self) -> Provenance {
745        self.provenance
746    }
747
748    /// The token-id space the mask covers (`max_token_id + 1`; EOS occupies a bit, never set).
749    #[must_use]
750    pub fn vocab_size(&self) -> usize {
751        self.vocab.mask_vocab_size()
752    }
753
754    /// The number of reference states, including DEAD. State ids span `0..state_count`.
755    #[must_use]
756    pub fn state_count(&self) -> usize {
757        self.engine.state_count()
758    }
759
760    /// Walks `bytes` from `from`; `None` once DEAD (or an out-of-range state) is reached.
761    #[must_use]
762    pub fn walk(&self, from: StateId, bytes: &[u8]) -> Option<StateId> {
763        self.engine.consume_token(from, bytes)
764    }
765
766    /// The number of packed `u32` words in one mask row: `ceil(vocab_size / 32)`.
767    #[must_use]
768    pub fn words_per_row(&self) -> usize {
769        self.vocab.mask_width_words()
770    }
771
772    /// Stopping at `s` yields a complete, valid instance.
773    #[must_use]
774    pub fn is_accepting(&self, s: StateId) -> bool {
775        self.engine.is_accepting(s)
776    }
777
778    /// Some byte leads out of `s` toward an accepting state.
779    #[must_use]
780    pub fn can_continue(&self, s: StateId) -> bool {
781        self.engine.can_continue(s)
782    }
783
784    /// `s` is the canonical dead sink.
785    #[must_use]
786    pub fn is_dead(&self, s: StateId) -> bool {
787        self.engine.is_dead(s)
788    }
789
790    /// EOS is legal exactly where `s` accepts. EOS is a side predicate, never a byte or a mask bit.
791    #[must_use]
792    pub fn eos_legal(&self, s: StateId) -> bool {
793        self.engine.eos_legal(s)
794    }
795
796    /// The bytes of `token`, or `None` if the id is not in the vocabulary.
797    #[must_use]
798    pub fn token_bytes(&self, token: TokenId) -> Option<&[u8]> {
799        self.vocab.token_bytes(token.get())
800    }
801
802    /// The ordinary token ids bound to this artifact (EOS excluded). Order is unspecified.
803    #[must_use]
804    pub fn token_ids(&self) -> Vec<TokenId> {
805        self.vocab
806            .iter()
807            .flat_map(|(_, ids)| ids.iter().copied())
808            .map(TokenId)
809            .collect()
810    }
811
812    /// How this artifact produces masks: the naive walk or a prebuilt trie index.
813    #[must_use]
814    pub fn bind_mode(&self) -> BindMode {
815        self.bind_mode
816    }
817
818    /// The allowed ordinary-token mask at `s` (EOS excluded). Routes through the selected bind mode;
819    /// every route is byte-identical to `Naive`.
820    pub fn allowed_mask(&self, s: StateId) -> Result<Bitmask, CompileError> {
821        match self.bind_mode {
822            BindMode::Naive => self.engine.allowed_tokens_from_records(
823                s,
824                self.vocab.mask_vocab_size(),
825                self.vocab.iter(),
826            ),
827            BindMode::TrieJointByte | BindMode::TrieJointClass => self.mask_from_index(s),
828            BindMode::TrieJointByteLazy => self.mask_from_lazy_walk(s),
829            BindMode::TrieJointBytePacked => Ok(self.packed_mask(s)),
830        }
831    }
832
833    /// The compatibility copying view of `s`'s packed mask (an out-of-range state is empty). Native
834    /// callers should use `mask_row_words`/the direct-buffer path, which borrow the row without a copy.
835    fn packed_mask(&self, s: StateId) -> Bitmask {
836        match self.packed.as_ref().and_then(|p| p.row(s)) {
837            Some(words) => Bitmask::from_words(words, self.vocab.mask_vocab_size()),
838            None => Bitmask::zeros(self.vocab.mask_vocab_size()),
839        }
840    }
841
842    /// Walks `s`'s row on demand from the shared byte-trie and packs it into a mask. An out-of-range
843    /// (forged) state is empty, exactly as the eager index and naive walk report it.
844    fn mask_from_lazy_walk(&self, s: StateId) -> Result<Bitmask, CompileError> {
845        let mut mask = Bitmask::zeros(self.vocab.mask_vocab_size());
846        if (s.get() as usize) >= self.engine.state_count() {
847            return Ok(mask);
848        }
849        let trie = self.lazy_trie.as_ref().ok_or_else(|| {
850            CompileError::new(
851                ErrorCode::ArtifactOutOfBounds,
852                Stage::L4Bind,
853                "the lazy bind mode requires a byte-trie",
854            )
855        })?;
856        crate::index::walk_state_mask_into(trie, &self.engine, s, &mut mask)?;
857        Ok(mask)
858    }
859
860    /// Reads the allowed-token mask for `s` from the trie index (an out-of-range state is empty).
861    fn mask_from_index(&self, s: StateId) -> Result<Bitmask, CompileError> {
862        let index = self.index.as_ref().ok_or_else(|| {
863            CompileError::new(
864                ErrorCode::ArtifactOutOfBounds,
865                Stage::L4Bind,
866                "a trie bind mode requires a compiled index",
867            )
868        })?;
869        let mut mask = Bitmask::zeros(self.vocab.mask_vocab_size());
870        if let Some(row) = index.row(s) {
871            for (token, _landing) in row.iter() {
872                mask.set(token)?;
873            }
874        }
875        Ok(mask)
876    }
877
878    /// The ordinary token ids allowed at `s` (EOS excluded), read from the per-state mask cache so
879    /// repeated queries on one state never re-walk the vocabulary. Ids are ascending.
880    pub fn allowed_ids(&self, s: StateId) -> Result<Vec<TokenId>, CompileError> {
881        let mask = self.cached_mask(s)?;
882        let mut out = Vec::with_capacity(mask.count_ones());
883        for (word_index, &word) in mask.as_words().iter().enumerate() {
884            let base = u32::try_from(word_index)
885                .expect("mask word count stays far below u32::MAX (bounded by MAX_MASK_TOKENS)")
886                * 32;
887            let mut bits = word;
888            while bits != 0 {
889                out.push(TokenId(base + bits.trailing_zeros()));
890                bits &= bits - 1;
891            }
892        }
893        Ok(out)
894    }
895
896    /// `allowed_mask(s)`, cached per state id. Out-of-range ids (DEAD-equivalent) never touch the
897    /// cache - only `0..state_count` has a slot.
898    fn cached_mask(&self, s: StateId) -> Result<Arc<Bitmask>, CompileError> {
899        if self.packed.is_some() {
900            return Ok(Arc::new(self.packed_mask(s)));
901        }
902        if (s.get() as usize) >= self.engine.state_count() {
903            return Ok(Arc::new(self.allowed_mask(s)?));
904        }
905        let MaskRows::Lazy(rows) = &self.mask_cache else {
906            unreachable!("non-packed artifact without a lazy mask cache");
907        };
908        let entry_bytes = self.words_per_row() * 4 + MASK_CACHE_ENTRY_OVERHEAD;
909        rows.get_or_compute(s.get() as usize, entry_bytes, || self.allowed_mask(s))
910    }
911
912    #[cfg(test)]
913    fn set_mask_cache_budget(&self, budget: usize) {
914        let MaskRows::Lazy(rows) = &self.mask_cache else {
915            panic!("set_mask_cache_budget called on a packed artifact");
916        };
917        rows.set_budget_for_test(budget);
918    }
919
920    #[cfg(test)]
921    fn mask_cache_filled_count(&self) -> usize {
922        match &self.mask_cache {
923            MaskRows::Lazy(rows) => rows.filled_count(),
924            MaskRows::Disabled => 0,
925        }
926    }
927
928    #[cfg(test)]
929    fn mask_cache_retained_bytes(&self) -> usize {
930        match &self.mask_cache {
931            MaskRows::Lazy(rows) => rows.retained_bytes(),
932            MaskRows::Disabled => 0,
933        }
934    }
935
936    #[cfg(test)]
937    fn mask_cache_is_disabled(&self) -> bool {
938        matches!(self.mask_cache, MaskRows::Disabled)
939    }
940
941    /// The exact packed `u32` length a mask buffer needs, all arithmetic checked (a length mismatch,
942    /// a boundary row index, or an overflowing product is a structured error). Callers size from here.
943    pub fn mask_buffer_len(
944        &self,
945        states_len: usize,
946        rows: Option<&[usize]>,
947    ) -> Result<usize, CompileError> {
948        let rows_count = match rows {
949            Some(r) => {
950                if r.len() != states_len {
951                    return Err(self.mask_err("rows length does not match states length"));
952                }
953                match r.iter().copied().max() {
954                    Some(m) => m
955                        .checked_add(1)
956                        .ok_or_else(|| self.mask_err("row index overflows"))?,
957                    None => 0,
958                }
959            }
960            None => states_len,
961        };
962        let words = rows_count
963            .checked_mul(self.words_per_row())
964            .ok_or_else(|| self.mask_err("mask buffer length overflows"))?;
965        if words > MAX_MASK_WORDS {
966            return Err(self.mask_err("mask buffer length exceeds the resource cap"));
967        }
968        Ok(words)
969    }
970
971    /// Fills caller-owned `out` with the packed masks for `states` (`rows` selects each target row,
972    /// identity when `None`). `out.len()` must equal [`Self::mask_buffer_len`]; unused bits stay zero.
973    pub fn write_mask_into(
974        &self,
975        states: &[StateId],
976        out: &mut [u32],
977        rows: Option<&[usize]>,
978    ) -> Result<(), CompileError> {
979        let wpr = self.words_per_row();
980        let expected = self.mask_buffer_len(states.len(), rows)?;
981        if out.len() != expected {
982            return Err(self.mask_err("mask buffer length does not match rows * words_per_row"));
983        }
984        let overwrite = rows.is_none();
985        if !overwrite {
986            out.iter_mut().for_each(|w| *w = 0);
987        }
988        for (i, &s) in states.iter().enumerate() {
989            let row = rows.map_or(i, |r| r[i]);
990            let base = row * wpr;
991            if let Some(words) = self.mask_row_words(s) {
992                write_words_u32(&mut out[base..base + wpr], words, overwrite);
993            } else {
994                let mask = self.cached_mask(s)?;
995                write_words_u32(&mut out[base..base + wpr], mask.as_words(), overwrite);
996            }
997        }
998        Ok(())
999    }
1000
1001    /// Like `write_mask_into`, but writes little-endian bytes straight from the cached mask into
1002    /// `out` - no `Vec<u32>` intermediate. `out.len()` must equal `4 * mask_buffer_len(...)`.
1003    pub fn write_mask_le_bytes_into(
1004        &self,
1005        states: &[StateId],
1006        out: &mut [u8],
1007        rows: Option<&[usize]>,
1008    ) -> Result<(), CompileError> {
1009        let wpr = self.words_per_row();
1010        let expected_words = self.mask_buffer_len(states.len(), rows)?;
1011        let expected_bytes = expected_words
1012            .checked_mul(4)
1013            .ok_or_else(|| self.mask_err("mask buffer length overflows"))?;
1014        if out.len() != expected_bytes {
1015            return Err(self.mask_err("mask buffer length does not match 4 * rows * words_per_row"));
1016        }
1017        let overwrite = rows.is_none();
1018        if !overwrite {
1019            out.iter_mut().for_each(|b| *b = 0);
1020        }
1021        for (i, &s) in states.iter().enumerate() {
1022            let row = rows.map_or(i, |r| r[i]);
1023            let base = row * wpr * 4;
1024            if let Some(words) = self.mask_row_words(s) {
1025                write_words_le(&mut out[base..base + wpr * 4], words, overwrite);
1026            } else {
1027                let mask = self.cached_mask(s)?;
1028                write_words_le(&mut out[base..base + wpr * 4], mask.as_words(), overwrite);
1029            }
1030        }
1031        Ok(())
1032    }
1033
1034    fn mask_err(&self, msg: &'static str) -> CompileError {
1035        CompileError::new(ErrorCode::ArtifactOutOfBounds, Stage::ArtifactValidate, msg)
1036    }
1037}
1038
1039/// Writes `words` into a `u32` row: overwrite (bulk copy) for a one-to-one row, else OR-accumulate.
1040fn write_words_u32(row: &mut [u32], words: &[u32], overwrite: bool) {
1041    if overwrite {
1042        row.copy_from_slice(words);
1043    } else {
1044        for (dst, &word) in row.iter_mut().zip(words) {
1045            *dst |= word;
1046        }
1047    }
1048}
1049
1050/// Writes `words` as little-endian bytes into a byte row: overwrite for a one-to-one row (no
1051/// read-back), else OR-accumulate. `row.len()` must be `4 * words.len()`.
1052fn write_words_le(row: &mut [u8], words: &[u32], overwrite: bool) {
1053    for (i, &word) in words.iter().enumerate() {
1054        let at = i * 4;
1055        let value = if overwrite {
1056            word
1057        } else {
1058            u32::from_le_bytes([row[at], row[at + 1], row[at + 2], row[at + 3]]) | word
1059        };
1060        row[at..at + 4].copy_from_slice(&value.to_le_bytes());
1061    }
1062}
1063
1064/// Validates a vocabulary binds without ambiguity (the same id -> bytes indexing an artifact runs),
1065/// so a handle checked here reuses across compilations without re-discovering a conflict.
1066pub fn validate_token_bytes(vocab: &Vocabulary) -> Result<(), CompileError> {
1067    CanonicalVocabulary::from_vocabulary(vocab)?;
1068    Ok(())
1069}
1070
1071/// A position inside a [`CompiledArtifact`]. Owns the position, shares the artifact by [`Arc`].
1072#[derive(Clone, Debug)]
1073pub struct Matcher {
1074    state: StateId,
1075    artifact: Arc<CompiledArtifact>,
1076    bytes_seen: u64,
1077}
1078
1079impl Matcher {
1080    /// Starts a matcher at the artifact's anchored start state.
1081    #[must_use]
1082    pub fn new(artifact: Arc<CompiledArtifact>) -> Self {
1083        let state = artifact.start();
1084        Self {
1085            state,
1086            artifact,
1087            bytes_seen: 0,
1088        }
1089    }
1090
1091    /// The current position.
1092    #[must_use]
1093    pub fn state(&self) -> StateId {
1094        self.state
1095    }
1096
1097    /// The number of bytes consumed so far.
1098    #[must_use]
1099    pub fn bytes_seen(&self) -> u64 {
1100        self.bytes_seen
1101    }
1102
1103    /// The shared artifact.
1104    #[must_use]
1105    pub fn artifact(&self) -> &Arc<CompiledArtifact> {
1106        &self.artifact
1107    }
1108
1109    /// Consumes `token`, advancing the position. Fails with `IllegalToken` if the id is unknown or
1110    /// no path exists (a DEAD matcher rejects every token forever; the position never moves on
1111    pub fn advance(&mut self, token: TokenId) -> Result<(), MatcherError> {
1112        let bytes = self
1113            .artifact
1114            .token_bytes(token)
1115            .ok_or(MatcherError::IllegalToken { token })?;
1116        match self.artifact.engine.consume_token(self.state, bytes) {
1117            Some(next) => {
1118                self.state = next;
1119                self.bytes_seen += bytes.len() as u64;
1120                Ok(())
1121            }
1122            None => Err(MatcherError::IllegalToken { token }),
1123        }
1124    }
1125
1126    /// Stopping now yields a complete, valid instance.
1127    #[must_use]
1128    pub fn is_accepting(&self) -> bool {
1129        self.artifact.is_accepting(self.state)
1130    }
1131
1132    /// Some token can still extend this position.
1133    #[must_use]
1134    pub fn can_continue(&self) -> bool {
1135        self.artifact.can_continue(self.state)
1136    }
1137
1138    /// The matcher has reached the permanent dead sink.
1139    #[must_use]
1140    pub fn is_dead(&self) -> bool {
1141        self.artifact.is_dead(self.state)
1142    }
1143
1144    /// EOS is legal here (equivalently, the position accepts).
1145    #[must_use]
1146    pub fn eos_legal(&self) -> bool {
1147        self.artifact.eos_legal(self.state)
1148    }
1149
1150    /// The allowed ordinary-token mask at the current position.
1151    pub fn allowed_mask(&self) -> Result<Bitmask, CompileError> {
1152        self.artifact.allowed_mask(self.state)
1153    }
1154}
1155
1156#[cfg(test)]
1157mod tests {
1158    use super::*;
1159    use crate::automaton::build_from_regex;
1160    use crate::correctness::corpus::corpus_cases;
1161    use crate::error::HashState;
1162    use crate::vocab::build_vocabulary;
1163    use rustc_hash::FxHashMap as Map;
1164
1165    #[cfg(feature = "huggingface-hub")]
1166    #[test]
1167    #[ignore = "timing; run with --release --features huggingface-hub -- --ignored --nocapture"]
1168    fn packed_bind_stage_isolated_gate_ab_on_real_tokenizers() {
1169        use std::time::Instant;
1170
1171        fn stats(mut xs: Vec<f64>) -> (f64, f64, f64, f64, f64, f64) {
1172            xs.sort_by(|a, b| a.partial_cmp(b).unwrap());
1173            let n = xs.len();
1174            let median = xs[n / 2];
1175            let p95 = xs[((n as f64 * 0.95) as usize).min(n - 1)];
1176            let p99 = xs[((n as f64 * 0.99) as usize).min(n - 1)];
1177            let mean = xs.iter().sum::<f64>() / n as f64;
1178            let var = xs.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / n as f64;
1179            (median, p95, p99, var.sqrt(), xs[0], xs[n - 1])
1180        }
1181        fn report(label: &str, xs: Vec<f64>) {
1182            let (median, p95, p99, stddev, min, max) = stats(xs.clone());
1183            println!(
1184                "    {label}: n={} median={median:.2} p95={p95:.2} p99={p99:.2} stddev={stddev:.2} min={min:.2} max={max:.2} us",
1185                xs.len()
1186            );
1187        }
1188
1189        const N: usize = 200;
1190        for repo in [
1191            "openai-community/gpt2",
1192            "Qwen/Qwen2.5-0.5B",
1193            "bigscience/bloom",
1194        ] {
1195            let vocab = Vocabulary::from_pretrained(repo, None).expect("vocab");
1196            let canonical =
1197                Arc::new(CanonicalVocabulary::from_vocabulary(&vocab).expect("canonical"));
1198            let trie = VocabTrie::build_byte(&vocab).expect("trie");
1199            let fresh_engine = || {
1200                corpus_cases()
1201                    .expect("corpus")
1202                    .into_iter()
1203                    .find(|c| c.name == "syn-object-nested")
1204                    .expect("case")
1205                    .engine
1206            };
1207            println!("=== {repo} ===");
1208
1209            let mut buckets: [[Vec<f64>; 5]; 2] = Default::default();
1210            let mut checksums: [Option<Vec<u32>>; 2] = [None, None];
1211            let order: Vec<bool> = (0..N / 2)
1212                .flat_map(|i| {
1213                    if i % 2 == 0 {
1214                        [true, false]
1215                    } else {
1216                        [false, true]
1217                    }
1218                })
1219                .collect();
1220            const WARMUP: usize = 20;
1221            for &use_gate in std::iter::repeat_n(&true, WARMUP).chain(order.iter()) {
1222                let gi = usize::from(!use_gate);
1223                let engine = fresh_engine();
1224                let n = engine.state_count();
1225                let wpr = canonical.mask_width_words();
1226                let payload = n.checked_mul(wpr).unwrap().checked_mul(4).unwrap();
1227                let scratch_fixed = crate::index::packed_scratch_fixed_bytes().unwrap();
1228                let stack_bytes =
1229                    crate::index::packed_stack_bytes(trie.max_dfs_frontier()).unwrap();
1230                let peak = payload + PACKED_HEADER_ALLOWANCE + scratch_fixed + stack_bytes;
1231
1232                let t_total = Instant::now();
1233                let t = Instant::now();
1234                let _permit = if use_gate {
1235                    Some(crate::mem_gate::acquire_vocab_build_bytes(peak).expect("permit"))
1236                } else {
1237                    None
1238                };
1239                let gate_dt = t.elapsed().as_secs_f64() * 1e6;
1240
1241                let t = Instant::now();
1242                let stack_fuel = trie.max_dfs_frontier();
1243                let word_count = payload / 4;
1244                let mut rows: Vec<u32> = Vec::new();
1245                rows.try_reserve_exact(word_count).expect("reserve");
1246                rows.resize(word_count, 0);
1247                let alloc_dt = t.elapsed().as_secs_f64() * 1e6;
1248
1249                let t = Instant::now();
1250                crate::index::fill_packed_rows(
1251                    &trie,
1252                    &engine,
1253                    &mut rows,
1254                    wpr,
1255                    canonical.mask_vocab_size(),
1256                    stack_fuel,
1257                )
1258                .expect("fill");
1259                let fill_dt = t.elapsed().as_secs_f64() * 1e6;
1260
1261                let t = Instant::now();
1262                let mut artifact = CompiledArtifact::new_with_canonical(
1263                    engine,
1264                    canonical.clone(),
1265                    Provenance::reference(HashState::Hash([7; 32])),
1266                )
1267                .expect("artifact");
1268                artifact.bind_mode = BindMode::TrieJointBytePacked;
1269                artifact.packed = Some(PackedMaskTable {
1270                    words_per_row: wpr,
1271                    rows: rows.into_boxed_slice(),
1272                });
1273                let construct_dt = t.elapsed().as_secs_f64() * 1e6;
1274                let total_dt = t_total.elapsed().as_secs_f64() * 1e6;
1275
1276                if checksums[gi].is_none() {
1277                    checksums[gi] = Some(artifact.packed_mask_table_words().unwrap().to_vec());
1278                }
1279                std::hint::black_box(&artifact);
1280
1281                buckets[gi][0].push(gate_dt);
1282                buckets[gi][1].push(alloc_dt);
1283                buckets[gi][2].push(fill_dt);
1284                buckets[gi][3].push(construct_dt);
1285                buckets[gi][4].push(total_dt);
1286            }
1287            for (gi, use_gate) in [(0, true), (1, false)] {
1288                println!("  use_gate={use_gate}:");
1289                let [gate_us, alloc_us, fill_us, construct_us, total_us] = buckets[gi].clone();
1290                report("gate_acquire", gate_us);
1291                report("payload_alloc", alloc_us);
1292                report("fill_packed_rows", fill_us);
1293                report("artifact_construct", construct_us);
1294                report("total", total_us);
1295            }
1296            assert_eq!(
1297                checksums[0], checksums[1],
1298                "gate on/off must produce byte-identical packed tables ({repo})"
1299            );
1300        }
1301    }
1302
1303    #[cfg(feature = "huggingface-hub")]
1304    #[test]
1305    #[ignore = "timing; run with --release --features huggingface-hub -- --ignored --nocapture"]
1306    fn warm_bind_stage_breakdown_on_real_tokenizers() {
1307        use std::time::Instant;
1308        let median = |mut xs: Vec<u128>| {
1309            xs.sort_unstable();
1310            xs[xs.len() / 2]
1311        };
1312        let fresh_engine = || {
1313            corpus_cases()
1314                .expect("corpus")
1315                .into_iter()
1316                .find(|c| c.name == "syn-object-nested")
1317                .expect("case")
1318                .engine
1319        };
1320        for repo in ["openai-community/gpt2", "Qwen/Qwen2.5-0.5B"] {
1321            let vocab = Arc::new(Vocabulary::from_pretrained(repo, None).expect("vocab"));
1322            let mut meta_us = Vec::new();
1323            let mut new_us = Vec::new();
1324            for _ in 0..11 {
1325                let t = Instant::now();
1326                let idx = CanonicalVocabulary::from_vocabulary(&vocab).expect("canonical");
1327                meta_us.push(t.elapsed().as_micros());
1328                std::hint::black_box(&idx);
1329                let t = Instant::now();
1330                let a = CompiledArtifact::new(
1331                    fresh_engine(),
1332                    vocab.clone(),
1333                    Provenance::reference(HashState::Hash([7; 32])),
1334                )
1335                .expect("artifact");
1336                new_us.push(t.elapsed().as_micros());
1337                std::hint::black_box(&a);
1338            }
1339            println!(
1340                "{repo}: canonical_build_us={} CompiledArtifact_new_us={} tokens={}",
1341                median(meta_us),
1342                median(new_us),
1343                vocab.len()
1344            );
1345        }
1346    }
1347
1348    /// Packs `vocab` into the four CSR buffers `VocabularyHandle::from_packed` expects.
1349    #[cfg(feature = "huggingface-hub")]
1350    fn pack_vocab(vocab: &Vocabulary) -> (Vec<u8>, Vec<u32>, Vec<u32>, Vec<u32>) {
1351        let mut token_bytes = Vec::new();
1352        let mut byte_offsets = vec![0u32];
1353        let mut token_ids = Vec::new();
1354        let mut id_offsets = vec![0u32];
1355        for (bytes, ids) in vocab.tokens() {
1356            token_bytes.extend_from_slice(bytes);
1357            byte_offsets.push(token_bytes.len() as u32);
1358            token_ids.extend_from_slice(ids);
1359            id_offsets.push(token_ids.len() as u32);
1360        }
1361        (token_bytes, byte_offsets, token_ids, id_offsets)
1362    }
1363
1364    /// Stage-isolated: tuple-vs-packed CONSTRUCTION (sort/flatten vs full handle build), then a
1365    /// cold BIND breakdown (schema IR, DFA compile, trie build, artifact assembly, first-vs-repeated
1366    #[cfg(feature = "huggingface-hub")]
1367    #[test]
1368    #[ignore = "timing; run with --release --features huggingface-hub -- --ignored --nocapture"]
1369    fn stage_isolated_construction_and_bind_on_real_tokenizers() {
1370        use std::time::Instant;
1371
1372        use crate::index::TrieCache;
1373
1374        fn stats(mut xs: Vec<f64>) -> (f64, f64, f64) {
1375            xs.sort_by(|a, b| a.partial_cmp(b).unwrap());
1376            let n = xs.len();
1377            (
1378                xs[n / 2],
1379                xs[((n as f64 * 0.95) as usize).min(n - 1)],
1380                xs[n - 1],
1381            )
1382        }
1383        fn report(label: &str, xs: Vec<f64>) {
1384            let (median, p95, max) = stats(xs.clone());
1385            println!(
1386                "    {label}: n={} median={median:.2} p95={p95:.2} max={max:.2} us",
1387                xs.len()
1388            );
1389        }
1390
1391        const N: usize = 30;
1392        for repo in [
1393            "openai-community/gpt2",
1394            "Qwen/Qwen2.5-0.5B",
1395            "bigscience/bloom",
1396        ] {
1397            let vocab = Vocabulary::from_pretrained(repo, None).expect("vocab");
1398            let (tb, bo, ti, io) = pack_vocab(&vocab);
1399            println!("=== {repo} ({} tokens) ===", vocab.len());
1400
1401            let mut tuple_sort_us = Vec::new();
1402            let mut tuple_total_us = Vec::new();
1403            let mut packed_sort_us = Vec::new();
1404            let mut packed_total_us = Vec::new();
1405            for _ in 0..N {
1406                let t = Instant::now();
1407                let prepared = crate::vocab::PreparedVocabulary::build(&vocab).expect("prepared");
1408                tuple_sort_us.push(t.elapsed().as_secs_f64() * 1e6);
1409                std::hint::black_box(&prepared);
1410
1411                let t = Instant::now();
1412                let canonical = CanonicalVocabulary::from_vocabulary(&vocab).expect("canonical");
1413                tuple_total_us.push(t.elapsed().as_secs_f64() * 1e6);
1414                std::hint::black_box(&canonical);
1415
1416                let t = Instant::now();
1417                let prepared = crate::vocab::PreparedVocabulary::build_from_packed(
1418                    &tb,
1419                    &bo,
1420                    &ti,
1421                    &io,
1422                    vocab.eos_token_id(),
1423                )
1424                .expect("prepared");
1425                packed_sort_us.push(t.elapsed().as_secs_f64() * 1e6);
1426                std::hint::black_box(&prepared);
1427
1428                let t = Instant::now();
1429                let handle =
1430                    VocabularyHandle::from_packed(&tb, &bo, &ti, &io, vocab.eos_token_id())
1431                        .expect("handle");
1432                packed_total_us.push(t.elapsed().as_secs_f64() * 1e6);
1433                std::hint::black_box(&handle);
1434            }
1435            println!("  construction (sort/flatten only, then the whole handle build):");
1436            report("tuple_sort_flatten", tuple_sort_us);
1437            report("tuple_total_handle_build", tuple_total_us);
1438            report("packed_sort_flatten", packed_sort_us);
1439            report("packed_total_handle_build", packed_total_us);
1440
1441            let schema = crate::correctness::corpus::corpus_cases()
1442                .expect("corpus")
1443                .into_iter()
1444                .find(|c| c.name == "syn-object-nested")
1445                .expect("case")
1446                .json_schema
1447                .to_string();
1448            let handle = VocabularyHandle::new(Arc::new(vocab.clone())).expect("handle");
1449
1450            let mut ir_us = Vec::new();
1451            let mut dfa_us = Vec::new();
1452            let mut trie_build_us = Vec::new();
1453            let mut artifact_us = Vec::new();
1454            let mut first_mask_us = Vec::new();
1455            let mut repeated_mask_us = Vec::new();
1456            for _ in 0..N {
1457                let t = Instant::now();
1458                let ir =
1459                    crate::schema_to_ir(&schema, crate::CompileOptions::default()).expect("ir");
1460                ir_us.push(t.elapsed().as_secs_f64() * 1e6);
1461
1462                let t = Instant::now();
1463                let engine = crate::compile_ir(&ir).expect("engine");
1464                dfa_us.push(t.elapsed().as_secs_f64() * 1e6);
1465
1466                let t = Instant::now();
1467                let cache = TrieCache::new();
1468                let bound = cache.bind(&handle).expect("bound");
1469                trie_build_us.push(t.elapsed().as_secs_f64() * 1e6);
1470
1471                let t = Instant::now();
1472                let artifact = CompiledArtifact::new_from_bound_trie(
1473                    &handle,
1474                    engine,
1475                    Provenance::reference(HashState::Hash([7; 32])),
1476                    BindMode::TrieJointBytePacked,
1477                    &bound,
1478                )
1479                .expect("artifact");
1480                artifact_us.push(t.elapsed().as_secs_f64() * 1e6);
1481
1482                let start = artifact.start();
1483                let wpr = artifact.words_per_row();
1484                let mut buf = vec![0u8; wpr * 4];
1485                let t = Instant::now();
1486                artifact
1487                    .write_mask_le_bytes_into(&[start], &mut buf, None)
1488                    .expect("mask");
1489                first_mask_us.push(t.elapsed().as_secs_f64() * 1e6);
1490
1491                let t = Instant::now();
1492                artifact
1493                    .write_mask_le_bytes_into(&[start], &mut buf, None)
1494                    .expect("mask");
1495                repeated_mask_us.push(t.elapsed().as_secs_f64() * 1e6);
1496            }
1497            println!("  cold bind (fresh trie every iteration):");
1498            report("schema_to_ir", ir_us);
1499            report("compile_ir (DFA)", dfa_us);
1500            report(
1501                "trie_build (cold, via fresh TrieCache::bind)",
1502                trie_build_us,
1503            );
1504            report("artifact_assembly (packed table fill)", artifact_us);
1505            report("first_direct_mask (same API as repeated)", first_mask_us);
1506            report("repeated_direct_mask", repeated_mask_us);
1507        }
1508    }
1509
1510    fn boolean_artifact() -> CompiledArtifact {
1511        let engine = corpus_cases()
1512            .expect("corpus")
1513            .into_iter()
1514            .find(|c| c.name == "suite-type-boolean")
1515            .expect("boolean case")
1516            .engine;
1517        let mut map: Map<Vec<u8>, Vec<u32>> = Map::default();
1518        map.insert(b"t".to_vec(), vec![0]);
1519        map.insert(b"true".to_vec(), vec![1]);
1520        map.insert(b"z".to_vec(), vec![2, 3]);
1521        let vocab = build_vocabulary(9, map).expect("vocab");
1522        CompiledArtifact::new(
1523            engine,
1524            Arc::new(vocab),
1525            Provenance::reference(HashState::Hash([7; 32])),
1526        )
1527        .expect("artifact")
1528    }
1529
1530    fn boolean_lazy_artifact() -> CompiledArtifact {
1531        let engine = corpus_cases()
1532            .expect("corpus")
1533            .into_iter()
1534            .find(|c| c.name == "suite-type-boolean")
1535            .expect("boolean case")
1536            .engine;
1537        let mut map: Map<Vec<u8>, Vec<u32>> = Map::default();
1538        map.insert(b"t".to_vec(), vec![0]);
1539        map.insert(b"true".to_vec(), vec![1]);
1540        map.insert(b"z".to_vec(), vec![2, 3]);
1541        let vocab = Arc::new(build_vocabulary(9, map).expect("vocab"));
1542        let fp = VocabFingerprint::of(&vocab).unwrap();
1543        let canonical = Arc::new(CanonicalVocabulary::from_vocabulary(&vocab).expect("canonical"));
1544        let trie = Arc::new(VocabTrie::build_byte(&vocab).expect("trie"));
1545        CompiledArtifact::new_lazy_trie(
1546            engine,
1547            canonical,
1548            Provenance::reference(HashState::Hash([7; 32])),
1549            trie,
1550            fp,
1551        )
1552        .expect("lazy artifact")
1553    }
1554
1555    #[test]
1556    fn packed_payload_reservation_is_fallible_not_an_abort() {
1557        let mut rows: Vec<u32> = Vec::new();
1558        assert!(rows.try_reserve_exact(usize::MAX / 2).is_err());
1559    }
1560
1561    #[test]
1562    fn packed_admission_charges_table_scratch_and_stack_and_rejects_before_allocating() {
1563        let vocab = Arc::new({
1564            let mut map: Map<Vec<u8>, Vec<u32>> = Map::default();
1565            map.insert(b"t".to_vec(), vec![0]);
1566            build_vocabulary(9, map).expect("vocab")
1567        });
1568        let handle = crate::index::VocabularyHandle::new(vocab.clone()).expect("handle");
1569        let trie = VocabTrie::build_byte(&vocab).expect("trie");
1570        let n = boolean_engine().state_count();
1571        let wpr = handle.canonical().mask_width_words();
1572        let scratch = crate::index::packed_scratch_fixed_bytes().unwrap();
1573        let stack = crate::index::packed_stack_bytes(trie.max_dfs_frontier()).unwrap();
1574        assert!(stack > 0, "the DFS stack charge is load-bearing");
1575        let payload_header_scratch = n * wpr * 4 + PACKED_HEADER_ALLOWANCE + scratch;
1576        let peak = payload_header_scratch + stack;
1577
1578        let budget = |b: usize| {
1579            CompiledArtifact::new_packed_from_trie_budgeted(
1580                boolean_engine(),
1581                handle.canonical().clone(),
1582                prov_bytes([4; 32]),
1583                &trie,
1584                handle.fingerprint(),
1585                b,
1586            )
1587        };
1588        assert_eq!(
1589            budget(payload_header_scratch).unwrap_err().code,
1590            ErrorCode::InternalLimitExceeded
1591        );
1592        assert_eq!(
1593            budget(peak - 1).unwrap_err().code,
1594            ErrorCode::InternalLimitExceeded
1595        );
1596        assert_eq!(
1597            budget(peak + 1).expect("peak + 1 binds").bind_mode(),
1598            BindMode::TrieJointBytePacked
1599        );
1600    }
1601
1602    #[test]
1603    fn packed_admission_rejects_when_scratch_alone_overflows_an_otherwise_fitting_payload() {
1604        let vocab = Arc::new({
1605            let mut map: Map<Vec<u8>, Vec<u32>> = Map::default();
1606            map.insert(b"a".to_vec(), vec![0]);
1607            build_vocabulary(1, map).expect("vocab")
1608        });
1609        let handle = crate::index::VocabularyHandle::new(vocab.clone()).expect("handle");
1610        let trie = VocabTrie::build_byte(&vocab).expect("trie");
1611        let n = build_from_regex("a{1,200}").expect("engine").state_count();
1612        let wpr = handle.canonical().mask_width_words();
1613        let scratch = crate::index::packed_scratch_fixed_bytes().unwrap();
1614        assert!(
1615            scratch > 4096,
1616            "the dense scratch is the load-bearing charge ({scratch} bytes)"
1617        );
1618        let budget = n * wpr * 4 + PACKED_HEADER_ALLOWANCE + (scratch / 2);
1619        let err = CompiledArtifact::new_packed_from_trie_budgeted(
1620            build_from_regex("a{1,200}").unwrap(),
1621            handle.canonical().clone(),
1622            prov_bytes([4; 32]),
1623            &trie,
1624            handle.fingerprint(),
1625            budget,
1626        )
1627        .expect_err("scratch charged on top of a fitting payload must reject");
1628        assert_eq!(err.code, ErrorCode::InternalLimitExceeded);
1629    }
1630
1631    #[test]
1632    fn packed_admission_holds_for_many_sparse_states_one_token_per_row() {
1633        let vocab = Arc::new({
1634            let mut map: Map<Vec<u8>, Vec<u32>> = Map::default();
1635            map.insert(b"a".to_vec(), vec![0]);
1636            build_vocabulary(1, map).expect("vocab")
1637        });
1638        let handle = crate::index::VocabularyHandle::new(vocab.clone()).expect("handle");
1639        let trie = VocabTrie::build_byte(&vocab).expect("trie");
1640        let n = build_from_regex("a{1,200}").expect("engine").state_count();
1641        assert!(n > 100, "the chain must have many states ({n})");
1642        let wpr = handle.canonical().mask_width_words();
1643        let scratch = crate::index::packed_scratch_fixed_bytes().unwrap();
1644        let fixed = n * wpr * 4 + PACKED_HEADER_ALLOWANCE + scratch;
1645
1646        let under = CompiledArtifact::new_packed_from_trie_budgeted(
1647            build_from_regex("a{1,200}").unwrap(),
1648            handle.canonical().clone(),
1649            prov_bytes([4; 32]),
1650            &trie,
1651            handle.fingerprint(),
1652            fixed - 1,
1653        )
1654        .expect_err("one byte under the table+scratch total must reject");
1655        assert_eq!(under.code, ErrorCode::InternalLimitExceeded);
1656
1657        let packed = CompiledArtifact::new_packed_from_trie(
1658            build_from_regex("a{1,200}").unwrap(),
1659            handle.canonical().clone(),
1660            prov_bytes([4; 32]),
1661            &trie,
1662            handle.fingerprint(),
1663        )
1664        .expect("generous budget binds the many-state chain");
1665        let naive = CompiledArtifact::new(
1666            build_from_regex("a{1,200}").unwrap(),
1667            vocab,
1668            prov_bytes([4; 32]),
1669        )
1670        .expect("naive");
1671        for raw in 0..naive.engine.state_count() {
1672            let s = StateId(raw as u32);
1673            assert_eq!(
1674                packed.allowed_mask(s).unwrap().as_words(),
1675                naive.allowed_mask(s).unwrap().as_words(),
1676                "packed != naive at state {raw}"
1677            );
1678        }
1679    }
1680
1681    #[test]
1682    fn recommended_bind_mode_is_a_no_op_for_every_non_packed_request() {
1683        let vocab = Arc::new({
1684            let mut map: Map<Vec<u8>, Vec<u32>> = Map::default();
1685            map.insert(b"a".to_vec(), vec![0]);
1686            build_vocabulary(1, map).expect("vocab")
1687        });
1688        let handle = crate::index::VocabularyHandle::new(vocab).expect("handle");
1689        let bound = crate::index::TrieCache::new().bind(&handle).expect("bound");
1690        let n = boolean_engine().state_count();
1691        for requested in [
1692            BindMode::Naive,
1693            BindMode::TrieJointByte,
1694            BindMode::TrieJointByteLazy,
1695            BindMode::TrieJointClass,
1696        ] {
1697            assert_eq!(
1698                CompiledArtifact::recommended_bind_mode(
1699                    n,
1700                    handle.mask_vocab_size(),
1701                    &bound,
1702                    requested
1703                ),
1704                requested,
1705                "{requested:?} must pass through unchanged"
1706            );
1707        }
1708    }
1709
1710    #[test]
1711    fn recommended_bind_mode_keeps_packed_when_the_estimate_fits_the_real_budget() {
1712        let vocab = Arc::new({
1713            let mut map: Map<Vec<u8>, Vec<u32>> = Map::default();
1714            map.insert(b"a".to_vec(), vec![0]);
1715            build_vocabulary(1, map).expect("vocab")
1716        });
1717        let handle = crate::index::VocabularyHandle::new(vocab).expect("handle");
1718        let bound = crate::index::TrieCache::new().bind(&handle).expect("bound");
1719        let n = boolean_engine().state_count();
1720        assert_eq!(
1721            CompiledArtifact::recommended_bind_mode(
1722                n,
1723                handle.mask_vocab_size(),
1724                &bound,
1725                BindMode::TrieJointBytePacked
1726            ),
1727            BindMode::TrieJointBytePacked,
1728            "a tiny grammar/vocabulary must stay packed"
1729        );
1730    }
1731
1732    #[test]
1733    fn recommended_bind_mode_downgrades_on_latency_alone_well_under_the_byte_budget() {
1734        let vocab = Arc::new({
1735            let mut map: Map<Vec<u8>, Vec<u32>> = Map::default();
1736            map.insert(b"a".to_vec(), vec![0]);
1737            build_vocabulary(1, map).expect("vocab")
1738        });
1739        let handle = crate::index::VocabularyHandle::new(vocab).expect("handle");
1740        let bound = crate::index::TrieCache::new().bind(&handle).expect("bound");
1741        let state_count =
1742            (PACKED_COMPILE_LATENCY_BUDGET_US * 1000) / PACKED_ROW_FILL_NS_PER_WORD + 1000;
1743        assert!(
1744            state_count * 4 < MAX_PACKED_MASK_BYTES,
1745            "fixture sanity: byte budget not the cause"
1746        );
1747        assert_eq!(
1748            CompiledArtifact::recommended_bind_mode(
1749                state_count,
1750                handle.mask_vocab_size(),
1751                &bound,
1752                BindMode::TrieJointBytePacked
1753            ),
1754            BindMode::TrieJointByteLazy,
1755            "latency alone must downgrade even when bytes are nowhere near the budget"
1756        );
1757    }
1758
1759    #[test]
1760    fn recommended_bind_mode_downgrades_to_lazy_when_the_payload_alone_exceeds_the_real_budget() {
1761        let vocab = Arc::new({
1762            let mut map: Map<Vec<u8>, Vec<u32>> = Map::default();
1763            map.insert(b"a".to_vec(), vec![0]);
1764            build_vocabulary(1, map).expect("vocab")
1765        });
1766        let handle = crate::index::VocabularyHandle::new(vocab).expect("handle");
1767        let bound = crate::index::TrieCache::new().bind(&handle).expect("bound");
1768        let wpr = handle.mask_vocab_size().div_ceil(32);
1769        assert_eq!(
1770            wpr, 1,
1771            "fixture sanity: one row word, so payload = state_count * 4"
1772        );
1773        let huge_state_count = MAX_PACKED_MASK_BYTES / 4 + 1; // payload alone > MAX_PACKED_MASK_BYTES
1774
1775        assert_eq!(
1776            CompiledArtifact::recommended_bind_mode(
1777                huge_state_count,
1778                handle.mask_vocab_size(),
1779                &bound,
1780                BindMode::TrieJointBytePacked,
1781            ),
1782            BindMode::TrieJointByteLazy,
1783            "a payload alone past the fixed packed budget must downgrade to lazy"
1784        );
1785    }
1786
1787    #[test]
1788    fn recommended_bind_mode_never_recommends_packed_when_the_real_build_would_reject() {
1789        let vocab = Arc::new({
1790            let mut map: Map<Vec<u8>, Vec<u32>> = Map::default();
1791            map.insert(b"a".to_vec(), vec![0]);
1792            build_vocabulary(1, map).expect("vocab")
1793        });
1794        let handle = crate::index::VocabularyHandle::new(vocab.clone()).expect("handle");
1795        let trie = VocabTrie::build_byte(&vocab).expect("trie");
1796        let n = build_from_regex("a{1,200}").expect("engine").state_count();
1797        let wpr = handle.canonical().mask_width_words();
1798        let peak = CompiledArtifact::packed_build_peak_bytes(n, wpr, &trie).expect("peak");
1799
1800        CompiledArtifact::new_packed_from_trie_budgeted(
1801            build_from_regex("a{1,200}").unwrap(),
1802            handle.canonical().clone(),
1803            prov_bytes([4; 32]),
1804            &trie,
1805            handle.fingerprint(),
1806            peak,
1807        )
1808        .expect("the real build must bind at exactly `peak`");
1809        CompiledArtifact::new_packed_from_trie_budgeted(
1810            build_from_regex("a{1,200}").unwrap(),
1811            handle.canonical().clone(),
1812            prov_bytes([4; 32]),
1813            &trie,
1814            handle.fingerprint(),
1815            peak - 1,
1816        )
1817        .expect_err("the real build must reject one byte under `peak`");
1818        assert!(
1819            peak <= MAX_PACKED_MASK_BYTES,
1820            "fixture sanity: this small chain must fit the real production budget"
1821        );
1822
1823        let bound = crate::index::TrieCache::new().bind(&handle).expect("bound");
1824        assert_eq!(
1825            CompiledArtifact::recommended_bind_mode(
1826                n,
1827                handle.mask_vocab_size(),
1828                &bound,
1829                BindMode::TrieJointBytePacked,
1830            ),
1831            BindMode::TrieJointBytePacked
1832        );
1833    }
1834
1835    #[test]
1836    fn new_trie_rejects_lazy_and_naive_with_a_typed_error() {
1837        let vocab = Arc::new({
1838            let mut map: Map<Vec<u8>, Vec<u32>> = Map::default();
1839            map.insert(b"t".to_vec(), vec![0]);
1840            build_vocabulary(9, map).expect("vocab")
1841        });
1842        let trie = VocabTrie::build_byte(&vocab).expect("trie");
1843        for mode in [BindMode::TrieJointByteLazy, BindMode::Naive] {
1844            let err = CompiledArtifact::new_trie(
1845                boolean_engine(),
1846                vocab.clone(),
1847                prov_bytes([2; 32]),
1848                mode,
1849                &trie,
1850            )
1851            .expect_err("lazy/naive are not constructible from a borrowed trie");
1852            assert_eq!(err.code, ErrorCode::ArtifactOutOfBounds);
1853        }
1854    }
1855
1856    #[test]
1857    fn new_trie_builds_packed_equal_to_naive() {
1858        let mut map: Map<Vec<u8>, Vec<u32>> = Map::default();
1859        map.insert(b"t".to_vec(), vec![0]);
1860        map.insert(b"true".to_vec(), vec![1]);
1861        map.insert(b"z".to_vec(), vec![2, 3]);
1862        let vocab = Arc::new(build_vocabulary(9, map).expect("vocab"));
1863        let trie = VocabTrie::build_byte(&vocab).expect("trie");
1864        let packed = CompiledArtifact::new_trie(
1865            boolean_engine(),
1866            vocab.clone(),
1867            prov_bytes([6; 32]),
1868            BindMode::TrieJointBytePacked,
1869            &trie,
1870        )
1871        .expect("packed via new_trie");
1872        let naive = boolean_artifact();
1873        assert_eq!(packed.bind_mode(), BindMode::TrieJointBytePacked);
1874        for raw in 0..packed.engine.state_count() {
1875            let s = StateId(raw as u32);
1876            assert_eq!(
1877                packed.allowed_mask(s).unwrap().as_words(),
1878                naive.allowed_mask(s).unwrap().as_words()
1879            );
1880        }
1881    }
1882
1883    fn object_with_ws_engine() -> crate::automaton::RefEngine {
1884        use crate::compile::compile_ir;
1885        use crate::ir::{Builder, CompileOptions, ObjectClosure};
1886        let mut b = Builder::new(CompileOptions::default());
1887        let boolean = b.boolean().unwrap();
1888        let obj = b
1889            .object(
1890                vec![("a".to_string(), boolean)],
1891                &[true],
1892                ObjectClosure::AssumeClosedProfile,
1893            )
1894            .unwrap();
1895        let ir = b.finish(obj).unwrap();
1896        compile_ir(&ir).unwrap()
1897    }
1898
1899    /// The Task 2 acceptance gate: a spaced instance must be accepted, and every bind mode must
1900    /// compute the identical mask, since whitespace lives in the shared engine, not the bind layer.
1901    #[test]
1902    fn whitespace_variants_agree_across_naive_packed_and_lazy_bind_modes() {
1903        let mut map: Map<Vec<u8>, Vec<u32>> = Map::default();
1904        for (bytes, id) in [
1905            (b"{".as_slice(), 0u32),
1906            (b"\"a\"".as_slice(), 1),
1907            (b":".as_slice(), 2),
1908            (b" ".as_slice(), 3),
1909            (b"true".as_slice(), 4),
1910            (b"}".as_slice(), 5),
1911        ] {
1912            map.insert(bytes.to_vec(), vec![id]);
1913        }
1914        let vocab = Arc::new(build_vocabulary(6, map).expect("vocab"));
1915
1916        let naive =
1917            CompiledArtifact::new(object_with_ws_engine(), vocab.clone(), prov_bytes([11; 32]))
1918                .expect("naive");
1919        let trie = VocabTrie::build_byte(&vocab).expect("trie");
1920        let packed = CompiledArtifact::new_trie(
1921            object_with_ws_engine(),
1922            vocab.clone(),
1923            prov_bytes([12; 32]),
1924            BindMode::TrieJointBytePacked,
1925            &trie,
1926        )
1927        .expect("packed");
1928        let fp = VocabFingerprint::of(&vocab).unwrap();
1929        let canonical = Arc::new(CanonicalVocabulary::from_vocabulary(&vocab).expect("canonical"));
1930        let lazy = CompiledArtifact::new_lazy_trie(
1931            object_with_ws_engine(),
1932            canonical,
1933            prov_bytes([13; 32]),
1934            Arc::new(VocabTrie::build_byte(&vocab).expect("trie")),
1935            fp,
1936        )
1937        .expect("lazy");
1938
1939        for raw in 0..naive.engine.state_count() {
1940            let s = StateId(raw as u32);
1941            let n = naive.allowed_mask(s).expect("naive mask");
1942            assert_eq!(
1943                n.as_words(),
1944                packed.allowed_mask(s).expect("packed mask").as_words()
1945            );
1946            assert_eq!(
1947                n.as_words(),
1948                lazy.allowed_mask(s).expect("lazy mask").as_words()
1949            );
1950        }
1951
1952        for artifact in [&naive, &packed, &lazy] {
1953            let mut s = artifact.start();
1954            for tok in [b"{".as_slice(), b"\"a\"", b":", b" ", b"true", b"}"] {
1955                s = artifact.walk(s, tok).expect("spaced token sequence walks");
1956            }
1957            assert!(
1958                artifact.is_accepting(s),
1959                "{:?} must accept {{\"a\": true}}",
1960                artifact.bind_mode()
1961            );
1962        }
1963    }
1964
1965    #[test]
1966    fn packed_mask_table_words_is_the_whole_table_row_major_matching_mask_row_words() {
1967        let mut map: Map<Vec<u8>, Vec<u32>> = Map::default();
1968        map.insert(b"t".to_vec(), vec![0]);
1969        map.insert(b"true".to_vec(), vec![1]);
1970        map.insert(b"z".to_vec(), vec![2, 3]);
1971        let vocab = Arc::new(build_vocabulary(9, map).expect("vocab"));
1972        let trie = VocabTrie::build_byte(&vocab).expect("trie");
1973        let packed = CompiledArtifact::new_trie(
1974            boolean_engine(),
1975            vocab,
1976            prov_bytes([7; 32]),
1977            BindMode::TrieJointBytePacked,
1978            &trie,
1979        )
1980        .expect("packed via new_trie");
1981
1982        let wpr = packed.words_per_row();
1983        let table = packed
1984            .packed_mask_table_words()
1985            .expect("packed mode always has a flat table");
1986        assert_eq!(table.len(), packed.state_count() * wpr);
1987
1988        for raw in 0..packed.state_count() {
1989            let s = StateId(raw as u32);
1990            let row_via_lookup = packed.mask_row_words(s).unwrap();
1991            let row_via_table = &table[raw * wpr..(raw + 1) * wpr];
1992            assert_eq!(
1993                row_via_table, row_via_lookup,
1994                "the flat table's row {raw} must match the single-row accessor byte for byte"
1995            );
1996        }
1997    }
1998
1999    #[test]
2000    fn packed_mask_table_words_is_none_outside_packed_mode() {
2001        let naive = boolean_artifact();
2002        assert_eq!(naive.bind_mode(), BindMode::Naive);
2003        assert!(naive.packed_mask_table_words().is_none());
2004    }
2005
2006    #[test]
2007    fn new_from_handle_equals_new_and_cannot_mispair() {
2008        let mut map: Map<Vec<u8>, Vec<u32>> = Map::default();
2009        map.insert(b"t".to_vec(), vec![0]);
2010        map.insert(b"true".to_vec(), vec![1]);
2011        map.insert(b"z".to_vec(), vec![2]);
2012        let vocab = Arc::new(build_vocabulary(9, map).expect("vocab"));
2013        let handle = crate::index::VocabularyHandle::new(vocab.clone()).expect("handle");
2014        let engine = boolean_engine();
2015        let shared = CompiledArtifact::new_from_handle(&handle, engine, prov_bytes([5; 32]))
2016            .expect("handle");
2017        let fresh =
2018            CompiledArtifact::new(boolean_engine(), vocab, prov_bytes([5; 32])).expect("fresh");
2019        for raw in 0..shared.engine.state_count() {
2020            let s = StateId(raw as u32);
2021            assert_eq!(
2022                shared.allowed_mask(s).unwrap().as_words(),
2023                fresh.allowed_mask(s).unwrap().as_words(),
2024                "handle-shared metadata must equal freshly-built metadata at {raw}"
2025            );
2026        }
2027    }
2028
2029    fn prov_bytes(h: [u8; 32]) -> Provenance {
2030        Provenance::reference(HashState::Hash(h))
2031    }
2032
2033    #[test]
2034    fn lazy_cached_mask_walks_once_then_reuses_the_arc_and_equals_naive() {
2035        let lazy = boolean_lazy_artifact();
2036        let naive = boolean_artifact();
2037        for raw in 0..lazy.engine.state_count() {
2038            let s = StateId(raw as u32);
2039            let first = lazy.cached_mask(s).expect("miss walks the trie");
2040            let second = lazy.cached_mask(s).expect("hit reuses the cache");
2041            assert!(
2042                Arc::ptr_eq(&first, &second),
2043                "a repeated lazy query must reuse the cached mask, not re-walk"
2044            );
2045            assert_eq!(
2046                first.as_words(),
2047                naive.allowed_mask(s).expect("naive").as_words(),
2048                "lazy direct-mask walk must equal naive at state {raw}"
2049            );
2050        }
2051    }
2052
2053    #[test]
2054    fn growing_lazy_rows_never_exceed_the_charged_upper_bound() {
2055        let lazy = boolean_lazy_artifact();
2056        let upper_bound = lazy.artifact_cache_charge_upper_bound();
2057        for raw in 0..lazy.engine.state_count() {
2058            let s = StateId(raw as u32);
2059            let _ = lazy.cached_mask(s).expect("mask");
2060            assert_eq!(
2061                lazy.artifact_cache_charge_upper_bound(),
2062                upper_bound,
2063                "the charged upper bound is fixed at build time, not measured live"
2064            );
2065            assert!(
2066                lazy.mask_cache_retained_bytes() <= upper_bound,
2067                "live retained bytes must never exceed what was charged"
2068            );
2069        }
2070    }
2071
2072    #[test]
2073    fn write_mask_into_matches_the_reference_mask_bit_for_bit() {
2074        let artifact = boolean_artifact();
2075        let s = artifact.start();
2076        let reference = artifact.allowed_mask(s).expect("reference mask");
2077        let mut out = vec![0u32; artifact.words_per_row()];
2078        artifact
2079            .write_mask_into(&[s], &mut out, None)
2080            .expect("fill");
2081        assert_eq!(out.as_slice(), reference.as_words());
2082    }
2083
2084    #[test]
2085    fn allowed_ids_matches_a_per_id_mask_scan_over_every_state() {
2086        let artifact = boolean_artifact();
2087        for raw in 0..artifact.state_count() {
2088            let s = StateId(raw as u32);
2089            let mask = artifact.allowed_mask(s).expect("mask");
2090            let reference: Vec<TokenId> = (0..mask.vocab_size())
2091                .map(|id| TokenId(id as u32))
2092                .filter(|&id| mask.get(id))
2093                .collect();
2094            assert_eq!(artifact.allowed_ids(s).expect("allowed_ids"), reference);
2095        }
2096    }
2097
2098    #[test]
2099    fn write_mask_le_bytes_into_matches_write_mask_into_byte_for_byte() {
2100        let artifact = boolean_artifact();
2101        let s = artifact.start();
2102        let mut words = vec![0u32; artifact.words_per_row()];
2103        artifact.write_mask_into(&[s], &mut words, None).unwrap();
2104        let mut bytes = vec![0u8; artifact.words_per_row() * 4];
2105        artifact
2106            .write_mask_le_bytes_into(&[s], &mut bytes, None)
2107            .unwrap();
2108        let expected: Vec<u8> = words.iter().flat_map(|w| w.to_le_bytes()).collect();
2109        assert_eq!(bytes, expected);
2110    }
2111
2112    #[test]
2113    fn write_mask_le_bytes_into_clears_a_dirty_reused_buffer() {
2114        let artifact = boolean_artifact();
2115        let s = artifact.start();
2116        let mut fresh = vec![0u8; artifact.words_per_row() * 4];
2117        artifact
2118            .write_mask_le_bytes_into(&[s], &mut fresh, None)
2119            .unwrap();
2120        let mut dirty = vec![0xFFu8; artifact.words_per_row() * 4];
2121        artifact
2122            .write_mask_le_bytes_into(&[s], &mut dirty, None)
2123            .unwrap();
2124        assert_eq!(dirty, fresh, "stale bits must be cleared, not OR-ed in");
2125    }
2126
2127    #[test]
2128    fn write_mask_le_bytes_into_rejects_a_mismatched_buffer_length() {
2129        let artifact = boolean_artifact();
2130        let s = artifact.start();
2131        let mut too_short = vec![0u8; artifact.words_per_row() * 4 - 1];
2132        let err = artifact
2133            .write_mask_le_bytes_into(&[s], &mut too_short, None)
2134            .expect_err("length mismatch is a structured error");
2135        assert_eq!(err.code, ErrorCode::ArtifactOutOfBounds);
2136    }
2137
2138    #[test]
2139    fn cached_mask_hit_returns_the_same_content_as_a_fresh_computation() {
2140        let artifact = boolean_artifact();
2141        let s = artifact.start();
2142        let fresh = artifact.allowed_mask(s).expect("uncached");
2143        let first = artifact.cached_mask(s).expect("miss, populates cache");
2144        let second = artifact.cached_mask(s).expect("hit, reads cache");
2145        assert_eq!(first.as_words(), fresh.as_words());
2146        assert_eq!(second.as_words(), fresh.as_words());
2147        assert!(
2148            Arc::ptr_eq(&first, &second),
2149            "a hit must reuse the cached Arc, not recompute"
2150        );
2151    }
2152
2153    #[test]
2154    fn cached_mask_does_not_confuse_two_different_states() {
2155        let artifact = boolean_artifact();
2156        let start = artifact.start();
2157        let dead = artifact.dead();
2158        let start_mask = artifact.cached_mask(start).expect("start");
2159        let dead_mask = artifact.cached_mask(dead).expect("dead");
2160        assert_ne!(
2161            start_mask.as_words(),
2162            dead_mask.as_words(),
2163            "a live start state must not share the dead state's (empty) mask"
2164        );
2165        assert_eq!(
2166            dead_mask.as_words(),
2167            artifact.allowed_mask(dead).unwrap().as_words()
2168        );
2169    }
2170
2171    #[test]
2172    fn cached_mask_is_isolated_per_artifact_not_shared_globally() {
2173        let a = boolean_artifact();
2174        let mut map: Map<Vec<u8>, Vec<u32>> = Map::default();
2175        map.insert(b"t".to_vec(), vec![0]);
2176        map.insert(b"true".to_vec(), vec![1]);
2177        let vocab_b = build_vocabulary(5, map).expect("vocab");
2178        let engine_b = corpus_cases()
2179            .expect("corpus")
2180            .into_iter()
2181            .find(|c| c.name == "suite-type-boolean")
2182            .expect("boolean case")
2183            .engine;
2184        let b = CompiledArtifact::new(
2185            engine_b,
2186            Arc::new(vocab_b),
2187            Provenance::reference(HashState::Hash([9; 32])),
2188        )
2189        .expect("artifact b");
2190
2191        let s = a.start();
2192        let mask_a = a.cached_mask(s).expect("a");
2193        let mask_b = b.cached_mask(s).expect("b");
2194        assert_eq!(mask_a.as_words().len(), a.words_per_row());
2195        assert_eq!(mask_b.as_words().len(), b.words_per_row());
2196        assert_ne!(
2197            a.vocab_size(),
2198            b.vocab_size(),
2199            "the two artifacts have different vocab sizes (EOS 9 vs EOS 5)"
2200        );
2201        assert_eq!(mask_a.as_words(), a.allowed_mask(s).unwrap().as_words());
2202        assert_eq!(mask_b.as_words(), b.allowed_mask(s).unwrap().as_words());
2203    }
2204
2205    #[test]
2206    fn cached_mask_agrees_with_uncached_mask_at_vocab_boundary_sizes() {
2207        for vocab_size in [0usize, 1, 31, 32, 33, 64, 65] {
2208            let mut map: Map<Vec<u8>, Vec<u32>> = Map::default();
2209            for i in 0..vocab_size {
2210                map.insert(vec![b'a', i as u8], vec![i as u32]);
2211            }
2212            let eos = vocab_size as u32;
2213            let vocab = build_vocabulary(eos, map).expect("vocab");
2214            let engine = corpus_cases()
2215                .expect("corpus")
2216                .into_iter()
2217                .find(|c| c.name == "suite-type-boolean")
2218                .expect("boolean case")
2219                .engine;
2220            let artifact = CompiledArtifact::new(
2221                engine,
2222                Arc::new(vocab),
2223                Provenance::reference(HashState::Hash([3; 32])),
2224            )
2225            .expect("artifact");
2226            let s = artifact.start();
2227            let cached = artifact.cached_mask(s).expect("cached");
2228            let uncached = artifact.allowed_mask(s).expect("uncached");
2229            assert_eq!(
2230                cached.as_words(),
2231                uncached.as_words(),
2232                "vocab_size={vocab_size}"
2233            );
2234        }
2235    }
2236
2237    #[test]
2238    fn write_mask_into_rejects_a_mismatched_buffer_length() {
2239        let artifact = boolean_artifact();
2240        let s = artifact.start();
2241        let mut too_short = vec![0u32; artifact.words_per_row().saturating_sub(1)];
2242        let err = artifact
2243            .write_mask_into(&[s], &mut too_short, None)
2244            .expect_err("length mismatch is a structured error");
2245        assert_eq!(err.code, ErrorCode::ArtifactOutOfBounds);
2246    }
2247
2248    #[test]
2249    fn write_mask_into_rejects_an_overflowing_row_index() {
2250        let artifact = boolean_artifact();
2251        let s = artifact.start();
2252        let mut out = vec![0u32; artifact.words_per_row()];
2253        let err = artifact
2254            .write_mask_into(&[s], &mut out, Some(&[usize::MAX]))
2255            .expect_err("a row index of usize::MAX must not overflow");
2256        assert_eq!(err.code, ErrorCode::ArtifactOutOfBounds);
2257    }
2258
2259    #[test]
2260    fn advance_tracks_consume_token_and_leaves_state_on_rejection() {
2261        let artifact = Arc::new(boolean_artifact());
2262        let mut m = Matcher::new(artifact);
2263        m.advance(TokenId(1)).expect("`true` is legal from start");
2264        assert!(m.is_accepting());
2265        assert!(!m.is_dead());
2266        assert_eq!(m.bytes_seen(), 4);
2267        let before = m.state();
2268        assert!(m.advance(TokenId(2)).is_err());
2269        assert_eq!(m.state(), before);
2270    }
2271
2272    #[test]
2273    fn empty_language_matcher_is_dead_and_rejects_every_token_forever() {
2274        let engine = build_from_regex("[a&&b]").expect("empty-language engine");
2275        let mut map: Map<Vec<u8>, Vec<u32>> = Map::default();
2276        map.insert(b"a".to_vec(), vec![0]);
2277        map.insert(b"b".to_vec(), vec![1]);
2278        let vocab = build_vocabulary(5, map).expect("vocab");
2279        let artifact = Arc::new(
2280            CompiledArtifact::new(
2281                engine,
2282                Arc::new(vocab),
2283                Provenance::reference(HashState::Hash([7; 32])),
2284            )
2285            .expect("artifact"),
2286        );
2287        assert!(artifact.start_is_dead());
2288        let mut m = Matcher::new(artifact);
2289        assert!(m.is_dead());
2290        for id in [0u32, 1, 2, 99] {
2291            assert!(m.advance(TokenId(id)).is_err());
2292            assert!(m.is_dead(), "is_dead must stay true after every attempt");
2293        }
2294    }
2295
2296    #[test]
2297    fn unknown_token_id_is_illegal() {
2298        let artifact = Arc::new(boolean_artifact());
2299        let mut m = Matcher::new(artifact);
2300        match m.advance(TokenId(999)) {
2301            Err(MatcherError::IllegalToken { token }) => assert_eq!(token, TokenId(999)),
2302            other => panic!("expected IllegalToken, got {other:?}"),
2303        }
2304    }
2305
2306    #[test]
2307    fn artifact_rejects_incomplete_provenance() {
2308        let engine = corpus_cases()
2309            .expect("corpus")
2310            .into_iter()
2311            .find(|c| c.name == "suite-type-null")
2312            .expect("null case")
2313            .engine;
2314        let vocab = build_vocabulary(1, Map::default()).expect("vocab");
2315        let err = CompiledArtifact::new(
2316            engine,
2317            Arc::new(vocab),
2318            Provenance::reference(HashState::Unhashed),
2319        )
2320        .expect_err("unhashed reference provenance is incomplete");
2321        assert_eq!(err.code, ErrorCode::ProvenanceIncomplete);
2322    }
2323
2324    #[test]
2325    fn token_bytes_returns_the_stored_bytes_for_multi_id_and_repeated_sequences() {
2326        let artifact = boolean_artifact();
2327        assert_eq!(artifact.token_bytes(TokenId(0)), Some(b"t".as_slice()));
2328        assert_eq!(artifact.token_bytes(TokenId(1)), Some(b"true".as_slice()));
2329        assert_eq!(artifact.token_bytes(TokenId(2)), Some(b"z".as_slice()));
2330        assert_eq!(artifact.token_bytes(TokenId(3)), Some(b"z".as_slice()));
2331        assert_eq!(
2332            artifact.token_bytes(TokenId(2)),
2333            artifact.token_bytes(TokenId(3)),
2334            "two ids sharing one byte sequence must return the same bytes"
2335        );
2336        assert_eq!(artifact.token_bytes(TokenId(42)), None);
2337    }
2338
2339    #[test]
2340    fn validate_token_bytes_accepts_a_clean_map_and_rejects_a_conflicting_id() {
2341        let mut ok: Map<Vec<u8>, Vec<u32>> = Map::default();
2342        ok.insert(b"a".to_vec(), vec![0]);
2343        ok.insert(b"bb".to_vec(), vec![1, 2]);
2344        let clean = build_vocabulary(9, ok).expect("vocab");
2345        validate_token_bytes(&clean).expect("a clean map has no id conflict");
2346
2347        let mut bad: Map<Vec<u8>, Vec<u32>> = Map::default();
2348        bad.insert(b"a".to_vec(), vec![5]);
2349        bad.insert(b"b".to_vec(), vec![5]);
2350        let conflicting = build_vocabulary(9, bad).expect("vocab");
2351        let err = validate_token_bytes(&conflicting)
2352            .expect_err("one id under two byte sequences is ambiguous");
2353        assert_eq!(err.code, ErrorCode::MalformedTokenizer);
2354    }
2355
2356    #[test]
2357    fn artifact_rejects_one_id_under_two_byte_sequences() {
2358        let engine = corpus_cases()
2359            .expect("corpus")
2360            .into_iter()
2361            .find(|c| c.name == "suite-type-boolean")
2362            .expect("boolean case")
2363            .engine;
2364        let mut map: Map<Vec<u8>, Vec<u32>> = Map::default();
2365        map.insert(b"a".to_vec(), vec![5]);
2366        map.insert(b"b".to_vec(), vec![5]);
2367        let vocab = build_vocabulary(9, map).expect("vocab");
2368        let err = CompiledArtifact::new(
2369            engine,
2370            Arc::new(vocab),
2371            Provenance::reference(HashState::Hash([7; 32])),
2372        )
2373        .expect_err("one id under two byte sequences is ambiguous");
2374        assert_eq!(err.code, ErrorCode::MalformedTokenizer);
2375    }
2376
2377    fn boolean_engine() -> crate::automaton::RefEngine {
2378        corpus_cases()
2379            .expect("corpus")
2380            .into_iter()
2381            .find(|c| c.name == "suite-type-boolean")
2382            .expect("boolean case")
2383            .engine
2384    }
2385
2386    #[test]
2387    fn a_forged_huge_token_id_is_rejected_before_allocating_a_mask() {
2388        let mut map: Map<Vec<u8>, Vec<u32>> = Map::default();
2389        map.insert(b"t".to_vec(), vec![u32::MAX - 1]);
2390        let vocab = build_vocabulary(9, map).expect("vocab");
2391        let err = CompiledArtifact::new(
2392            boolean_engine(),
2393            Arc::new(vocab),
2394            Provenance::reference(HashState::Hash([7; 32])),
2395        )
2396        .expect_err("an oversized mask width must be rejected");
2397        assert_eq!(err.code, ErrorCode::InternalLimitExceeded);
2398    }
2399
2400    #[test]
2401    fn forged_out_of_range_state_ids_do_not_grow_the_mask_cache() {
2402        let artifact = boolean_artifact();
2403        let n = artifact.state_count() as u32;
2404        for id in n..n + 1000 {
2405            let _ = artifact
2406                .cached_mask(StateId(id))
2407                .expect("out-of-range mask");
2408        }
2409        assert_eq!(
2410            artifact.mask_cache_filled_count(),
2411            0,
2412            "a forged out-of-range id must never be cached"
2413        );
2414        artifact
2415            .cached_mask(artifact.start())
2416            .expect("in-range mask");
2417        let cached = artifact.mask_cache_filled_count();
2418        assert!(
2419            (1..=artifact.state_count()).contains(&cached),
2420            "only in-range states are cached, bounded by state_count"
2421        );
2422    }
2423
2424    #[test]
2425    fn the_mask_cache_respects_its_byte_budget() {
2426        let artifact = boolean_artifact();
2427        artifact.set_mask_cache_budget(0);
2428        artifact
2429            .cached_mask(artifact.start())
2430            .expect("mask still computed");
2431        assert_eq!(
2432            artifact.mask_cache_filled_count(),
2433            0,
2434            "a zero budget caches nothing, but the mask is still returned"
2435        );
2436        let one_entry =
2437            artifact.words_per_row() * std::mem::size_of::<u32>() + MASK_CACHE_ENTRY_OVERHEAD;
2438        artifact.set_mask_cache_budget(one_entry);
2439        artifact.cached_mask(artifact.start()).expect("first fits");
2440        artifact
2441            .cached_mask(artifact.dead())
2442            .expect("second does not fit");
2443        assert!(
2444            artifact.mask_cache_retained_bytes() <= one_entry,
2445            "cached bytes never exceed the budget"
2446        );
2447        assert!(
2448            artifact.mask_cache_filled_count() <= 1,
2449            "only what fits in the budget is retained"
2450        );
2451    }
2452
2453    #[test]
2454    fn concurrent_lazy_queries_across_many_threads_agree_with_naive() {
2455        let lazy = boolean_lazy_artifact();
2456        let naive = boolean_artifact();
2457        let n = lazy.engine.state_count();
2458        std::thread::scope(|scope| {
2459            for t in 0..32 {
2460                let lazy = &lazy;
2461                let naive = &naive;
2462                scope.spawn(move || {
2463                    for raw in 0..n {
2464                        let s = StateId(((raw + t) % n) as u32);
2465                        assert_eq!(
2466                            lazy.cached_mask(s).unwrap().as_words(),
2467                            naive.cached_mask(s).unwrap().as_words(),
2468                            "thread {t} state {raw} diverged from naive"
2469                        );
2470                    }
2471                });
2472            }
2473        });
2474    }
2475
2476    #[test]
2477    fn concurrent_race_on_one_state_charges_exactly_one_entry() {
2478        let artifact = boolean_lazy_artifact();
2479        let entry_bytes = artifact.words_per_row() * 4 + MASK_CACHE_ENTRY_OVERHEAD;
2480        std::thread::scope(|scope| {
2481            for _ in 0..32 {
2482                let artifact = &artifact;
2483                scope.spawn(move || {
2484                    artifact.cached_mask(artifact.start()).unwrap();
2485                });
2486            }
2487        });
2488        assert_eq!(artifact.mask_cache_retained_bytes(), entry_bytes);
2489    }
2490
2491    #[test]
2492    fn concurrent_race_across_m_states_charges_exactly_m_entries() {
2493        let artifact = boolean_lazy_artifact();
2494        let m = artifact.engine.state_count();
2495        let entry_bytes = artifact.words_per_row() * 4 + MASK_CACHE_ENTRY_OVERHEAD;
2496        std::thread::scope(|scope| {
2497            for _ in 0..32 {
2498                let artifact = &artifact;
2499                scope.spawn(move || {
2500                    for raw in 0..m {
2501                        artifact.cached_mask(StateId(raw as u32)).unwrap();
2502                    }
2503                });
2504            }
2505        });
2506        assert_eq!(artifact.mask_cache_retained_bytes(), m * entry_bytes);
2507    }
2508
2509    #[test]
2510    fn a_failed_compute_refunds_exactly_one_reservation() {
2511        let rows = LazyRows::try_new(4, 1024).expect("rows");
2512        let err = || {
2513            Err(CompileError::new(
2514                ErrorCode::Unsupported,
2515                Stage::L4Bind,
2516                "injected",
2517            ))
2518        };
2519        let result = rows.get_or_compute(0, 100, err);
2520        assert!(result.is_err());
2521        assert_eq!(
2522            rows.retained_bytes(),
2523            0,
2524            "a failed compute must refund its reservation"
2525        );
2526    }
2527
2528    #[test]
2529    fn packed_artifact_allocates_no_lazy_rows() {
2530        let mut map: Map<Vec<u8>, Vec<u32>> = Map::default();
2531        map.insert(b"t".to_vec(), vec![0]);
2532        let vocab = Arc::new(build_vocabulary(1, map).expect("vocab"));
2533        let trie = VocabTrie::build_byte(&vocab).expect("trie");
2534        let packed = CompiledArtifact::new_trie(
2535            boolean_engine(),
2536            vocab,
2537            prov_bytes([9; 32]),
2538            BindMode::TrieJointBytePacked,
2539            &trie,
2540        )
2541        .expect("packed");
2542        assert!(packed.mask_cache_is_disabled());
2543        let lazy_share = packed
2544            .artifact_cache_charge_upper_bound()
2545            .saturating_sub(packed.engine_retained_bytes())
2546            .saturating_sub(packed.packed_retained_bytes())
2547            .saturating_sub(packed.index_retained_bytes());
2548        assert_eq!(
2549            lazy_share, 0,
2550            "packed's charge must carry zero lazy-row bytes"
2551        );
2552    }
2553
2554    #[test]
2555    fn oversized_lazy_row_array_fails_structurally_before_allocating() {
2556        let huge_state_count = MAX_LAZY_ROWS_BYTES / std::mem::size_of::<MaskSlot>() + 1;
2557        let err =
2558            LazyRows::try_new(huge_state_count, 1024).expect_err("must reject before allocating");
2559        assert_eq!(err.code, ErrorCode::InternalLimitExceeded);
2560    }
2561
2562    #[test]
2563    fn mask_slot_stays_compact_via_arc_compile_error() {
2564        let slot = std::mem::size_of::<MaskSlot>();
2565        let inline = std::mem::size_of::<OnceLock<Result<Arc<Bitmask>, CompileError>>>();
2566        assert!(
2567            slot <= 24,
2568            "MaskSlot grew past 24 bytes: {slot} - CompileError may be inlined again"
2569        );
2570        assert!(
2571            slot < inline,
2572            "Arc<CompileError> ({slot}B) must stay smaller than inlining it ({inline}B)"
2573        );
2574    }
2575
2576    #[test]
2577    fn packed_forced_lazy_and_adaptive_masks_all_equal_naive() {
2578        let mut map: Map<Vec<u8>, Vec<u32>> = Map::default();
2579        map.insert(b"t".to_vec(), vec![0]);
2580        map.insert(b"true".to_vec(), vec![1]);
2581        map.insert(b"z".to_vec(), vec![2, 3]);
2582        let vocab = Arc::new(build_vocabulary(9, map).expect("vocab"));
2583        let handle = crate::index::VocabularyHandle::new(vocab.clone()).expect("handle");
2584        let bound = crate::index::TrieCache::new().bind(&handle).expect("bound");
2585        let naive = boolean_artifact();
2586        let n = boolean_engine().state_count();
2587
2588        let packed = CompiledArtifact::new_from_bound_trie(
2589            &handle,
2590            boolean_engine(),
2591            prov_bytes([8; 32]),
2592            BindMode::TrieJointBytePacked,
2593            &bound,
2594        )
2595        .expect("packed");
2596        let lazy = CompiledArtifact::new_from_bound_trie(
2597            &handle,
2598            boolean_engine(),
2599            prov_bytes([8; 32]),
2600            BindMode::TrieJointByteLazy,
2601            &bound,
2602        )
2603        .expect("lazy");
2604        let adaptive_mode = CompiledArtifact::recommended_bind_mode(
2605            n,
2606            handle.mask_vocab_size(),
2607            &bound,
2608            BindMode::TrieJointBytePacked,
2609        );
2610        let adaptive = CompiledArtifact::new_from_bound_trie(
2611            &handle,
2612            boolean_engine(),
2613            prov_bytes([8; 32]),
2614            adaptive_mode,
2615            &bound,
2616        )
2617        .expect("adaptive");
2618
2619        for raw in 0..n {
2620            let s = StateId(raw as u32);
2621            let expected = naive.allowed_mask(s).unwrap();
2622            assert_eq!(
2623                packed.allowed_mask(s).unwrap().as_words(),
2624                expected.as_words()
2625            );
2626            assert_eq!(
2627                lazy.allowed_mask(s).unwrap().as_words(),
2628                expected.as_words()
2629            );
2630            assert_eq!(
2631                adaptive.allowed_mask(s).unwrap().as_words(),
2632                expected.as_words()
2633            );
2634        }
2635    }
2636
2637    #[test]
2638    fn a_shared_canonical_vocabulary_is_not_rebuilt_per_artifact() {
2639        let mut map: Map<Vec<u8>, Vec<u32>> = Map::default();
2640        map.insert(b"t".to_vec(), vec![0]);
2641        let vocab = Arc::new(build_vocabulary(9, map).expect("vocab"));
2642        let handle = crate::index::VocabularyHandle::new(vocab).expect("handle");
2643        let before = Arc::strong_count(handle.canonical());
2644        let a = CompiledArtifact::new_from_handle(
2645            &handle,
2646            boolean_engine(),
2647            Provenance::reference(HashState::Hash([7; 32])),
2648        )
2649        .expect("artifact a");
2650        let b = CompiledArtifact::new_from_handle(
2651            &handle,
2652            boolean_engine(),
2653            Provenance::reference(HashState::Hash([8; 32])),
2654        )
2655        .expect("artifact b");
2656        assert_eq!(
2657            Arc::strong_count(handle.canonical()),
2658            before + 2,
2659            "each artifact shares the canonical Arc, never rebuilds it"
2660        );
2661        assert_eq!(a.vocab_size(), b.vocab_size());
2662    }
2663
2664    #[test]
2665    fn new_trie_rejects_a_trie_built_from_a_different_vocabulary() {
2666        use crate::index::{BindMode, VocabTrie};
2667        let mut map_a: Map<Vec<u8>, Vec<u32>> = Map::default();
2668        map_a.insert(b"a".to_vec(), vec![0]);
2669        let vocab_a = build_vocabulary(9, map_a).expect("vocab a");
2670
2671        let mut map_b: Map<Vec<u8>, Vec<u32>> = Map::default();
2672        map_b.insert(b"b".to_vec(), vec![0]);
2673        let vocab_b = Arc::new(build_vocabulary(9, map_b).expect("vocab b"));
2674
2675        let trie_a = VocabTrie::build_byte(&vocab_a).expect("trie a");
2676        let err = CompiledArtifact::new_trie(
2677            boolean_engine(),
2678            vocab_b,
2679            Provenance::reference(HashState::Hash([7; 32])),
2680            BindMode::TrieJointByte,
2681            &trie_a,
2682        )
2683        .expect_err("a trie from a different vocabulary must be rejected");
2684        assert_eq!(err.code, ErrorCode::ArtifactMismatch);
2685    }
2686
2687    #[test]
2688    fn new_trie_prevalidated_still_rejects_a_genuine_mismatch() {
2689        use crate::index::{BindMode, VocabFingerprint, VocabTrie};
2690        let mut map_a: Map<Vec<u8>, Vec<u32>> = Map::default();
2691        map_a.insert(b"a".to_vec(), vec![0]);
2692        let vocab_a = build_vocabulary(9, map_a).expect("vocab a");
2693        let fingerprint_a = VocabFingerprint::of(&vocab_a).unwrap();
2694
2695        let mut map_b: Map<Vec<u8>, Vec<u32>> = Map::default();
2696        map_b.insert(b"b".to_vec(), vec![0]);
2697        let vocab_b = Arc::new(build_vocabulary(9, map_b).expect("vocab b"));
2698        let trie_b = VocabTrie::build_byte(&vocab_b).expect("trie b");
2699        let canonical_b =
2700            Arc::new(CanonicalVocabulary::from_vocabulary(&vocab_b).expect("canonical b"));
2701
2702        let err = CompiledArtifact::new_trie_prevalidated(
2703            boolean_engine(),
2704            canonical_b,
2705            Provenance::reference(HashState::Hash([7; 32])),
2706            BindMode::TrieJointByte,
2707            &trie_b,
2708            fingerprint_a, // wrong fingerprint for this trie
2709        )
2710        .expect_err(
2711            "a trie whose OWN fingerprint disagrees with the supplied one must be rejected",
2712        );
2713        assert_eq!(err.code, ErrorCode::ArtifactMismatch);
2714    }
2715
2716    #[test]
2717    fn new_trie_prevalidated_agrees_with_new_trie_on_a_genuine_pairing() {
2718        use crate::index::{BindMode, VocabFingerprint, VocabTrie};
2719        let mut map: Map<Vec<u8>, Vec<u32>> = Map::default();
2720        map.insert(b"true".to_vec(), vec![0]);
2721        map.insert(b"false".to_vec(), vec![1]);
2722        let vocab = Arc::new(build_vocabulary(9, map).expect("vocab"));
2723        let trie = VocabTrie::build_byte(&vocab).expect("trie");
2724
2725        let a = CompiledArtifact::new_trie(
2726            boolean_engine(),
2727            vocab.clone(),
2728            Provenance::reference(HashState::Hash([3; 32])),
2729            BindMode::TrieJointByte,
2730            &trie,
2731        )
2732        .expect("generic new_trie");
2733
2734        let fingerprint = VocabFingerprint::of(&vocab).unwrap();
2735        let canonical = Arc::new(CanonicalVocabulary::from_vocabulary(&vocab).expect("canonical"));
2736        let b = CompiledArtifact::new_trie_prevalidated(
2737            boolean_engine(),
2738            canonical,
2739            Provenance::reference(HashState::Hash([3; 32])),
2740            BindMode::TrieJointByte,
2741            &trie,
2742            fingerprint,
2743        )
2744        .expect("prevalidated new_trie");
2745
2746        let s = a.start();
2747        assert_eq!(
2748            a.allowed_mask(s).unwrap().as_words(),
2749            b.allowed_mask(s).unwrap().as_words(),
2750            "the fast and generic constructors must produce byte-identical masks"
2751        );
2752    }
2753
2754    fn handle_of(pairs: &[(&[u8], u32)], eos: u32) -> crate::index::VocabularyHandle {
2755        let mut map: Map<Vec<u8>, Vec<u32>> = Map::default();
2756        for &(bytes, id) in pairs {
2757            map.insert(bytes.to_vec(), vec![id]);
2758        }
2759        let vocab = Arc::new(build_vocabulary(eos, map).expect("vocab"));
2760        crate::index::VocabularyHandle::new(vocab).expect("handle")
2761    }
2762
2763    #[test]
2764    fn new_from_bound_trie_rejects_a_trie_bound_to_a_different_handle() {
2765        use crate::index::{BindMode, TrieCache};
2766        let cache = TrieCache::new();
2767        let a = handle_of(&[(b"a", 0)], 9);
2768        let b = handle_of(&[(b"b", 0)], 9);
2769        let bound_a = cache.bind(&a).expect("bind a");
2770        let err = CompiledArtifact::new_from_bound_trie(
2771            &b,
2772            boolean_engine(),
2773            Provenance::reference(HashState::Hash([7; 32])),
2774            BindMode::TrieJointByte,
2775            &bound_a,
2776        )
2777        .expect_err("a trie bound to handle a must not bind under handle b");
2778        assert_eq!(err.code, ErrorCode::ArtifactMismatch);
2779    }
2780
2781    #[test]
2782    fn new_from_bound_trie_rejects_same_bytes_different_eos() {
2783        use crate::index::{BindMode, TrieCache};
2784        let cache = TrieCache::new();
2785        let a = handle_of(&[(b"a", 0), (b"b", 1)], 9);
2786        let b = handle_of(&[(b"a", 0), (b"b", 1)], 10);
2787        let bound_a = cache.bind(&a).expect("bind a");
2788        let err = CompiledArtifact::new_from_bound_trie(
2789            &b,
2790            boolean_engine(),
2791            Provenance::reference(HashState::Hash([7; 32])),
2792            BindMode::TrieJointByte,
2793            &bound_a,
2794        )
2795        .expect_err("identical token bytes under a different EOS must not bind");
2796        assert_eq!(err.code, ErrorCode::ArtifactMismatch);
2797    }
2798
2799    #[test]
2800    fn new_from_bound_trie_rejects_same_ids_different_bytes() {
2801        use crate::index::{BindMode, TrieCache};
2802        let cache = TrieCache::new();
2803        let a = handle_of(&[(b"a", 0), (b"b", 1)], 9);
2804        let b = handle_of(&[(b"x", 0), (b"y", 1)], 9);
2805        let bound_a = cache.bind(&a).expect("bind a");
2806        let err = CompiledArtifact::new_from_bound_trie(
2807            &b,
2808            boolean_engine(),
2809            Provenance::reference(HashState::Hash([7; 32])),
2810            BindMode::TrieJointByte,
2811            &bound_a,
2812        )
2813        .expect_err("same token ids under different byte payloads must not bind");
2814        assert_eq!(err.code, ErrorCode::ArtifactMismatch);
2815    }
2816
2817    #[test]
2818    fn new_from_bound_trie_accepts_a_genuinely_matching_pair() {
2819        use crate::index::{BindMode, TrieCache};
2820        let cache = TrieCache::new();
2821        let handle = handle_of(&[(b"true", 0), (b"false", 1)], 9);
2822        let bound = cache.bind(&handle).expect("bind");
2823        let artifact = CompiledArtifact::new_from_bound_trie(
2824            &handle,
2825            boolean_engine(),
2826            Provenance::reference(HashState::Hash([7; 32])),
2827            BindMode::TrieJointByte,
2828            &bound,
2829        )
2830        .expect("a genuinely matching handle and bound trie must bind");
2831        assert!(artifact.allowed_mask(artifact.start()).is_ok());
2832    }
2833
2834    #[test]
2835    fn cache_bind_reuses_the_same_arc_across_two_handles_of_equal_content() {
2836        use crate::index::TrieCache;
2837        let cache = TrieCache::new();
2838        let a = handle_of(&[(b"a", 0), (b"b", 1)], 9);
2839        let b = handle_of(&[(b"a", 0), (b"b", 1)], 9);
2840        let bound_a = cache.bind(&a).expect("bind a");
2841        let bound_b = cache.bind(&b).expect("bind b");
2842        assert_eq!(
2843            bound_a.fingerprint(),
2844            bound_b.fingerprint(),
2845            "equal content under two distinct handles must share one cache entry"
2846        );
2847    }
2848}