Skip to main content

bitcannon/
lib.rs

1//! `bitcannon` — GPT-family pre-tokenization as a **bitstream program**, replacing the scalar FSMs.
2//!
3//! Follows *Interleaved Bitstream Execution for Multi-Pattern Regex Matching on GPUs*
4//! (MICRO'25, doi 10.1145/3725843.3756052). The paper's two ideas that carry over to a CPU:
5//!
6//!   1. **Bit-parallel regex.** Compile the grammar into character-class *bitstreams* (one bit per
7//!      input byte) plus boolean ops and carry-propagating adds. 64 input bytes are decided per
8//!      64-bit register op, branchlessly — the FSM's per-token unpredictable branch disappears.
9//!   2. **Interleaved execution.** Do NOT run one loop per bitstream instruction over the whole
10//!      input (that is the "sequential" baseline the paper beats: every intermediate stream is
11//!      materialised and re-read). Instead fuse *all* instructions into ONE block-wise loop, so an
12//!      intermediate stream lives in a register for the ~100 ops it is needed and dies there. The
13//!      only thing that reaches memory is the final `starts` bitmap (n/8 bytes).
14//!
15//! The paper's third contribution (dependency-aware thread-data mapping) is a GPU concern — it
16//! resolves cross-block dependencies by recomputing them on other SMs. On one core the blocks are
17//! visited in order, so those dependencies are carried in a handful of scalar registers and, for
18//! the one genuinely backward-in-time rule, patched into the already-written bitmap.
19//!
20//! We do not re-derive character classes: [`classify`] already emits one `Atom` tag per byte. The
21//! builder folds those 16 atoms into a grammar-specific **dense 3-bit code** and extracts 3
22//! bit-planes of it; every class stream is then a 2–3 op boolean function of the planes.
23//!
24//! Grammars: [`bitcannon_deepseek`], [`bitcannon_byte_level`] (GPT-2), [`bitcannon_cl100k`]. All three
25//! byte-exact with the oniguruma oracle over a block-phase sweep — see `tests/parity.rs`.
26
27/// Declares a grammar's block-local class streams, and with them the two carried shifts every
28/// rule needs: `back` — what held one position earlier, carrying in the previous block's last
29/// bit — and `fwd`, one position later, carrying in the next block's first bit.
30///
31/// This is the whole cross-block story. Keeping the carry in the *stream* instead of passing it at
32/// every call site is what lets a rule read as plain stream algebra: "a run starts where this class
33/// holds and the position before it did not" is `x & lead & !bk.x`, exact at a block edge. Before
34/// this, each grammar rebuilt the edge byte's class into a `u16` mask (`code_bits`/`bits_at`) and
35/// threaded it through every shift by hand — so the boundary byte got classified twice, once by the
36/// builder and again scalar-side.
37macro_rules! streams {
38    ($(#[$m:meta])* $name:ident { $($f:ident),+ $(,)? }) => {
39        $(#[$m])*
40        #[derive(Clone, Copy, Default)]
41        pub(crate) struct $name { $(pub(crate) $f: u64,)+ }
42
43        impl $name {
44            /// Every stream shifted one position later, carrying in `p`'s last bit.
45            #[inline]
46            pub(crate) fn back(&self, p: &Self) -> Self {
47                Self { $($f: (self.$f << 1) | (p.$f >> 63),)+ }
48            }
49
50            /// Every stream shifted one position earlier, carrying in `n`'s first bit.
51            #[inline]
52            pub(crate) fn fwd(&self, n: &Self) -> Self {
53                Self { $($f: (self.$f >> 1) | (n.$f << 63),)+ }
54            }
55        }
56
57        impl crate::Streams for $name {
58            #[inline]
59            fn back(&self, p: &Self) -> Self { $name::back(self, p) }
60            #[inline]
61            fn fwd(&self, n: &Self) -> Self { $name::fwd(self, n) }
62            #[inline]
63            fn lead(&self) -> u64 { self.lead }
64        }
65    };
66}
67
68/// A grammar's class-stream set, as generated by [`streams!`]. Every grammar declares `lead` and
69/// `cont`, because the driver needs them for the char-boundary geometry.
70pub(crate) trait Streams: Copy + Default {
71    fn back(&self, p: &Self) -> Self;
72    fn fwd(&self, n: &Self) -> Self;
73    fn lead(&self) -> u64;
74}
75
76/// Everything a rule needs about one block, prepared once by [`blocks`] instead of re-derived by
77/// every grammar: the previous / current / next blocks, both carried shifts, and the geometry.
78pub(crate) struct Ctx<C> {
79    /// the previous block's streams (all-zero before the first)
80    pub pv: C,
81    pub cur: C,
82    /// the next block's streams (all-zero past the last)
83    pub nx: C,
84    /// `cur` one position later, carrying in `pv`'s last bit — "what held just before here"
85    pub bk: C,
86    /// `cur` one position earlier, carrying in `nx`'s first bit — "what holds just after here"
87    pub fw: C,
88    pub bi: usize,
89    pub base: usize,
90    pub len: usize,
91    /// bits inside this block (`!0` for a full one)
92    pub valid: u64,
93    pub last_blk: bool,
94    /// this position is the LAST byte of its char (the input's final byte counts)
95    pub lb: u64,
96    /// the input's final byte as a bit, else 0
97    pub eof: u64,
98}
99
100/// What a rule returns for its block: the token-start bitmap, a patch OR-ed into the *previous*
101/// block (a char straddling the edge lands its lead in the word already written), and a mask
102/// AND-ed with the FINISHED starts to fill the `flag` bitmap. Returning the mask rather than
103/// writing `flag` keeps `& lead` and the "position 0 always opens a token" rule in one place —
104/// three grammars were redoing both just to see the final `st`.
105pub(crate) struct Out {
106    pub st: u64,
107    pub patch: u64,
108    pub flag: u64,
109}
110
111/// The edge bits of a carried stream: `was` is the previous block's last, `will` the next block's
112/// first. Every grammar defined these as local closures.
113#[inline]
114pub(crate) fn was(v: u64) -> bool {
115    v >> 63 != 0
116}
117
118#[inline]
119pub(crate) fn will(v: u64) -> bool {
120    v & 1 != 0
121}
122
123/// Drive a grammar block by block with one block of lookahead.
124///
125/// This owns everything that is *not* grammar — block geometry, the fill-seed threading (inside
126/// `build`'s captures), the previous/current/next rotation, both carried shifts, `lb`, the eof bit,
127/// and the starts assembly, plus the fill seed. A grammar supplies its atom table, its plane decode
128/// and its rules; nothing else. `rule` also gets `starts` because one rule (deepseek's `anl`) retracts a
129/// start it committed in an earlier block.
130#[inline]
131pub(crate) fn blocks<C, S, B, R>(
132    ntext: usize,
133    starts: &mut [u64],
134    mut flag: Option<&mut [u64]>,
135    seed0: S,
136    mut build: B,
137    mut rule: R,
138) where
139    C: Streams,
140    B: FnMut(usize, usize, S) -> (C, S),
141    R: FnMut(&Ctx<C>, &mut [u64]) -> Out,
142{
143    // The fill seed — the class of the byte before a block, so one opening mid-char keeps
144    // inheriting its lead's. Threaded here rather than through a mutable capture in every grammar.
145    let mut seed = seed0;
146    let nblk = ntext.div_ceil(64);
147    let blen = |bi: usize| (ntext - bi * 64).min(64);
148    let mut pv = C::default();
149    let (mut cur, s) = build(0, blen(0), seed);
150    seed = s;
151
152    for bi in 0..nblk {
153        let base = bi * 64;
154        let len = blen(bi);
155        let last_blk = base + len == ntext;
156        // built BEFORE this block is processed, which is what makes `fw` a real carry
157        let nx = if last_blk {
158            C::default()
159        } else {
160            let (c, s) = build(base + 64, blen(bi + 1), seed);
161            seed = s;
162            c
163        };
164        let valid = if len == 64 { !0u64 } else { (1u64 << len) - 1 };
165        let eof = if last_blk { 1u64 << (len - 1) } else { 0 };
166        let bk = cur.back(&pv);
167        let fw = cur.fwd(&nx);
168        let ctx = Ctx {
169            pv,
170            cur,
171            nx,
172            bk,
173            fw,
174            bi,
175            base,
176            len,
177            valid,
178            last_blk,
179            lb: (fw.lead() & valid) | eof,
180            eof,
181        };
182        let o = rule(&ctx, starts);
183        let mut st = o.st & cur.lead();
184        if bi == 0 {
185            st |= 1; // position 0 always opens a token
186        }
187        starts[bi] = st;
188        if let Some(f) = flag.as_deref_mut() {
189            f[bi] = st & o.flag;
190        }
191        if bi > 0 {
192            starts[bi - 1] |= o.patch;
193        }
194        pv = cur;
195        cur = nx;
196    }
197}
198
199pub mod classes;
200pub mod classify;
201pub mod literal;
202pub mod models;
203pub mod regexes;
204mod simd;
205
206pub use models::cl100k::{bitcannon_cl100k, bitcannon_qwen};
207pub use models::deepseek::bitcannon_deepseek;
208pub use models::gpt2::bitcannon_byte_level;
209pub use models::kimi::bitcannon_kimi;
210pub use models::o200k::bitcannon_o200k;
211pub use models::tekken::bitcannon_tekken;
212
213/// A token span: byte offsets `[start, end)` into the input. `#[repr(C)]` so the output buffer has a
214/// stable `[start, end]` layout — the pipeline reuses it with zero conversion, and it can be
215/// reinterpreted as bytes / handed across the crate boundary.
216#[repr(C)]
217#[derive(Clone, Copy, PartialEq, Eq, Debug, Default, Hash, PartialOrd, Ord)]
218pub struct Span {
219    pub start: u32,
220    pub end: u32,
221}
222
223impl Span {
224    #[inline]
225    #[must_use]
226    pub const fn new(start: u32, end: u32) -> Self {
227        Self { start, end }
228    }
229
230    /// `[start, end)` as a `usize` range — for slicing the input text.
231    #[inline]
232    #[must_use]
233    pub fn range(self) -> core::ops::Range<usize> {
234        self.start as usize..self.end as usize
235    }
236}
237
238impl PartialEq<(u32, u32)> for Span {
239    fn eq(&self, other: &(u32, u32)) -> bool {
240        self.start == other.0 && self.end == other.1
241    }
242}
243
244/// Whether this target builds its bitstreams with SIMD. Without it the builder is the portable
245/// byte-at-a-time reference, which is slower than the FSM this replaces -- so a caller should keep
246/// its FSM path rather than route here. On x86 the kernel needs SSSE3, so this is a runtime check.
247#[must_use]
248pub fn fast_builder() -> bool {
249    #[cfg(target_arch = "aarch64")]
250    {
251        true
252    }
253    #[cfg(target_arch = "x86_64")]
254    {
255        has_ssse3()
256    }
257    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
258    {
259        false
260    }
261}
262
263#[cfg(target_arch = "x86_64")]
264pub(crate) fn has_ssse3() -> bool {
265    use std::sync::atomic::{AtomicU8, Ordering};
266    static CACHED: AtomicU8 = AtomicU8::new(0);
267    match CACHED.load(Ordering::Relaxed) {
268        0 => {
269            let yes = std::arch::is_x86_feature_detected!("ssse3");
270            CACHED.store(1 + u8::from(yes), Ordering::Relaxed);
271            yes
272        }
273        n => n == 2,
274    }
275}
276
277pub(crate) const CONT: u8 = 15; // Atom::Cont
278pub(crate) const CODE_CONT: u8 = 7; // every grammar's dense code for a continuation byte
279
280/// What a grammar wants in the text-derived `Blk::aux` stream. Tag-derived classes go through the
281/// LUT; these three need the raw bytes, so they are the one thing the builder reads `text` for.
282pub(crate) const AUX_NONE: u8 = 0;
283pub(crate) const AUX_CJK: u8 = 1; // deepseek Split-2: Han U+4E00..9FA5 ∪ kana U+3040..30FF
284pub(crate) const AUX_SLASH: u8 = 2; // o200k rule 4's `[\r\n/]*` tail
285
286/// The bitstreams for one 64-byte block. `p0`..`p3` are the bit-planes of the **filled** dense
287/// code — filled meaning a multi-byte char sets its bits on *all* of its bytes, so "previous char's
288/// class" is a plain `<< 1` and no rule does char-width arithmetic. `cont` is the one un-filled
289/// stream (it defines the fill); `aux` is text-derived and only built when a grammar asks for it.
290///
291/// `p3` exists because o200k needs 9 tag classes (U/L/C/N/NL/SP/WSO/OTHER + cont) and code 7 is
292/// reserved — the SIMD kernels find continuation lanes by testing `lut[tag] == 7`, so "other"
293/// cannot be the leftover code. It is const-gated off for the 3-plane grammars.
294#[derive(Default, Clone, Copy)]
295pub(crate) struct Blk {
296    pub cont: u64,
297    pub p0: u64,
298    pub p1: u64,
299    pub p2: u64,
300    pub p3: u64,
301    pub aux: u64,
302}
303
304/// deepseek Split-2's isolated range: Han U+4E00..9FA5 ∪ Hiragana/Katakana U+3040..30FF (all
305/// 3-byte, leads E3..E9). Same predicate as `fsm_deepseek`'s `ds_is_cjk_at`.
306#[inline]
307pub(crate) fn is_cjk_at(text: &[u8], p: usize) -> bool {
308    let b = text[p];
309    if !(0xE3..=0xE9).contains(&b) || p + 2 >= text.len() {
310        return false;
311    }
312    let cp = ((b as u32 & 0x0F) << 12)
313        | ((text[p + 1] as u32 & 0x3F) << 6)
314        | (text[p + 2] as u32 & 0x3F);
315    (0x4E00..=0x9FA5).contains(&cp) || (0x3040..=0x30FF).contains(&cp)
316}
317
318/// The `AUX` predicate at a lead byte.
319#[inline]
320pub(crate) fn aux_at<const AUX: u8>(text: &[u8], p: usize) -> bool {
321    match AUX {
322        AUX_CJK => is_cjk_at(text, p),
323        AUX_SLASH => text[p] == b'/',
324        _ => false,
325    }
326}
327
328/// Build one block. Full blocks go through the NEON kernel; the ragged tail (and every other
329/// target) uses the portable byte-at-a-time reference. Both produce the identical `Blk`.
330/// `lut` is the grammar's tag → dense code table.
331#[inline]
332pub(crate) fn build_block<const AUX: u8, const P3: bool>(
333    text: &[u8],
334    tags: &[u8],
335    base: usize,
336    len: usize,
337    lut: &[u8; 64],
338    cur_code: u8,
339    cur_aux: bool,
340) -> (Blk, u8) {
341    #[cfg(target_arch = "aarch64")]
342    if len == 64 {
343        // SAFETY: `len == 64` means `base + 64 <= text.len() == tags.len()`.
344        return unsafe {
345            crate::simd::neon::build64::<AUX, P3>(text, tags, base, lut, cur_code, cur_aux)
346        };
347    }
348    #[cfg(target_arch = "x86_64")]
349    if len == 64 && has_ssse3() {
350        // SAFETY: `len == 64` bounds both reads, and SSSE3 is checked above.
351        return unsafe {
352            crate::simd::x86::build64::<AUX, P3>(text, tags, base, lut, cur_code, cur_aux)
353        };
354    }
355    build_block_scalar::<AUX, P3>(text, tags, base, len, lut, cur_code, cur_aux)
356}
357
358/// Portable reference builder: one byte at a time.
359pub(crate) fn build_block_scalar<const AUX: u8, const P3: bool>(
360    text: &[u8],
361    tags: &[u8],
362    base: usize,
363    len: usize,
364    lut: &[u8; 64],
365    mut cur_code: u8,
366    mut cur_aux: bool,
367) -> (Blk, u8) {
368    let mut b = Blk::default();
369    for i in 0..len {
370        let p = base + i;
371        let bit = 1u64 << i;
372        if tags[p] == CONT {
373            b.cont |= bit;
374        } else {
375            cur_code = lut[tags[p] as usize];
376            cur_aux = AUX != AUX_NONE && aux_at::<AUX>(text, p);
377        }
378        b.p0 |= bit * u64::from(cur_code & 1 != 0);
379        b.p1 |= bit * u64::from(cur_code & 2 != 0);
380        b.p2 |= bit * u64::from(cur_code & 4 != 0);
381        if P3 {
382            b.p3 |= bit * u64::from(cur_code & 8 != 0);
383        }
384        b.aux |= bit * u64::from(cur_aux);
385    }
386    (b, cur_code)
387}
388
389// ── bitstream primitives ────────────────────────────────────────────────────────────────────────
390
391/// `ScanThru`: move every marker in `m` forward past the run of `c` it sits in, landing on the
392/// first position not in `c`. The classic Parabix carry-propagation trick — one add. Done in `u128`
393/// so a run reaching bit 63 puts its landing bit at 64 instead of vanishing: that bit *is* the
394/// block's carry-out, and `e - m` still yields the right in-block span.
395#[inline]
396pub(crate) const fn scanthru(m: u128, c: u128) -> u128 {
397    m.wrapping_add(c) & !c
398}
399
400/// Advance markers by ONE CHAR: shift one byte, then scan through the continuation bytes. Markers
401/// that leave the block are dropped — their state is reconstructed from the scalar carry instead.
402#[inline]
403pub(crate) const fn adv(m: u64, cont: u64) -> u64 {
404    scanthru((m as u128) << 1, cont as u128) as u64
405}
406
407/// Move each bit back to the lead byte of its char. Rules that look at the *next* char are stated
408/// at a char's last byte; token starts must sit on leads. ≤3 steps (UTF-8 chars are ≤4 bytes).
409/// Returns `(in-block, patch for the previous word)` — a char straddling the block boundary lands
410/// its lead in the word we already wrote.
411#[inline]
412pub(crate) fn to_lead(x: u64, cont: u64, prev_cont: u64) -> (u64, u64) {
413    let c = ((cont as u128) << 64) | prev_cont as u128;
414    let mut y = (x as u128) << 64;
415    for _ in 0..3 {
416        let m = y & c;
417        if m == 0 {
418            break;
419        }
420        y = (y & !c) | (m >> 1);
421    }
422    ((y >> 64) as u64, y as u64)
423}
424
425/// The span from each marker of `m` forward to the END of the run of `c` it sits in, landing
426/// EXCLUDED. `c` may be any combination of class streams — it is just a bitmask.
427///
428/// This is the closure operator, and the one the spec called `MatchStar`. Its doc used to have the
429/// direction backwards ("fill from the run start through the LAST marker"), which is why callers
430/// reached for it through `reverse_bits` pairs and hand-rolled their smears instead of calling it.
431/// Reversed, "forward to the end of the run" reads as "is there a marker at or after me, still
432/// inside this run?" — that is [`later_in_run`].
433///
434/// Parabix's `MatchStar` includes the position the run stops on. No rule here wants that: a rule
435/// needs either the span without the landing (this) or the landing alone ([`scanthru`]). Should one
436/// ever need both, it is `fill_to_last(m, c) | scanthru(m, c) as u64`.
437#[inline]
438pub(crate) fn fill_to_last(m: u64, c: u64) -> u64 {
439    if m == 0 {
440        return 0;
441    }
442    let (m128, c128) = (m as u128, c as u128);
443    (scanthru(m128, c128).wrapping_sub(m128) | m128) as u64
444}
445
446/// **Backward.** For every position, is there a marker of `m` at or after it, inside the same run of
447/// `c`? [`fill_to_last`] spans forward from a marker, so reversed it answers exactly this. All three
448/// grammars spelled it out as a `reverse_bits` triple.
449#[inline]
450pub(crate) fn later_in_run(m: u64, c: u64) -> u64 {
451    fill_to_last(m.reverse_bits(), c.reverse_bits()).reverse_bits()
452}
453
454/// Absolute position of the highest set bit of `x`, for a block starting at `base` (`x != 0`).
455#[inline]
456pub(crate) fn last_pos(base: usize, x: u64) -> usize {
457    base + 63 - x.leading_zeros() as usize
458}
459
460/// `\p{N}{1,cap}` across a block edge: how many digits of the current group the run has consumed by
461/// the time it leaves this block. `groups` are this block's group starts, `tn` its trailing digit
462/// run. Identical in cl100k and the o200k family, so it lives here.
463#[inline]
464fn digits_since(groups: u64, tn: u64, n: u64, lead: u64, since: u32, cap: usize) -> u32 {
465    let g = groups & tn;
466    let counted = if g == 0 {
467        since + (n & lead & tn).count_ones()
468    } else {
469        (n & lead & tn & !((1u64 << (63 - g.leading_zeros())) - 1)).count_ones()
470    };
471    if cap == 0 || cap >= 64 {
472        counted
473    } else {
474        counted % cap as u32
475    }
476}
477
478/// `\p{N}{1,cap}` across block edges: whether a digit run is still open at the edge, and how far
479/// into its current group it got. All three grammars kept exactly this pair, recomputed it with the
480/// same `trail_run`, and two of them spelled the counter out identically.
481#[derive(Default)]
482pub(crate) struct Digits {
483    open: bool,
484    since: u32,
485}
486
487impl Digits {
488    /// This block's group starts: a boundary every `cap` chars from the run start.
489    pub(crate) fn starts(
490        &self,
491        cap: usize,
492        n: u64,
493        lead: u64,
494        cont: u64,
495        prev_is_digit: bool,
496    ) -> u64 {
497        digit_groups(cap, n, lead, cont, prev_is_digit, self.open, self.since)
498    }
499
500    /// Carry the run into the next block. `next_is_digit` is that block's first byte.
501    // Every one of these is a live bitstream register at the call site; bundling them into a
502    // struct would only put an indirection between the driver and the carry it just computed.
503    #[allow(clippy::too_many_arguments)]
504    pub(crate) fn commit(
505        &mut self,
506        cap: usize,
507        groups: u64,
508        n: u64,
509        lead: u64,
510        valid: u64,
511        len: usize,
512        next_is_digit: bool,
513    ) {
514        let tn = trail_run(n, valid, len);
515        self.open = tn != 0 && next_is_digit;
516        self.since = if self.open {
517            digits_since(groups, tn, n, lead, self.since, cap)
518        } else {
519            0
520        };
521    }
522}
523
524/// `\s*[\r\n]+` versus `\s+(?!\S)`: the one backward-in-time rule these grammars have. `\s*[\r\n]+`
525/// runs through a whitespace run's LAST newline, so a token opens right after it — but a newline
526/// arriving in a LATER block moves which newline is last, retracting a start already written. All
527/// three grammars carried this verbatim, comments included; it lives here instead.
528#[derive(Default)]
529pub(crate) struct Anl(Option<usize>);
530
531impl Anl {
532    /// Top of a block: a newline has arrived inside the run that was open at the edge, so the start
533    /// committed for it is no longer after the last newline.
534    pub(crate) fn retract(
535        &mut self,
536        starts: &mut [u64],
537        prev_ws: bool,
538        ws: u64,
539        nl: u64,
540        valid: u64,
541    ) {
542        if let Some(p) = self.0
543            && prev_ws
544            && ws & 1 != 0
545            && nl & lead_run(ws, valid) != 0
546        {
547            starts[p / 64] &= !(1u64 << (p % 64));
548            self.0 = None;
549        }
550    }
551
552    /// Bottom of a block: remember the start a later newline could still retract. `absorbed` is what
553    /// a punct run's `[\r\n]*` tail already claimed — a bit that tail made a run start cannot be
554    /// reached back past, because the absorption cut the whitespace run.
555    pub(crate) fn commit(
556        &mut self,
557        base: usize,
558        after_nl: u64,
559        tws: u64,
560        absorbed: u64,
561        ws_next: bool,
562        prev_ws: bool,
563    ) {
564        if tws != 0 && ws_next {
565            let a = after_nl & tws & !absorbed;
566            if a != 0 {
567                self.0 = Some(last_pos(base, a));
568            } else if !(tws & 1 != 0 && prev_ws) {
569                self.0 = None; // a fresh run with no newline yet — nothing left to retract
570            }
571        } else {
572            self.0 = None;
573        }
574    }
575}
576
577/// Run of `x` that starts at bit 0.
578#[inline]
579pub(crate) fn lead_run(x: u64, valid: u64) -> u64 {
580    let z = !x & valid;
581    if z == 0 {
582        valid
583    } else {
584        (z & z.wrapping_neg()) - 1
585    }
586}
587
588/// Run of `x` that ends at bit `len - 1` (0 if that bit is clear).
589#[inline]
590pub(crate) fn trail_run(x: u64, valid: u64, len: usize) -> u64 {
591    if x & (1u64 << (len - 1)) == 0 {
592        return 0;
593    }
594    let z = !x & valid;
595    if z == 0 {
596        valid
597    } else {
598        valid & !((1u64 << (64 - z.leading_zeros())) - 1)
599    }
600}
601
602/// `\p{N}{1,cap}` — a group boundary every `cap` chars from the run start. The one non-local rule
603/// in every tiktoken-family grammar, so it lives here rather than three times over.
604///
605/// `cap == 0` or `>= 64` means an unbounded `\p{N}+`: the run is one token and there is nothing to
606/// do (and the shifts below would be UB). `dig_since` resumes a run that crossed the block edge —
607/// re-masking with `n` at every hop matters, because the carry only says the byte *at* the edge was
608/// a digit, not that the run survived it.
609#[inline]
610fn digit_groups(
611    cap: usize,
612    n: u64,
613    lead: u64,
614    cont: u64,
615    prev_is_digit: bool,
616    dig_run: bool,
617    dig_since: u32,
618) -> u64 {
619    let mut m = n & lead & !((n << 1) | u64::from(prev_is_digit));
620    if cap == 0 || cap >= 64 {
621        return m; // an unbounded `\p{N}+`: the run is one token
622    }
623    let capu = cap as u32;
624    if dig_run && prev_is_digit {
625        let mut s = lead & lead.wrapping_neg() & n; // first lead of the block
626        for _ in 0..((capu - dig_since % capu) % capu) {
627            s = adv(s, cont) & n & lead;
628        }
629        m |= s;
630    }
631    let mut groups = m;
632    if n & cont == 0 {
633        // Fast path: every digit here is single-byte, so "CAP chars on" is a plain shift against a
634        // mask asking that the skipped positions were digits too — what the `adv` chain checks.
635        let mut nk = n;
636        let mut k = 1;
637        while k < cap {
638            nk &= n << k;
639            k += 1;
640        }
641        while m != 0 {
642            m = (m << cap) & nk;
643            groups |= m;
644        }
645    } else {
646        while m != 0 {
647            let mut e = m;
648            for _ in 0..cap {
649                e = adv(e, cont) & n & lead;
650            }
651            if e == 0 {
652                break;
653            }
654            groups |= e;
655            m = e;
656        }
657    }
658    groups
659}
660
661// ── emit ────────────────────────────────────────────────────────────────────────────────────────
662
663/// `starts` bitmap → spans. Each set bit closes the previous token and opens the next; `tzcnt`
664/// walks them at ~3 ops per token. A bit also set in `fake` closes but opens nothing — that is how
665/// the class-run family expresses a *dropped* run: it is a boundary, not a token. Pass `&[]` when
666/// every start is a real one.
667pub(crate) fn emit(starts: &[u64], fake: &[u64], nblk: usize, n: usize, out: &mut [Span]) -> usize {
668    let (mut w, mut open) = (0usize, u32::MAX);
669    for (bi, &word) in starts.iter().enumerate().take(nblk) {
670        let mut m = word;
671        while m != 0 {
672            let t = m.trailing_zeros() as usize;
673            let pos = (bi * 64 + t) as u32;
674            if open != u32::MAX {
675                out[w] = Span::new(open, pos);
676                w += 1;
677            }
678            open = if fake.get(bi).is_some_and(|f| f >> t & 1 != 0) {
679                u32::MAX
680            } else {
681                pos
682            };
683            m &= m - 1;
684        }
685    }
686    if open != u32::MAX {
687        out[w] = Span::new(open, n as u32);
688        w += 1;
689    }
690    w
691}
692
693/// Byte length of the contraction at `i` (2 or 3), or 0. `ci` picks cl100k/o200k's `(?i:)` form
694/// over GPT-2's case-sensitive one.
695#[inline]
696pub(crate) fn contr_len(text: &[u8], i: usize, ci: bool) -> usize {
697    let n = text.len();
698    if i + 1 >= n || text[i] != b'\'' || text[i + 1] >= 0x80 {
699        return 0;
700    }
701    let c1 = if ci { text[i + 1] | 0x20 } else { text[i + 1] };
702    match c1 {
703        b's' | b't' | b'm' | b'd' => 2,
704        b'r' | b'v' | b'l' if i + 2 < n && text[i + 2] < 0x80 => {
705            let c2 = if ci { text[i + 2] | 0x20 } else { text[i + 2] };
706            usize::from((matches!(c1, b'r' | b'v') && c2 == b'e') || (c1 == b'l' && c2 == b'l')) * 3
707        }
708        _ => 0,
709    }
710}
711
712/// Builder-only cost probe: runs just the byte→bitstream transpose over every block and folds the
713/// streams so nothing is optimised away. The difference against a full grammar is what the
714/// bitstream program itself costs.
715#[doc(hidden)]
716#[must_use]
717pub fn build_only(text: &[u8], tags: &[u8]) -> u64 {
718    let (mut acc, mut code, mut cjk) = (0u64, CODE_CONT, false);
719    for base in (0..text.len()).step_by(64) {
720        let len = (text.len() - base).min(64);
721        let (b, c) = build_block::<{ AUX_CJK }, false>(
722            text,
723            tags,
724            base,
725            len,
726            &models::deepseek::LUT,
727            code,
728            cjk,
729        );
730        code = c;
731        cjk = b.aux >> (len - 1) & 1 != 0;
732        acc ^= b.cont ^ b.p0 ^ b.p1 ^ b.p2 ^ b.aux;
733    }
734    acc
735}
736
737/// Convenience wrapper: classify + deepseek bitcannon over caller-owned scratch.
738#[must_use]
739pub fn pre_tokenize(text: &[u8], tags: &mut [u8], starts: &mut [u64], out: &mut [Span]) -> usize {
740    classify::classify(text, tags);
741    bitcannon_deepseek(text, tags, starts, out)
742}
743
744#[cfg(test)]
745mod op_tests {
746    use super::{fill_to_last, scanthru};
747
748    /// Brute force: from every marker, walk forward while `c` holds, marking what we cross.
749    fn star_ref(m: u64, c: u64) -> u64 {
750        let mut out = 0u64;
751        for p in 0..64 {
752            if m >> p & 1 == 0 {
753                continue;
754            }
755            out |= 1 << p;
756            let mut q = p;
757            while q < 63 && c >> (q + 1) & 1 != 0 {
758                q += 1;
759                out |= 1 << q;
760            }
761        }
762        out
763    }
764
765    /// The smear shape the grammars use: step into the run, span to its end, add the marker back.
766    #[test]
767    fn fill_to_last_is_the_closure() {
768        for &(m, c) in &[
769            (0b0001u64, 0b1110u64),
770            (0b0010, 0b1110),
771            (0b0001, 0b0000),
772            (1, !0u64 >> 1),
773            (0x0000_0001_0000_0100, 0x0000_00FE_0000_FE00),
774            (0xAAAA_AAAA_AAAA_AAAA, 0x5555_5555_5555_5555),
775        ] {
776            let stepped = (m << 1) & c;
777            let got = fill_to_last(stepped, c) | m;
778            assert_eq!(got, star_ref(m, c), "m={m:#018x} c={c:#018x}");
779        }
780    }
781
782    /// The two operators against each other, so the distinction stays pinned.
783    #[test]
784    fn scanthru_lands_and_fill_to_last_spans() {
785        let (m, c) = (0b0010u64, 0b1110u64);
786        assert_eq!(
787            scanthru(m as u128, c as u128) as u64,
788            0b10000,
789            "landing only"
790        );
791        assert_eq!(fill_to_last(m, c), 0b01110, "span, landing excluded");
792    }
793}