macro_rules! streams {
($(#[$m:meta])* $name:ident { $($f:ident),+ $(,)? }) => {
$(#[$m])*
#[derive(Clone, Copy, Default)]
pub(crate) struct $name { $(pub(crate) $f: u64,)+ }
impl $name {
#[inline]
pub(crate) fn back(&self, p: &Self) -> Self {
Self { $($f: (self.$f << 1) | (p.$f >> 63),)+ }
}
#[inline]
pub(crate) fn fwd(&self, n: &Self) -> Self {
Self { $($f: (self.$f >> 1) | (n.$f << 63),)+ }
}
}
impl crate::Streams for $name {
#[inline]
fn back(&self, p: &Self) -> Self { $name::back(self, p) }
#[inline]
fn fwd(&self, n: &Self) -> Self { $name::fwd(self, n) }
#[inline]
fn lead(&self) -> u64 { self.lead }
}
};
}
pub(crate) trait Streams: Copy + Default {
fn back(&self, p: &Self) -> Self;
fn fwd(&self, n: &Self) -> Self;
fn lead(&self) -> u64;
}
pub(crate) struct Ctx<C> {
pub pv: C,
pub cur: C,
pub nx: C,
pub bk: C,
pub fw: C,
pub bi: usize,
pub base: usize,
pub len: usize,
pub valid: u64,
pub last_blk: bool,
pub lb: u64,
pub eof: u64,
}
pub(crate) struct Out {
pub st: u64,
pub patch: u64,
pub flag: u64,
}
#[inline]
pub(crate) fn was(v: u64) -> bool {
v >> 63 != 0
}
#[inline]
pub(crate) fn will(v: u64) -> bool {
v & 1 != 0
}
#[inline]
pub(crate) fn blocks<C, S, B, R>(
ntext: usize,
starts: &mut [u64],
mut flag: Option<&mut [u64]>,
seed0: S,
mut build: B,
mut rule: R,
) where
C: Streams,
B: FnMut(usize, usize, S) -> (C, S),
R: FnMut(&Ctx<C>, &mut [u64]) -> Out,
{
let mut seed = seed0;
let nblk = ntext.div_ceil(64);
let blen = |bi: usize| (ntext - bi * 64).min(64);
let mut pv = C::default();
let (mut cur, s) = build(0, blen(0), seed);
seed = s;
for bi in 0..nblk {
let base = bi * 64;
let len = blen(bi);
let last_blk = base + len == ntext;
let nx = if last_blk {
C::default()
} else {
let (c, s) = build(base + 64, blen(bi + 1), seed);
seed = s;
c
};
let valid = if len == 64 { !0u64 } else { (1u64 << len) - 1 };
let eof = if last_blk { 1u64 << (len - 1) } else { 0 };
let bk = cur.back(&pv);
let fw = cur.fwd(&nx);
let ctx = Ctx {
pv,
cur,
nx,
bk,
fw,
bi,
base,
len,
valid,
last_blk,
lb: (fw.lead() & valid) | eof,
eof,
};
let o = rule(&ctx, starts);
let mut st = o.st & cur.lead();
if bi == 0 {
st |= 1; }
starts[bi] = st;
if let Some(f) = flag.as_deref_mut() {
f[bi] = st & o.flag;
}
if bi > 0 {
starts[bi - 1] |= o.patch;
}
pv = cur;
cur = nx;
}
}
pub mod classes;
pub mod classify;
pub mod literal;
pub mod models;
pub mod regexes;
mod simd;
pub use models::cl100k::{bitcannon_cl100k, bitcannon_qwen};
pub use models::deepseek::bitcannon_deepseek;
pub use models::gpt2::bitcannon_byte_level;
pub use models::kimi::bitcannon_kimi;
pub use models::o200k::bitcannon_o200k;
pub use models::tekken::bitcannon_tekken;
#[repr(C)]
#[derive(Clone, Copy, PartialEq, Eq, Debug, Default, Hash, PartialOrd, Ord)]
pub struct Span {
pub start: u32,
pub end: u32,
}
impl Span {
#[inline]
#[must_use]
pub const fn new(start: u32, end: u32) -> Self {
Self { start, end }
}
#[inline]
#[must_use]
pub fn range(self) -> core::ops::Range<usize> {
self.start as usize..self.end as usize
}
}
impl PartialEq<(u32, u32)> for Span {
fn eq(&self, other: &(u32, u32)) -> bool {
self.start == other.0 && self.end == other.1
}
}
#[must_use]
pub fn fast_builder() -> bool {
#[cfg(target_arch = "aarch64")]
{
true
}
#[cfg(target_arch = "x86_64")]
{
has_ssse3()
}
#[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
{
false
}
}
#[cfg(target_arch = "x86_64")]
pub(crate) fn has_ssse3() -> bool {
use std::sync::atomic::{AtomicU8, Ordering};
static CACHED: AtomicU8 = AtomicU8::new(0);
match CACHED.load(Ordering::Relaxed) {
0 => {
let yes = std::arch::is_x86_feature_detected!("ssse3");
CACHED.store(1 + u8::from(yes), Ordering::Relaxed);
yes
}
n => n == 2,
}
}
pub(crate) const CONT: u8 = 15; pub(crate) const CODE_CONT: u8 = 7;
pub(crate) const AUX_NONE: u8 = 0;
pub(crate) const AUX_CJK: u8 = 1; pub(crate) const AUX_SLASH: u8 = 2;
#[derive(Default, Clone, Copy)]
pub(crate) struct Blk {
pub cont: u64,
pub p0: u64,
pub p1: u64,
pub p2: u64,
pub p3: u64,
pub aux: u64,
}
#[inline]
pub(crate) fn is_cjk_at(text: &[u8], p: usize) -> bool {
let b = text[p];
if !(0xE3..=0xE9).contains(&b) || p + 2 >= text.len() {
return false;
}
let cp = ((b as u32 & 0x0F) << 12)
| ((text[p + 1] as u32 & 0x3F) << 6)
| (text[p + 2] as u32 & 0x3F);
(0x4E00..=0x9FA5).contains(&cp) || (0x3040..=0x30FF).contains(&cp)
}
#[inline]
pub(crate) fn aux_at<const AUX: u8>(text: &[u8], p: usize) -> bool {
match AUX {
AUX_CJK => is_cjk_at(text, p),
AUX_SLASH => text[p] == b'/',
_ => false,
}
}
#[inline]
pub(crate) fn build_block<const AUX: u8, const P3: bool>(
text: &[u8],
tags: &[u8],
base: usize,
len: usize,
lut: &[u8; 64],
cur_code: u8,
cur_aux: bool,
) -> (Blk, u8) {
#[cfg(target_arch = "aarch64")]
if len == 64 {
return unsafe {
crate::simd::neon::build64::<AUX, P3>(text, tags, base, lut, cur_code, cur_aux)
};
}
#[cfg(target_arch = "x86_64")]
if len == 64 && has_ssse3() {
return unsafe {
crate::simd::x86::build64::<AUX, P3>(text, tags, base, lut, cur_code, cur_aux)
};
}
build_block_scalar::<AUX, P3>(text, tags, base, len, lut, cur_code, cur_aux)
}
pub(crate) fn build_block_scalar<const AUX: u8, const P3: bool>(
text: &[u8],
tags: &[u8],
base: usize,
len: usize,
lut: &[u8; 64],
mut cur_code: u8,
mut cur_aux: bool,
) -> (Blk, u8) {
let mut b = Blk::default();
for i in 0..len {
let p = base + i;
let bit = 1u64 << i;
if tags[p] == CONT {
b.cont |= bit;
} else {
cur_code = lut[tags[p] as usize];
cur_aux = AUX != AUX_NONE && aux_at::<AUX>(text, p);
}
b.p0 |= bit * u64::from(cur_code & 1 != 0);
b.p1 |= bit * u64::from(cur_code & 2 != 0);
b.p2 |= bit * u64::from(cur_code & 4 != 0);
if P3 {
b.p3 |= bit * u64::from(cur_code & 8 != 0);
}
b.aux |= bit * u64::from(cur_aux);
}
(b, cur_code)
}
#[inline]
pub(crate) const fn scanthru(m: u128, c: u128) -> u128 {
m.wrapping_add(c) & !c
}
#[inline]
pub(crate) const fn adv(m: u64, cont: u64) -> u64 {
scanthru((m as u128) << 1, cont as u128) as u64
}
#[inline]
pub(crate) fn to_lead(x: u64, cont: u64, prev_cont: u64) -> (u64, u64) {
let c = ((cont as u128) << 64) | prev_cont as u128;
let mut y = (x as u128) << 64;
for _ in 0..3 {
let m = y & c;
if m == 0 {
break;
}
y = (y & !c) | (m >> 1);
}
((y >> 64) as u64, y as u64)
}
#[inline]
pub(crate) fn fill_to_last(m: u64, c: u64) -> u64 {
if m == 0 {
return 0;
}
let (m128, c128) = (m as u128, c as u128);
(scanthru(m128, c128).wrapping_sub(m128) | m128) as u64
}
#[inline]
pub(crate) fn later_in_run(m: u64, c: u64) -> u64 {
fill_to_last(m.reverse_bits(), c.reverse_bits()).reverse_bits()
}
#[inline]
pub(crate) fn last_pos(base: usize, x: u64) -> usize {
base + 63 - x.leading_zeros() as usize
}
#[inline]
fn digits_since(groups: u64, tn: u64, n: u64, lead: u64, since: u32, cap: usize) -> u32 {
let g = groups & tn;
let counted = if g == 0 {
since + (n & lead & tn).count_ones()
} else {
(n & lead & tn & !((1u64 << (63 - g.leading_zeros())) - 1)).count_ones()
};
if cap == 0 || cap >= 64 {
counted
} else {
counted % cap as u32
}
}
#[derive(Default)]
pub(crate) struct Digits {
open: bool,
since: u32,
}
impl Digits {
pub(crate) fn starts(
&self,
cap: usize,
n: u64,
lead: u64,
cont: u64,
prev_is_digit: bool,
) -> u64 {
digit_groups(cap, n, lead, cont, prev_is_digit, self.open, self.since)
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn commit(
&mut self,
cap: usize,
groups: u64,
n: u64,
lead: u64,
valid: u64,
len: usize,
next_is_digit: bool,
) {
let tn = trail_run(n, valid, len);
self.open = tn != 0 && next_is_digit;
self.since = if self.open {
digits_since(groups, tn, n, lead, self.since, cap)
} else {
0
};
}
}
#[derive(Default)]
pub(crate) struct Anl(Option<usize>);
impl Anl {
pub(crate) fn retract(
&mut self,
starts: &mut [u64],
prev_ws: bool,
ws: u64,
nl: u64,
valid: u64,
) {
if let Some(p) = self.0
&& prev_ws
&& ws & 1 != 0
&& nl & lead_run(ws, valid) != 0
{
starts[p / 64] &= !(1u64 << (p % 64));
self.0 = None;
}
}
pub(crate) fn commit(
&mut self,
base: usize,
after_nl: u64,
tws: u64,
absorbed: u64,
ws_next: bool,
prev_ws: bool,
) {
if tws != 0 && ws_next {
let a = after_nl & tws & !absorbed;
if a != 0 {
self.0 = Some(last_pos(base, a));
} else if !(tws & 1 != 0 && prev_ws) {
self.0 = None; }
} else {
self.0 = None;
}
}
}
#[inline]
pub(crate) fn lead_run(x: u64, valid: u64) -> u64 {
let z = !x & valid;
if z == 0 {
valid
} else {
(z & z.wrapping_neg()) - 1
}
}
#[inline]
pub(crate) fn trail_run(x: u64, valid: u64, len: usize) -> u64 {
if x & (1u64 << (len - 1)) == 0 {
return 0;
}
let z = !x & valid;
if z == 0 {
valid
} else {
valid & !((1u64 << (64 - z.leading_zeros())) - 1)
}
}
#[inline]
fn digit_groups(
cap: usize,
n: u64,
lead: u64,
cont: u64,
prev_is_digit: bool,
dig_run: bool,
dig_since: u32,
) -> u64 {
let mut m = n & lead & !((n << 1) | u64::from(prev_is_digit));
if cap == 0 || cap >= 64 {
return m; }
let capu = cap as u32;
if dig_run && prev_is_digit {
let mut s = lead & lead.wrapping_neg() & n; for _ in 0..((capu - dig_since % capu) % capu) {
s = adv(s, cont) & n & lead;
}
m |= s;
}
let mut groups = m;
if n & cont == 0 {
let mut nk = n;
let mut k = 1;
while k < cap {
nk &= n << k;
k += 1;
}
while m != 0 {
m = (m << cap) & nk;
groups |= m;
}
} else {
while m != 0 {
let mut e = m;
for _ in 0..cap {
e = adv(e, cont) & n & lead;
}
if e == 0 {
break;
}
groups |= e;
m = e;
}
}
groups
}
pub(crate) fn emit(starts: &[u64], fake: &[u64], nblk: usize, n: usize, out: &mut [Span]) -> usize {
let (mut w, mut open) = (0usize, u32::MAX);
for (bi, &word) in starts.iter().enumerate().take(nblk) {
let mut m = word;
while m != 0 {
let t = m.trailing_zeros() as usize;
let pos = (bi * 64 + t) as u32;
if open != u32::MAX {
out[w] = Span::new(open, pos);
w += 1;
}
open = if fake.get(bi).is_some_and(|f| f >> t & 1 != 0) {
u32::MAX
} else {
pos
};
m &= m - 1;
}
}
if open != u32::MAX {
out[w] = Span::new(open, n as u32);
w += 1;
}
w
}
#[inline]
pub(crate) fn contr_len(text: &[u8], i: usize, ci: bool) -> usize {
let n = text.len();
if i + 1 >= n || text[i] != b'\'' || text[i + 1] >= 0x80 {
return 0;
}
let c1 = if ci { text[i + 1] | 0x20 } else { text[i + 1] };
match c1 {
b's' | b't' | b'm' | b'd' => 2,
b'r' | b'v' | b'l' if i + 2 < n && text[i + 2] < 0x80 => {
let c2 = if ci { text[i + 2] | 0x20 } else { text[i + 2] };
usize::from((matches!(c1, b'r' | b'v') && c2 == b'e') || (c1 == b'l' && c2 == b'l')) * 3
}
_ => 0,
}
}
#[doc(hidden)]
#[must_use]
pub fn build_only(text: &[u8], tags: &[u8]) -> u64 {
let (mut acc, mut code, mut cjk) = (0u64, CODE_CONT, false);
for base in (0..text.len()).step_by(64) {
let len = (text.len() - base).min(64);
let (b, c) = build_block::<{ AUX_CJK }, false>(
text,
tags,
base,
len,
&models::deepseek::LUT,
code,
cjk,
);
code = c;
cjk = b.aux >> (len - 1) & 1 != 0;
acc ^= b.cont ^ b.p0 ^ b.p1 ^ b.p2 ^ b.aux;
}
acc
}
#[must_use]
pub fn pre_tokenize(text: &[u8], tags: &mut [u8], starts: &mut [u64], out: &mut [Span]) -> usize {
classify::classify(text, tags);
bitcannon_deepseek(text, tags, starts, out)
}
#[cfg(test)]
mod op_tests {
use super::{fill_to_last, scanthru};
fn star_ref(m: u64, c: u64) -> u64 {
let mut out = 0u64;
for p in 0..64 {
if m >> p & 1 == 0 {
continue;
}
out |= 1 << p;
let mut q = p;
while q < 63 && c >> (q + 1) & 1 != 0 {
q += 1;
out |= 1 << q;
}
}
out
}
#[test]
fn fill_to_last_is_the_closure() {
for &(m, c) in &[
(0b0001u64, 0b1110u64),
(0b0010, 0b1110),
(0b0001, 0b0000),
(1, !0u64 >> 1),
(0x0000_0001_0000_0100, 0x0000_00FE_0000_FE00),
(0xAAAA_AAAA_AAAA_AAAA, 0x5555_5555_5555_5555),
] {
let stepped = (m << 1) & c;
let got = fill_to_last(stepped, c) | m;
assert_eq!(got, star_ref(m, c), "m={m:#018x} c={c:#018x}");
}
}
#[test]
fn scanthru_lands_and_fill_to_last_spans() {
let (m, c) = (0b0010u64, 0b1110u64);
assert_eq!(
scanthru(m as u128, c as u128) as u64,
0b10000,
"landing only"
);
assert_eq!(fill_to_last(m, c), 0b01110, "span, landing excluded");
}
}