Skip to main content

maskforge_core/index/
mod.rs

1//! Binds vocabularies to automata for allowed-token mask generation.
2
3mod artifact;
4mod build;
5mod cache;
6mod slice;
7mod trie;
8
9pub use artifact::CompiledIndex;
10pub use build::{build_delta, DeltaRow};
11pub(crate) use build::{
12    fill_packed_rows, packed_scratch_fixed_bytes, packed_stack_bytes, walk_state_mask_into,
13};
14pub use cache::{BoundByteTrie, TrieCache, VocabFingerprint, VocabularyHandle};
15pub(crate) use slice::{BoundedMask, TokenSliceCatalog};
16pub use trie::VocabTrie;
17
18use crate::error::CompileError;
19
20#[must_use]
21pub fn packed_ingestion_peak_bytes(
22    token_bytes_len: usize,
23    byte_offsets_len: usize,
24    token_ids_len: usize,
25    id_offsets_len: usize,
26) -> Option<usize> {
27    let decoded = token_bytes_len
28        .checked_add(byte_offsets_len)?
29        .checked_add(token_ids_len)?
30        .checked_add(id_offsets_len)?;
31    let token_count = (byte_offsets_len / 4).saturating_sub(1);
32    let ids_count = token_ids_len / 4;
33    let core_peak = crate::vocab::projected_peak_bytes(token_bytes_len, ids_count, token_count);
34    decoded.checked_add(core_peak)
35}
36
37#[must_use]
38pub fn direct_csr_ingestion_peak_bytes(
39    total_bytes: usize,
40    total_ids: usize,
41    token_count: usize,
42) -> Option<usize> {
43    let offsets_len = token_count.checked_add(1)?.checked_mul(4)?;
44    let buffers = total_bytes
45        .checked_add(total_ids.checked_mul(4)?)?
46        .checked_add(offsets_len)?
47        .checked_add(offsets_len)?;
48    let core_peak = crate::vocab::projected_peak_bytes(total_bytes, total_ids, token_count);
49    buffers.checked_add(core_peak)
50}
51
52#[must_use]
53pub fn dense_id_ordered_ingestion_peak_bytes(
54    total_bytes: usize,
55    present_count: usize,
56    token_count: usize,
57) -> Option<usize> {
58    let offsets_len = token_count.checked_add(1)?.checked_mul(4)?;
59    let present_len = token_count;
60    let buffers = total_bytes
61        .checked_add(offsets_len)?
62        .checked_add(present_len)?;
63    let core_peak = crate::vocab::projected_peak_bytes(total_bytes, present_count, present_count);
64    buffers.checked_add(core_peak)
65}
66
67pub use crate::mem_gate::PermitGuard as VocabBuildPermit;
68
69pub fn acquire_vocab_build_permit(
70    estimated_bytes: usize,
71) -> Result<VocabBuildPermit<'static>, CompileError> {
72    crate::mem_gate::acquire_vocab_build_bytes(estimated_bytes)
73}
74
75#[must_use]
76pub fn vocab_build_gate_stats() -> (usize, usize, usize, u64) {
77    crate::mem_gate::vocab_build_gate_stats()
78}
79
80#[must_use]
81pub fn vocab_serving_gate_stats() -> (usize, usize, usize, u64) {
82    crate::mem_gate::vocab_serving_gate_stats()
83}
84
85#[must_use]
86pub fn automaton_build_gate_stats() -> (usize, usize, usize, u64) {
87    crate::mem_gate::automaton_build_gate_stats()
88}
89
90pub use crate::mem_gate::GateConfigError;
91
92pub fn configure_vocab_build_gate(
93    max_concurrent: usize,
94    byte_budget: usize,
95) -> Result<(), GateConfigError> {
96    crate::mem_gate::configure_vocab_build_gate(max_concurrent, byte_budget)
97}
98
99pub fn configure_vocab_serving_gate(
100    max_concurrent: usize,
101    byte_budget: usize,
102) -> Result<(), GateConfigError> {
103    crate::mem_gate::configure_vocab_serving_gate(max_concurrent, byte_budget)
104}
105
106pub fn configure_automaton_build_gate(
107    max_concurrent: usize,
108    byte_budget: usize,
109) -> Result<(), GateConfigError> {
110    crate::mem_gate::configure_automaton_build_gate(max_concurrent, byte_budget)
111}
112
113#[non_exhaustive]
114#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
115/// Selects the strategy used to build allowed-token masks.
116pub enum BindMode {
117    /// Rewalks every token from each automaton state.
118    Naive,
119    /// Eagerly builds state rows with a shared byte trie.
120    TrieJointByte,
121    /// Builds byte-trie rows lazily on first use.
122    TrieJointByteLazy,
123    /// Precomputes packed mask rows for all states.
124    TrieJointBytePacked,
125    /// Uses a byte-class trie for per-compile binding.
126    TrieJointClass,
127}