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}