From 2bd61b7b8d2c2d3c7eeb62595d2f09a6cc36e550 Mon Sep 17 00:00:00 2001 From: AmitMY Date: Thu, 23 Jul 2026 13:24:13 +0200 Subject: [PATCH 1/2] perf(fast): batch ingestion in Rust with gigatoken's O200k pretokenizer trainer_from_texts (fast/src/ingest.rs) takes the raw text list once and does pretokenization, unique-word dedup, graph building, doc assembly, and Trainer construction entirely in Rust with the GIL released. The Python shim gates it on the default configuration (utf8_clusters units, default GPT pretokenizer, no registered script handlers); anything else falls back to the unchanged per-document path. Pretokenization uses gigatoken's O200k implementation (our GPT_PATTERN is byte-for-byte the gpt-oss/O200k pattern), vendored under fast/src/pretok/ with MIT attribution pinned to rev 542367a3: a cargo git dependency is impossible on stable (gigatoken's manifest uses a nightly-only profile rustflags key), so the O200k subset (~3.3k lines) is vendored with the trimmed pieces documented. Only new crate dep: icu_properties. Pretokenize ladder at 20k rows (1.19M tokens, identical counts everywhere): Python HF pre_tokenize_str 576ms -> HF tokenizers crate parallel 72ms -> gigatoken O200k parallel 5ms. 20k rows / 100 merges, best of 3, digests byte-identical to the reference: - BPE: 1.73s -> 0.36s (4.8x) - BNE n=4: 2.13s -> 0.69s (3.1x) - Boundless: 3.36s -> 2.20s (1.5x; reference 68.2s -> 31x) - SuperBPE: 4.91s -> 2.07s (2.4x; reference 42.9s -> 21x) Setup collapsed 0.95s -> 0.04s, and loop times also dropped because ingestion Arc-shares each unique word graph across occurrences (1.19M occurrence trees -> 74k unique + Arc clones). Co-Authored-By: Claude Fable 5 --- fast/Cargo.lock | 205 +++ fast/Cargo.toml | 2 + .../complex_tokenization_fast/tokenizer.py | 17 +- fast/src/ingest.rs | 148 ++ fast/src/lib.rs | 5 + fast/src/pretok/mask.rs | 849 ++++++++++ fast/src/pretok/mod.rs | 398 +++++ fast/src/pretok/o200k.rs | 73 + fast/src/pretok/o200k_family.rs | 1495 +++++++++++++++++ fast/src/pretok/unicode.rs | 448 +++++ fast/src/trainer.rs | 12 + fast/src/units.rs | 15 + 12 files changed, 3664 insertions(+), 3 deletions(-) create mode 100644 fast/src/ingest.rs create mode 100644 fast/src/pretok/mask.rs create mode 100644 fast/src/pretok/mod.rs create mode 100644 fast/src/pretok/o200k.rs create mode 100644 fast/src/pretok/o200k_family.rs create mode 100644 fast/src/pretok/unicode.rs diff --git a/fast/Cargo.lock b/fast/Cargo.lock index 111723d..5812e04 100644 --- a/fast/Cargo.lock +++ b/fast/Cargo.lock @@ -27,6 +27,7 @@ checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" name = "complex-tokenization" version = "0.0.1" dependencies = [ + "icu_properties", "pyo3", "rayon", "regex", @@ -59,6 +60,17 @@ version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" +[[package]] +name = "displaydoc" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "either" version = "1.15.0" @@ -71,6 +83,68 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "icu_collections" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2984d1cd16c883d7935b9e07e44071dca8d917fd52ecc02c04d5fa0b5a3f191c" +dependencies = [ + "displaydoc", + "potential_utf", + "utf8_iter", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_locale_core" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92219b62b3e2b4d88ac5119f8904c10f8f61bf7e95b640d25ba3075e6cac2c29" +dependencies = [ + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", +] + +[[package]] +name = "icu_properties" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bee3b67d0ea5c2cca5003417989af8996f8604e34fb9ddf96208a033901e70de" +dependencies = [ + "icu_collections", + "icu_locale_core", + "icu_properties_data", + "icu_provider", + "zerotrie", + "zerovec", +] + +[[package]] +name = "icu_properties_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e2bbb201e0c04f7b4b3e14382af113e17ba4f63e2c9d2ee626b720cbce54a14" + +[[package]] +name = "icu_provider" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "139c4cf31c8b5f33d7e199446eff9c1e02decfc2f0eec2c8d71f65befa45b421" +dependencies = [ + "displaydoc", + "icu_locale_core", + "writeable", + "yoke", + "zerofrom", + "zerotrie", + "zerovec", +] + [[package]] name = "indoc" version = "2.0.7" @@ -86,6 +160,12 @@ version = "0.2.184" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "48f5d2a454e16a5ea0f4ced81bd44e4cfc7bd3a507b61887c99fd3538b28e4af" +[[package]] +name = "litemap" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" + [[package]] name = "memchr" version = "2.8.0" @@ -113,6 +193,15 @@ version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" +[[package]] +name = "potential_utf" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0103b1cef7ec0cf76490e969665504990193874ea05c85ff9bab8b911d0a0564" +dependencies = [ + "zerovec", +] + [[package]] name = "proc-macro2" version = "1.0.106" @@ -255,6 +344,12 @@ version = "1.0.22" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" +[[package]] +name = "stable_deref_trait" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" + [[package]] name = "syn" version = "2.0.117" @@ -266,12 +361,33 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "target-lexicon" version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" +[[package]] +name = "tinystr" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8323304221c2a851516f22236c5722a72eaa19749016521d6dff0824447d96d" +dependencies = [ + "displaydoc", + "zerovec", +] + [[package]] name = "unicode-ident" version = "1.0.24" @@ -289,3 +405,92 @@ name = "unindent" version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" + +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + +[[package]] +name = "writeable" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" + +[[package]] +name = "yoke" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "709fe23a0424b6a435d82152b1bd3fdfb0833487d5fa90d05d42762a9891fef5" +dependencies = [ + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zerofrom" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zerotrie" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f9152d31db0792fa83f70fb2f83148effb5c1f5b8c7686c3459e361d9bc20bf" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", +] + +[[package]] +name = "zerovec" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90f911cbc359ab6af17377d242225f4d75119aec87ea711a880987b18cd7b239" +dependencies = [ + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] diff --git a/fast/Cargo.toml b/fast/Cargo.toml index 461e71d..b6c8215 100644 --- a/fast/Cargo.toml +++ b/fast/Cargo.toml @@ -14,3 +14,5 @@ unicode-segmentation = "1.12" regex = "1" rayon = "1" rustc-hash = "2" +# For the vendored gigatoken O200k pretokenizer (see src/pretok/). +icu_properties = { version = "2", features = ["compiled_data"] } diff --git a/fast/python/complex_tokenization_fast/tokenizer.py b/fast/python/complex_tokenization_fast/tokenizer.py index 3fd2fb1..bf075dd 100644 --- a/fast/python/complex_tokenization_fast/tokenizer.py +++ b/fast/python/complex_tokenization_fast/tokenizer.py @@ -1,6 +1,6 @@ from functools import reduce -from complex_tokenization_fast._rs import Node, Trainer +from complex_tokenization_fast._rs import Node, Trainer, has_cluster_handlers_py, trainer_from_texts from complex_tokenization_fast.graphs.settings import GraphSettings from complex_tokenization_fast.graphs.units import characters, register_script, utf8, utf8_clusters from complex_tokenization_fast.graphs.words import GPTPretokenizer, words @@ -42,8 +42,19 @@ def make_trainer(self, texts): GraphSettings.ONLY_MINIMAL_MERGES = True GraphSettings.MAX_MERGE_SIZE = self.merge_size - graphs = self._build_graphs(texts) - trainer = Trainer(graphs=graphs) + # Default configuration ingests entirely in Rust (one boundary + # crossing; gigatoken's O200k pretokenizer is byte-identical to + # GPTPretokenizer). Custom pretokenizers/units/script handlers take + # the per-document Python path. + if ( + self.units is utf8_clusters + and self.pretokenizer is GPTPretokenizer + and not has_cluster_handlers_py() + ): + trainer = trainer_from_texts(list(texts), connected=self.connected) + else: + graphs = self._build_graphs(texts) + trainer = Trainer(graphs=graphs) for merge_strs in self.merges: nodes = tuple(Node(value=s.encode("utf-8")) for s in merge_strs) diff --git a/fast/src/ingest.rs b/fast/src/ingest.rs new file mode 100644 index 0000000..75fe7f0 --- /dev/null +++ b/fast/src/ingest.rs @@ -0,0 +1,148 @@ +//! Batch corpus ingestion for the default-configuration fast path: one +//! boundary crossing takes the raw text list and does pretokenization +//! (vendored gigatoken O200k — byte-identical to the GPT_PATTERN Split the +//! Python side uses), word dedup, graph building, and Trainer construction +//! entirely in Rust. The Python shim gates this on the default pretokenizer, +//! `utf8_clusters` units, and no registered script handlers; anything else +//! takes the per-document Python path. + +use pyo3::prelude::*; +use rayon::prelude::*; +use rustc_hash::FxHashMap; +use std::sync::Arc; +use unicode_segmentation::UnicodeSegmentation; + +use crate::graph::GraphV; +use crate::pretok::FastO200kPretokenizer; +use crate::trainer::Trainer; +use crate::units::{extend_cluster_cache, snapshot_cluster_cache, utf8_inner}; + +/// `utf8_clusters` for one grapheme cluster with no script handlers: a +/// single-char cluster is its UTF-8 graph, a multi-char cluster is a Seq of +/// per-char UTF-8 graphs. Mirrors `resolve_cluster`'s no-handler branch. +fn cluster_graph_pure(cluster: &str) -> GraphV { + if cluster.chars().nth(1).is_none() { + return utf8_inner(cluster); + } + let char_nodes: Vec = cluster + .chars() + .map(|c| utf8_inner(c.encode_utf8(&mut [0u8; 4]))) + .collect(); + GraphV::new_seq(char_nodes) +} + +pub(crate) fn build_corpus_graph(texts: &[String], connected: bool) -> GraphV { + // Pretokenize every document. Match boundaries are char boundaries, so + // the byte slices are valid &str. + let doc_tokens: Vec> = texts + .par_iter() + .map(|t| { + FastO200kPretokenizer::new(t.as_bytes()) + .map(|p| std::str::from_utf8(p.0).expect("pretoken on char boundary")) + .collect() + }) + .collect(); + + // Unique words, first-occurrence order. + let mut word_id: FxHashMap<&str, u32> = FxHashMap::default(); + let mut unique_words: Vec<&str> = Vec::new(); + for tokens in &doc_tokens { + for &w in tokens { + word_id.entry(w).or_insert_with(|| { + unique_words.push(w); + (unique_words.len() - 1) as u32 + }); + } + } + + // One graph per unique grapheme cluster, corpus-wide. Existing cluster + // cache entries are reused and new ones written back, exactly as the + // per-word `utf8_clusters` path would. + let mut cluster_of: FxHashMap<&str, GraphV> = FxHashMap::default(); + let cache = snapshot_cluster_cache(); + let mut new_clusters: Vec<(String, GraphV)> = Vec::new(); + for &w in &unique_words { + for cl in w.graphemes(true) { + if !cluster_of.contains_key(cl) { + let g = match cache.get(cl) { + Some(g) => g.clone(), + None => { + let g = cluster_graph_pure(cl); + new_clusters.push((cl.to_string(), g.clone())); + g + } + }; + cluster_of.insert(cl, g); + } + } + } + extend_cluster_cache(new_clusters); + + // One graph per unique word, shared by Arc across its occurrences + // (candidate counting is content-based, so sharing does not affect + // results; the trainer re-canonicalizes by content anyway). + let word_graphs: Vec = unique_words + .par_iter() + .map(|w| { + let mut nodes: Vec = + w.graphemes(true).map(|cl| cluster_of[cl].clone()).collect(); + if nodes.len() == 1 { + nodes.pop().unwrap() + } else { + GraphV::new_seq(nodes) + } + }) + .collect(); + + // Assemble docs the way `words()` + `Trainer(graphs=...)` would: + // connected docs are Seqs of their words (a single-word doc is the word + // itself); disconnected docs flatten into one word-occurrence list. + let subs: Vec = if connected { + doc_tokens + .par_iter() + .map(|tokens| { + let mut nodes: Vec = tokens + .iter() + .map(|w| word_graphs[word_id[w] as usize].clone()) + .collect(); + if nodes.len() == 1 { + nodes.pop().unwrap() + } else { + GraphV::new_seq(nodes) + } + }) + .collect() + } else { + doc_tokens + .iter() + .flat_map(|tokens| tokens.iter()) + .map(|w| word_graphs[word_id[w] as usize].clone()) + .collect() + }; + + GraphV::Unconn(Arc::new(subs)) +} + +#[pyfunction] +#[pyo3(signature = (texts, connected=false))] +pub fn trainer_from_texts(py: Python<'_>, texts: Vec, connected: bool) -> Trainer { + let graph = py.allow_threads(|| build_corpus_graph(&texts, connected)); + Trainer::from_graph(graph) +} + +#[pyfunction] +pub fn has_cluster_handlers_py() -> bool { + crate::units::has_cluster_handlers() +} + +/// Total pretoken count over `texts` — a benchmark hook for measuring the +/// vendored pretokenizer alone. +#[pyfunction] +pub fn pretokenize_count(py: Python<'_>, texts: Vec) -> usize { + py.allow_threads(|| { + texts + .par_iter() + .map(|t| FastO200kPretokenizer::new(t.as_bytes()).count()) + .sum() + }) +} diff --git a/fast/src/lib.rs b/fast/src/lib.rs index 842a1b9..1c67fc5 100644 --- a/fast/src/lib.rs +++ b/fast/src/lib.rs @@ -1,4 +1,6 @@ mod graph; +mod ingest; +mod pretok; mod settings; mod trainer; mod units; @@ -22,6 +24,9 @@ fn _rs(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(units::set_ids_reverse_dict_py, m)?)?; m.add_function(wrap_pyfunction!(units::warm_word_cache_py, m)?)?; m.add_function(wrap_pyfunction!(units::clear_word_cache, m)?)?; + m.add_function(wrap_pyfunction!(ingest::trainer_from_texts, m)?)?; + m.add_function(wrap_pyfunction!(ingest::has_cluster_handlers_py, m)?)?; + m.add_function(wrap_pyfunction!(ingest::pretokenize_count, m)?)?; m.add_function(wrap_pyfunction!(sync_settings, m)?)?; Ok(()) } diff --git a/fast/src/pretok/mask.rs b/fast/src/pretok/mask.rs new file mode 100644 index 0000000..f25ede4 --- /dev/null +++ b/fast/src/pretok/mask.rs @@ -0,0 +1,849 @@ +// Vendored from gigatoken (https://github.com/marcelroed/gigatoken), +// rev 542367a3efed134883fb4f1140b49c04e6fad3a3, MIT license. +// See src/pretok/mod.rs for what was trimmed and why. + +//! Shared infrastructure for mask-scanner pretokenizers. +//! +//! A mask-scanner pretokenizer processes 64-byte batches: SIMD classifies +//! every byte, bitmask algebra derives "a token starts here" bits, and +//! `next()` pops one bit per token — no per-token dispatch branches, which +//! is what makes it ~2x the serial scalar scanners (see +//! pretokenizer_optimization_log.md step 15). +//! +//! A scheme plugs in two functions ([`MaskScheme`]): +//! - `advance`: the scalar ground truth (also the no-SIMD iterator), +//! - `batch_masks`: `(usable, bad)` bitmasks for a 64-byte batch. `usable` +//! bits are trustworthy token starts; `bad` marks zones (non-ASCII the +//! scheme doesn't classify in-mask, batch-edge ambiguities) that +//! [`MaskState`] re-derives through `advance`, never emitting a token +//! across an unresolved zone. +//! +//! Layering, bottom to top: +//! 1. Platform SIMD primitives (`movemask64`, `ascii_masks` on NEON; +//! `ascii_masks_avx512` / `ascii_masks_avx2` on x86-64) — the only +//! per-platform code. +//! 2. Bit-domain helpers shared across schemes — platform-independent +//! u64 algebra and per-char table classification +//! (`classify_uni_chars`, `char_through`, `nn_at_full`, +//! `digit_run_splits3`), parameterized by each scheme's codepoint +//! classifier. +//! 3. Per-scheme `batch_masks` boundary algebra (in the scheme's module). +//! 4. [`MaskState`] — the scheme-agnostic batch walker: segments, bad-zone +//! gaps, scalar tail, one-batch-ahead precompute; scalar overruns stay +//! on the 64-byte grid so the precompute survives them. +//! 5. [`MaskState::fill_spans_two_phase`] — the chunked pull the encode +//! loop uses: the same masks and trust rules as `next_span`, but +//! harvested a chunk at a time into a flat boundary buffer and emitted +//! in a branch-free counted loop. + +use super::unicode::{self, CharClass}; + +// ----------------------------------------------------------------------- +// Platform SIMD primitives: aarch64 NEON (compile-time, always present) +// and x86_64 AVX-512 or AVX2 (runtime-detected; scalar fallback +// otherwise). +// ----------------------------------------------------------------------- + +/// Does this x86_64 CPU have the full AVX-512 tier (Zen 4/5, Ice +/// Lake+)? Schemes dispatch their batch classifier on this: the AVX-512 +/// front-end when true, the AVX2 one otherwise. +#[cfg(target_arch = "x86_64")] +#[inline] +pub(crate) fn avx512_scanner_available() -> bool { + // std's feature cache makes this an atomic load + bit test after the + // first call. + std::arch::is_x86_feature_detected!("avx512f") + && std::arch::is_x86_feature_detected!("avx512bw") + && std::arch::is_x86_feature_detected!("avx512vl") + && std::arch::is_x86_feature_detected!("bmi1") + && std::arch::is_x86_feature_detected!("bmi2") + && std::arch::is_x86_feature_detected!("lzcnt") + && std::arch::is_x86_feature_detected!("popcnt") +} + +/// Does this x86_64 CPU also have AVX-512 VBMI2 (native 512-bit +/// `vpcompressb`: Zen 4/5, Ice Lake+ — i.e. nearly every AVX-512 CPU, +/// but the bit is detected, not assumed: Skylake-X lacks it and stays on +/// the plain AVX-512 tier)? Gates the `X86_TIER_AVX512_VBMI2` fill tier +/// ([`MaskState::fill_spans_two_phase`]'s `_avx512_vbmi2_crc` wrapper), +/// whose `flatten_bits_avx512` needs VBMI2 on top of the scanner tier. +#[cfg(target_arch = "x86_64")] +#[inline] +pub(crate) fn avx512_fill_available() -> bool { + avx512_scanner_available() && std::arch::is_x86_feature_detected!("avx512vbmi2") +} + +/// Does this x86_64 CPU have the AVX2 tier (Haswell+, all Zen)? The bit +/// features (BMI1/2, LZCNT, POPCNT) arrived with or before AVX2 on every +/// AVX2 CPU, but are detected explicitly since the boundary algebra's +/// codegen relies on them. +#[cfg(target_arch = "x86_64")] +#[inline] +pub(crate) fn avx2_scanner_available() -> bool { + std::arch::is_x86_feature_detected!("avx2") + && std::arch::is_x86_feature_detected!("bmi1") + && std::arch::is_x86_feature_detected!("bmi2") + && std::arch::is_x86_feature_detected!("lzcnt") + && std::arch::is_x86_feature_detected!("popcnt") +} + +/// Is the SIMD mask scanner usable on this machine? aarch64 always has +/// NEON; x86_64 requires AVX-512 (Zen 4/5, Ice Lake+) or AVX2 (Haswell+, +/// Zen 1-3), detected at runtime. When this returns false, [`MaskState`] +/// runs every token through the scheme's scalar `advance`. +#[cfg(target_arch = "x86_64")] +#[inline] +pub(crate) fn simd_scanner_available() -> bool { + avx512_scanner_available() || avx2_scanner_available() +} + +#[cfg(not(target_arch = "x86_64"))] +#[inline] +pub(crate) fn simd_scanner_available() -> bool { + cfg!(target_arch = "aarch64") +} + +// The x86-64 batch classifiers are annotated +// `#[target_feature(enable = "avx512f,avx512bw,avx512vl,bmi1,bmi2,lzcnt,popcnt")]` +// (AVX-512 tier) or `#[target_feature(enable = "avx2,bmi1,bmi2,lzcnt,popcnt")]` +// (AVX2 tier). Besides the wide byte ops, the scalar-visible bit features +// (BMI1/2, LZCNT, POPCNT) are enabled so the boundary algebra inlined +// into those functions compiles to tzcnt/lzcnt/blsr instead of +// baseline-x86 bsf sequences. The sets must stay in sync with +// [`avx512_scanner_available`] / [`avx2_scanner_available`]. + +/// simdjson-style movemask: 4 mask vectors (64 lanes of 0x00/0xFF) -> u64, +/// bit i = lane i. +#[cfg(target_arch = "aarch64")] +#[inline(always)] +pub(crate) unsafe fn movemask64( + v0: std::arch::aarch64::uint8x16_t, + v1: std::arch::aarch64::uint8x16_t, + v2: std::arch::aarch64::uint8x16_t, + v3: std::arch::aarch64::uint8x16_t, +) -> u64 { + use std::arch::aarch64::*; + unsafe { + const W: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128]; + let w = vld1q_u8(W.as_ptr()); + let mut a0 = vandq_u8(v0, w); + let a1 = vandq_u8(v1, w); + let mut a2 = vandq_u8(v2, w); + let a3 = vandq_u8(v3, w); + // The 4-`addp` reduction tree (simdjson's arm64 movemask), pinned + // as asm. Written with `vpaddq_u8`, LLVM rewrites every pairwise + // add into a uzp1/uzp2/orr triple — adjacent weighted lanes have + // disjoint bits, so add == or, and the canonical or-form never + // re-forms addp — inflating each call from 9 to 17 vector ops + // (4-7 calls per 64-byte batch across the schemes). The weighted + // `and`s stay outside so the scheduler still interleaves + // neighboring calls. `addp(x, x)` lane 0..7 equals the old + // `addp(x, zero)` lanes 0..7; only lane u64 0 is read. + core::arch::asm!( + "addp {a0:v}.16b, {a0:v}.16b, {a1:v}.16b", + "addp {a2:v}.16b, {a2:v}.16b, {a3:v}.16b", + "addp {a0:v}.16b, {a0:v}.16b, {a2:v}.16b", + "addp {a0:v}.16b, {a0:v}.16b, {a0:v}.16b", + a0 = inout(vreg) a0, + a1 = in(vreg) a1, + a2 = inout(vreg) a2, + a3 = in(vreg) a3, + options(pure, nomem, nostack, preserves_flags), + ); + vgetq_lane_u64::<0>(vreinterpretq_u64_u8(a0)) + } +} + +/// One u64 mask (bit i = byte scan+i) per byte predicate, for 64 bytes. +/// The working currency of scheme boundary algebra: everything after this +/// is platform-independent u64 bit math. +#[derive(Clone, Copy, Default)] +pub(crate) struct AsciiMasks { + /// ASCII letters. + pub l: u64, + /// ASCII digits. + pub d: u64, + /// Space (0x20) only. + pub s: u64, + /// Non-newline ASCII whitespace: \t, \x0b, \x0c. + pub wt: u64, + /// Newlines: \r, \n. + pub n: u64, + /// Non-ASCII bytes (>= 0x80). + pub hi: u64, + /// ASCII apostrophes. + pub ap: u64, +} + +/// Classify `bytes[scan..scan+64]` (requires `scan + 64 <= bytes.len()`). +#[cfg(target_arch = "aarch64")] +#[inline(always)] +pub(crate) fn ascii_masks(bytes: &[u8], scan: usize) -> AsciiMasks { + use std::arch::aarch64::*; + unsafe { + let p = bytes.as_ptr().add(scan); + let mut l = [vdupq_n_u8(0); 4]; + let mut d = [vdupq_n_u8(0); 4]; + let mut s = [vdupq_n_u8(0); 4]; + let mut wt = [vdupq_n_u8(0); 4]; + let mut n = [vdupq_n_u8(0); 4]; + let mut hi = [vdupq_n_u8(0); 4]; + let mut ap = [vdupq_n_u8(0); 4]; + for i in 0..4 { + let v = vld1q_u8(p.add(16 * i)); + let lowered = vorrq_u8(v, vdupq_n_u8(0x20)); + l[i] = vcleq_u8(vsubq_u8(lowered, vdupq_n_u8(b'a')), vdupq_n_u8(25)); + d[i] = vcleq_u8(vsubq_u8(v, vdupq_n_u8(b'0')), vdupq_n_u8(9)); + s[i] = vceqq_u8(v, vdupq_n_u8(b' ')); + n[i] = vorrq_u8( + vceqq_u8(v, vdupq_n_u8(b'\r')), + vceqq_u8(v, vdupq_n_u8(b'\n')), + ); + // \t (9), \x0b (11), \x0c (12): ascii ws minus \r\n and space. + wt[i] = vbicq_u8( + vcleq_u8(vsubq_u8(v, vdupq_n_u8(9)), vdupq_n_u8(4)), + n[i], + ); + hi[i] = vcltzq_s8(vreinterpretq_s8_u8(v)); + ap[i] = vceqq_u8(v, vdupq_n_u8(b'\'')); + } + AsciiMasks { + l: movemask64(l[0], l[1], l[2], l[3]), + d: movemask64(d[0], d[1], d[2], d[3]), + s: movemask64(s[0], s[1], s[2], s[3]), + wt: movemask64(wt[0], wt[1], wt[2], wt[3]), + n: movemask64(n[0], n[1], n[2], n[3]), + hi: movemask64(hi[0], hi[1], hi[2], hi[3]), + ap: movemask64(ap[0], ap[1], ap[2], ap[3]), + } + } +} + +/// Classify `bytes[scan..scan+64]` with AVX-512 (requires +/// `scan + 64 <= bytes.len()`). One 64-byte load and one k-register +/// compare per predicate: a `__mmask64` IS the u64 the bit algebra wants, +/// so there is no movemask ladder and no lazy any-tests — every field +/// (including `hi` and `ap`) is computed unconditionally. +/// +/// Runtime-gated: callers reach this only after +/// [`simd_scanner_available`] reported AVX-512 support (enforced by +/// [`MaskState`], which otherwise never leaves the scalar path). +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx512f,avx512bw,avx512vl,bmi1,bmi2,lzcnt,popcnt")] +#[inline] +pub(crate) fn ascii_masks_avx512(bytes: &[u8], scan: usize) -> AsciiMasks { + use std::arch::x86_64::*; + unsafe { + let v = _mm512_loadu_si512(bytes.as_ptr().add(scan) as *const _); + let lowered = _mm512_or_si512(v, _mm512_set1_epi8(0x20)); + let l = _mm512_cmple_epu8_mask( + _mm512_sub_epi8(lowered, _mm512_set1_epi8(b'a' as i8)), + _mm512_set1_epi8(25), + ); + let d = _mm512_cmple_epu8_mask( + _mm512_sub_epi8(v, _mm512_set1_epi8(b'0' as i8)), + _mm512_set1_epi8(9), + ); + let s = _mm512_cmpeq_epi8_mask(v, _mm512_set1_epi8(b' ' as i8)); + let n = _mm512_cmpeq_epi8_mask(v, _mm512_set1_epi8(b'\r' as i8)) + | _mm512_cmpeq_epi8_mask(v, _mm512_set1_epi8(b'\n' as i8)); + // \t (9), \x0b (11), \x0c (12): ascii ws minus \r\n and space. + let wt = _mm512_cmple_epu8_mask( + _mm512_sub_epi8(v, _mm512_set1_epi8(9)), + _mm512_set1_epi8(4), + ) & !n; + let hi = _mm512_movepi8_mask(v) as u64; + let ap = _mm512_cmpeq_epi8_mask(v, _mm512_set1_epi8(b'\'' as i8)); + AsciiMasks { l, d, s, wt, n, hi, ap } + } +} + +/// Classify `bytes[scan..scan+64]` with AVX2 (requires +/// `scan + 64 <= bytes.len()`). Two 32-byte loads; each predicate is one +/// vector compare per half plus a `vpmovmskb` ladder into the u64 the bit +/// algebra wants — more mask-extraction traffic than the AVX-512 version +/// (whose k-register compares ARE the u64s), but the output currency is +/// identical, so everything downstream is shared. AVX2 has no unsigned +/// byte compare; `x <= lim` is `min_epu8(x, lim) == x`. +/// +/// Runtime-gated: callers reach this only after +/// [`avx2_scanner_available`] reported AVX2 support (enforced by the +/// schemes' dispatch, behind [`MaskState`]'s `simd_scanner_available` +/// gate). +/// +/// `#[inline(never)]` is load-bearing: inlined, LLVM's vector combiner +/// sees the compare vectors behind the returned u64s and pulls the +/// caller's scalar boundary algebra back into the byte-vector domain, +/// expanding every mask<->vector crossing into vpinsrb/vpextrb ladders +/// (~240 byte ops per batch, measured 3.5x slower end to end on Zen 2). +/// The AVX-512 tier has no such domain to return to (k-register compares +/// ARE the u64s), so it stays inline. +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx2,bmi1,bmi2,lzcnt,popcnt")] +#[inline(never)] +pub(crate) fn ascii_masks_avx2(bytes: &[u8], scan: usize) -> AsciiMasks { + use std::arch::x86_64::*; + unsafe { + // Closures inherit the enclosing fn's target features. + let le = |v: __m256i, lim: __m256i| -> __m256i { + _mm256_cmpeq_epi8(_mm256_min_epu8(v, lim), v) + }; + let mm = |m0: __m256i, m1: __m256i| -> u64 { + (_mm256_movemask_epi8(m0) as u32 as u64) + | ((_mm256_movemask_epi8(m1) as u32 as u64) << 32) + }; + + let p = bytes.as_ptr().add(scan); + let v0 = _mm256_loadu_si256(p as *const _); + let v1 = _mm256_loadu_si256(p.add(32) as *const _); + + let x20 = _mm256_set1_epi8(0x20); + let ca = _mm256_set1_epi8(b'a' as i8); + let c25 = _mm256_set1_epi8(25); + let l = mm( + le(_mm256_sub_epi8(_mm256_or_si256(v0, x20), ca), c25), + le(_mm256_sub_epi8(_mm256_or_si256(v1, x20), ca), c25), + ); + let c0 = _mm256_set1_epi8(b'0' as i8); + let c9 = _mm256_set1_epi8(9); + let d = mm( + le(_mm256_sub_epi8(v0, c0), c9), + le(_mm256_sub_epi8(v1, c0), c9), + ); + let sp = _mm256_set1_epi8(b' ' as i8); + let s = mm(_mm256_cmpeq_epi8(v0, sp), _mm256_cmpeq_epi8(v1, sp)); + let cr = _mm256_set1_epi8(b'\r' as i8); + let lf = _mm256_set1_epi8(b'\n' as i8); + let n = mm( + _mm256_or_si256(_mm256_cmpeq_epi8(v0, cr), _mm256_cmpeq_epi8(v0, lf)), + _mm256_or_si256(_mm256_cmpeq_epi8(v1, cr), _mm256_cmpeq_epi8(v1, lf)), + ); + // \t (9), \x0b (11), \x0c (12): ascii ws minus \r\n and space. + let c4 = _mm256_set1_epi8(4); + let wt = mm( + le(_mm256_sub_epi8(v0, c9), c4), + le(_mm256_sub_epi8(v1, c9), c4), + ) & !n; + let hi = mm(v0, v1); // vpmovmskb takes the sign bit directly + let apc = _mm256_set1_epi8(b'\'' as i8); + let ap = mm(_mm256_cmpeq_epi8(v0, apc), _mm256_cmpeq_epi8(v1, apc)); + AsciiMasks { l, d, s, wt, n, hi, ap } + } +} + +// ----------------------------------------------------------------------- +// Bit-domain helpers (platform-independent) +// ----------------------------------------------------------------------- + +/// Is the char starting at `idx` NOT whitespace (`\S` for a `(?!\S)` +/// lookahead)? Full answer via the packed table. +/// +/// # Safety +/// +/// `idx < bytes.len()`, and when `bytes[idx]` is non-ASCII, +/// `idx + 4 <= bytes.len()` (the guardless [`decode_cp_inbounds`] read). +/// The batch classifiers' `scan + 70 <= len` guard covers every call +/// site's worst case (`idx = scan + 64`). +#[inline(always)] +pub(crate) unsafe fn nn_at_full(bytes: &[u8], idx: usize) -> bool { + use super::{decode_cp_inbounds, is_ascii_ws}; + let b = bytes[idx]; + if b < 0x80 { + return !is_ascii_ws(b); + } + // SAFETY: caller guarantees idx + 4 <= len for a non-ASCII byte here + // (this fn's contract). + let (cp, _) = unsafe { decode_cp_inbounds(bytes, idx) }; + unicode::class_of(cp) != CharClass::Whitespace +} + +/// The char containing byte `pos - 1` (`pos > 0`, valid UTF-8): its +/// class, lead index, and end (exclusive). `end > pos` iff the char +/// straddles across `pos`. ASCII classifies with the byte predicates; +/// multi-byte chars walk back to their lead (at most 3 bytes) and use +/// the packed table — this is what lets a batch after a unicode char +/// compute true boundary carries instead of deferring to a bad zone. +/// `class`: the scheme's codepoint classifier (`unicode::class_of`, or a +/// mark-folding view like `unicode::class_of_marks_join`). +/// +/// # Safety +/// +/// `pos > 0`, and when `bytes[pos - 1]` is non-ASCII, +/// `pos + 3 <= bytes.len()`: the walk-back lead `j` satisfies +/// `j <= pos - 1`, so the guardless [`decode_cp_inbounds`] read needs +/// `j + 4 <= pos + 3` in-bounds bytes. The batch classifiers' `scan + 70 +/// <= len` guard covers every call site (`pos <= scan + 64`). +#[inline(always)] +pub(crate) unsafe fn char_through( + bytes: &[u8], + pos: usize, + class: impl Fn(u32) -> CharClass, +) -> (CharClass, usize, usize) { + use super::{decode_cp_inbounds, is_ascii_ws, is_digit, is_letter}; + let b = bytes[pos - 1]; + if b < 0x80 { + let cls = if is_letter(b) { + CharClass::Letter + } else if is_digit(b) { + CharClass::Number + } else if is_ascii_ws(b) { + CharClass::Whitespace + } else { + CharClass::Other + }; + return (cls, pos - 1, pos); + } + let mut j = pos - 1; + while j > 0 && bytes[j] & 0xC0 == 0x80 { + j -= 1; + } + // SAFETY: j < pos and pos + 3 <= len (this fn's contract), so + // j + 4 <= len. + let (cp, l) = unsafe { decode_cp_inbounds(bytes, j) }; + (class(cp), j, j + l) +} + +/// Per-byte class masks for a batch's unicode chars, classified with the +/// packed table (`unicode::class_of`) — the same lookup the scalar paths +/// do. Every byte of a classified char carries the char's class, so +/// byte-adjacency == char-adjacency and the schemes' u64 boundary +/// algebra applies unchanged. +#[derive(Clone, Copy, Default)] +pub(crate) struct UniClasses { + /// Letter / number / other / whitespace bytes. + pub l: u64, + pub n: u64, + pub o: u64, + pub ws: u64, + /// Whitespace lead bits by char length, for the char-length-aware + /// `(?!\S)` shift tests. Deferred ws chars (see `resid`) are not + /// included. + pub w2: u64, + pub w3: u64, + /// Lead bits of all classified chars by length, for schemes that + /// shift a test by the previous char's length (the cl100k family's + /// two-chars-back rule). + pub lead2: u64, + pub lead3: u64, + pub lead4: u64, + /// Continuation bytes of classified chars. + pub cont: u64, + /// Bytes only the scalar path can decide: whitespace chars straddling + /// the batch end (their run-split bookkeeping crosses the boundary), + /// number chars when `NUMBERS` is false, and stray continuation + /// bytes. Class masks stay truthful for these bytes so neighbors' + /// algebra is exact; callers turn `resid` into bad zones (±1 smear). + pub resid: u64, +} + +/// Classify every unicode char whose lead bit is in `m` (typically +/// `hi & !claimed-straddle-in-bytes`) for `bytes[scan..scan+64]`. +/// A char spilling off the batch end is classified via the lookahead +/// bytes; only its in-batch bytes get class bits, and the next +/// batch's `char_through` walk-back covers the remainder. `NUMBERS`: +/// false for schemes whose digit grouping is char-counted (`\p{N}{1,3}` +/// byte masks can't express multi-byte chars), true otherwise. +/// `LEADS`: whether to fill the per-length lead masks (only schemes with +/// a shift-by-prev-char-length rule need them). +/// +/// The loop stays branchy on purpose: a branchless csel-selected +/// decode/classify body measured 0.986x (predicted branches beat data +/// chains, log step 13/17). 2-byte chars (nearly all non-ASCII in western +/// corpora) take a dedicated lane with an inline decode; 3/4-byte chars +/// pay the general ladder. +/// +/// # Safety +/// +/// `scan + 70 <= bytes.len()` (the batch classifiers' lookahead guard): +/// a lead bit at position 63 puts the guardless [`decode_cp_inbounds`] +/// read at `scan + 63`, which may touch through `scan + 67`. +#[inline(always)] +pub(crate) unsafe fn classify_uni_chars( + bytes: &[u8], + scan: usize, + mut m: u64, + class: impl Fn(u32) -> CharClass, +) -> UniClasses { + use super::decode_cp_inbounds; + let mut u = UniClasses::default(); + while m != 0 { + let i = m.trailing_zeros() as usize; + m &= m - 1; + let b = bytes[scan + i]; + if b < 0xE0 { + // 2-byte lane (leads 0xC2..0xDF, cp < 0x800): nearly every + // non-ASCII char in western corpora, so this branch predicts + // taken and skips the length ladder + general decode. + if b < 0xC2 { + u.resid |= 1 << i; // stray continuation byte (invalid UTF-8) + continue; + } + let lead = 1u64 << i; + let chm = 3u64 << i; // in-batch bytes (excess drops at bit 63) + // SAFETY: scan + 70 <= len (this fn's # Safety contract), + // i <= 63, so scan + i + 1 <= scan + 64 < len. + let b1 = unsafe { *bytes.get_unchecked(scan + i + 1) }; + let cp = ((b as u32 & 0x1F) << 6) | (b1 as u32 & 0x3F); + match class(cp) { + CharClass::Letter => u.l |= chm, + CharClass::Number => { + u.n |= chm; + if !NUMBERS { + u.resid |= chm; + } + } + CharClass::Other => u.o |= chm, + CharClass::Whitespace => { + u.ws |= chm; + if i + 2 > 64 { + // Straddling-out ws stays a bad zone; its true + // class marks keep neighbors' `(?!\S)` tests + // exact. + u.resid |= chm; + } else { + u.w2 |= lead; + } + } + } + if LEADS { + u.lead2 |= lead; + } + u.cont |= chm & !lead; + m &= !chm; + continue; + } + let l = if b < 0xF0 { 3 } else { 4 }; + let chm = ((1u64 << l) - 1) << i; // in-batch bytes (excess drops) + let lead = 1u64 << i; + // SAFETY: scan + 70 <= len (this fn's # Safety contract), i <= 63, + // so scan + i + 4 <= len even for a 4-byte lead at bit 63. + let (cp, _) = unsafe { decode_cp_inbounds(bytes, scan + i) }; + match class(cp) { + CharClass::Letter => u.l |= chm, + CharClass::Number => { + u.n |= chm; + if !NUMBERS { + u.resid |= chm; + } + } + CharClass::Other => u.o |= chm, + CharClass::Whitespace => { + u.ws |= chm; + if i + l > 64 || l == 4 { + // Straddling-out ws (and defensively: no 4-byte cp + // is ws in Unicode) stays a bad zone; its true class + // marks keep neighbors' `(?!\S)` tests exact. + u.resid |= chm; + } else { + u.w3 |= lead; + } + } + } + if LEADS { + if l == 3 { + u.lead3 |= lead; + } else { + u.lead4 |= lead; + } + } + u.cont |= chm & !lead; + m &= !chm; + } + u +} + +/// Token-start bits inside ASCII digit runs for `\p{N}{1,3}`: each run +/// splits into 3-char tokens, so boundaries sit at run start + 3k. (For a +/// plain `\p{N}` scheme every digit is a start — no helper needed.) +#[inline(always)] +pub(crate) fn digit_run_splits3(d: u64) -> u64 { + let mut b = d & !(d << 1); // run starts + // A start at p re-arms at p+3 while the run continues: hop condition + // c = "p..p+3 all digits". Log-doubling covers 64-bit runs in 5 steps. + let mut c = d & (d >> 1) & (d >> 2) & (d >> 3); + let mut sh = 3u32; + while sh < 64 { + b |= (b & c) << sh; + c &= c >> sh; + sh <<= 1; + } + b +} + +// ----------------------------------------------------------------------- +// The batch walker +// ----------------------------------------------------------------------- + +/// The two per-scheme hooks of a mask-scanner pretokenizer. +pub(crate) trait MaskScheme { + /// Scalar ground truth: end of the token starting at `pos` + /// (`pos < bytes.len()`, `pos` on a token boundary). + fn advance(bytes: &[u8], pos: usize) -> usize; + + /// `(usable, bad)` for `bytes[scan..scan+64]` (`scan+64 <= len`): + /// `usable` bit k = trustworthy token start at scan+k; `bad` bit k = + /// byte scan+k needs the scalar path. `usable & bad` must be 0. + #[cfg(target_arch = "aarch64")] + fn batch_masks(bytes: &[u8], scan: usize) -> (u64, u64); + + /// The x86_64 batch classifier, monomorphized on the SIMD tier + /// (`AVX512` = true → the AVX-512 front-end, false → AVX2); same + /// `(usable, bad)` contract as the aarch64 `batch_masks`. The fill + /// wrappers instantiate this inside a matching `#[target_feature]` + /// region, so the tier function inlines into the fill loop and no + /// per-batch dispatch survives (the codegen a `-C target-cpu=native` + /// build gets). + /// + /// # Safety + /// + /// The selected tier must have been runtime-detected: + /// [`avx512_scanner_available`] for `AVX512` = true, + /// [`avx2_scanner_available`] for `AVX512` = false. + #[cfg(target_arch = "x86_64")] + unsafe fn batch_masks_x86(bytes: &[u8], scan: usize) -> (u64, u64); + + /// Runtime-dispatched form of [`Self::batch_masks_x86`] for call + /// sites outside a tier-monomorphized region (`next_span`): a cached + /// tier check plus a non-inlined call per batch into a per-tier + /// `#[target_feature]` wrapper ([`batch_masks_dyn_avx512`] / + /// [`batch_masks_dyn_avx2`]), so the classifier body still compiles + /// under the full tier feature set. Must only be called when + /// [`simd_scanner_available`] is true — [`MaskState`] guarantees this + /// by never leaving the scalar path otherwise. + #[cfg(target_arch = "x86_64")] + #[inline(always)] + fn batch_masks(bytes: &[u8], scan: usize) -> (u64, u64) + where + Self: Sized, + { + debug_assert!(simd_scanner_available()); + // The tier check is a cached atomic load + bit test and the + // branch is perfectly predicted, so it is noise next to the + // batch classification it selects. + if avx512_scanner_available() { + // SAFETY: runtime AVX-512 detection right above. + unsafe { batch_masks_dyn_avx512::(bytes, scan) } + } else { + // SAFETY: MaskState enables the mask-scanner path only after + // runtime detection (simd_scanner_available); without AVX-512 + // that detection was the AVX2 tier's. + unsafe { batch_masks_dyn_avx2::(bytes, scan) } + } + } +} + +/// AVX-512 feature region for the runtime-dispatched +/// `MaskScheme::batch_masks`: the scheme's `#[inline(always)]` +/// `batch_masks_x86` body fuses into this wrapper, so the per-batch call +/// `next_span` pays runs full-tier codegen (without this region the body +/// would inline into the plain-feature caller, where the inner +/// `#[target_feature]` mask classifiers can't inline and the boundary +/// algebra loses BMI/LZCNT codegen — measured ~25% slower). +/// +/// # Safety +/// +/// The CPU must support the AVX-512 scanner tier +/// ([`avx512_scanner_available`]). +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx512f,avx512bw,avx512vl,bmi1,bmi2,lzcnt,popcnt")] +#[inline] +unsafe fn batch_masks_dyn_avx512(bytes: &[u8], scan: usize) -> (u64, u64) { + // SAFETY: the caller detected the AVX-512 tier (fn contract). + unsafe { S::batch_masks_x86::(bytes, scan) } +} + +/// AVX2 counterpart of [`batch_masks_dyn_avx512`]. +/// +/// # Safety +/// +/// The CPU must support the AVX2 scanner tier +/// ([`avx2_scanner_available`]). +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx2,bmi1,bmi2,lzcnt,popcnt")] +#[inline] +unsafe fn batch_masks_dyn_avx2(bytes: &[u8], scan: usize) -> (u64, u64) { + // SAFETY: the caller detected the AVX2 tier (fn contract). + unsafe { S::batch_masks_x86::(bytes, scan) } +} + +/// x86 SIMD-tier selector for the monomorphized fill bodies +/// ([`MaskState::fill_spans_two_phase_impl`]): `DYN` keeps the per-batch +/// runtime dispatch of the provided `MaskScheme::batch_masks`; `AVX2` / +/// `AVX512` pin the tier, chosen once per fill inside a matching +/// `#[target_feature]` wrapper. `AVX512_VBMI2` is the AVX-512 tier plus +/// VBMI2 ([`avx512_fill_available`]): same batch classifiers, but phase +/// A's flatten runs `vpcompressb` (`flatten_bits_avx512`) — its only +/// divergence. Meaningless (and always `DYN`) off x86_64. +pub(crate) const X86_TIER_DYN: u8 = 0; +pub(crate) const X86_TIER_AVX2: u8 = 1; +pub(crate) const X86_TIER_AVX512: u8 = 2; +pub(crate) const X86_TIER_AVX512_VBMI2: u8 = 3; + +/// Scheme-agnostic mask-scanner state: pops trusted boundary bits, walks +/// bad zones through the scheme's scalar `advance`, runs the buffer tail +/// scalar, and precomputes one batch ahead so the SIMD chain retires under +/// the previous batch's pops. Without SIMD support (non-aarch64/x86_64 +/// targets, or an x86_64 CPU without AVX-512 or AVX2) `scalar_until` +/// starts at `usize::MAX`, so every token takes the scalar path. +pub(crate) struct MaskState { + /// Start of the pending (not yet emitted) token. + pub pos: usize, + /// Base of the next batch to scan. + scan: usize, + /// Base the `rem`/`batch_*` bits refer to. + mask_base: usize, + /// Boundary bits of the current segment (trusted, pop-ready). + rem: u64, + /// Full usable mask of the current batch (later segments). + batch_usable: u64, + /// Bad zones of the current batch not yet passed. + batch_bad: u64, + /// Emit tokens via the scalar advance while `pos < scalar_until`. + scalar_until: usize, + /// Eagerly computed masks for the batch at `pre_base` (usize::MAX = + /// none). + pre_base: usize, + pre_usable: u64, + pre_bad: u64, +} + +impl MaskState { + #[inline] + pub(crate) fn new(pos: usize) -> Self { + let scalar_until = if simd_scanner_available() { pos } else { usize::MAX }; + Self { + pos, + scan: pos, + mask_base: pos, + rem: 0, + batch_usable: 0, + batch_bad: 0, + scalar_until, + pre_base: usize::MAX, + pre_usable: 0, + pre_bad: 0, + } + } + + /// Load the segment of `batch_usable` bits in [from_bit, next bad run) + /// into `rem` and aim `scalar_until` past that bad run at the next + /// trusted boundary (or the batch end). + #[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))] + #[inline(always)] + fn load_segment(&mut self, from_bit: u32) { + let live = u64::MAX << from_bit; + let seg_bad = self.batch_bad & live; + if seg_bad == 0 { + self.rem = self.batch_usable & live; + self.batch_bad = 0; + } else { + let nb = seg_bad.trailing_zeros(); + self.rem = self.batch_usable & live & ((1u64 << nb) - 1); + let rest = self.batch_usable & (u64::MAX << nb); + self.scalar_until = if rest != 0 { + self.mask_base + rest.trailing_zeros() as usize + } else { + self.mask_base + 64 + }; + } + // A bit at the pending token's own start is not an end. Branchless: + // whether the pending token starts exactly at this segment's first + // bit is a ~20% coin flip on natural text. + let at_start = self.pos == self.mask_base + from_bit as usize; + self.rem &= !(u64::from(at_start) << from_bit); + } + + /// The next token's byte range, or None at end of input. + #[inline(always)] + pub(crate) fn next_span(&mut self, bytes: &[u8]) -> Option<(usize, usize)> { + let len = bytes.len(); + loop { + if self.rem != 0 { + let tz = self.rem.trailing_zeros() as usize; + let end = self.mask_base + tz; + self.rem &= self.rem - 1; + let start = self.pos; + self.pos = end; + return Some((start, end)); + } + if self.pos < self.scalar_until { + if self.pos >= len { + return None; + } + let start = self.pos; + let end = S::advance(bytes, start); + self.pos = end; + return Some((start, end)); + } + #[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))] + { + // Continue with the current batch's next trusted segment + // after a scalar gap (each batch is computed exactly once). + if self.batch_bad != 0 && self.pos < self.mask_base + 64 { + self.load_segment((self.pos - self.mask_base) as u32); + continue; + } + self.batch_bad = 0; + // Resume after a scalar overrun WITHOUT leaving the + // 64-byte grid: the precomputed next batch (and the + // prefetch chain behind it) stays valid, where rebasing + // to the token boundary invalidated it on every bad-zone + // overrun — a large part of a deferral's ~800-cycle + // cost. Grid bits below `pos` may be stale run-internal + // bits (a ws or digit run the scalar walked through can + // cross the grid base); they are masked by the + // `from_bit` passed to load_segment below, and every + // path that puts `pos` inside such a run goes through a + // deferral first, so those bits are never trusted. + while self.scan + 64 <= self.pos { + self.scan += 64; + } + if self.scan + 64 > len { + // Tail: scalar to the end of the buffer. + self.scalar_until = usize::MAX; + continue; + } + let (usable, bad) = if self.pre_base == self.scan { + (self.pre_usable, self.pre_bad) + } else { + S::batch_masks(bytes, self.scan) + }; + self.mask_base = self.scan; + self.scan += 64; + self.batch_usable = usable; + self.batch_bad = bad; + // Kick off the next batch now; its SIMD chain overlaps this + // batch's pops instead of stalling the next refill. Also + // done for dirty batches: a scalar overrun past the batch + // end just leaves the precompute unused (`pre_base` misses), + // while gaps that resolve inside the batch — the common + // case — keep the pipeline primed. Dirty batches used to + // skip this, and paying the whole SIMD chain latency at the + // next refill was a large part of their ~270-cycle cost. + if self.scan + 64 <= len { + let (u2, b2) = S::batch_masks(bytes, self.scan); + self.pre_base = self.scan; + self.pre_usable = u2; + self.pre_bad = b2; + } else { + self.pre_base = usize::MAX; + } + // An overrun may have left `pos` inside this grid batch; + // start from its bit so stale bits below never pop. The + // no-overrun case keeps the constant argument (and its + // folded codegen) — schemes with few bad zones take that + // branch essentially always. + if self.pos > self.mask_base { + self.load_segment((self.pos - self.mask_base) as u32); + } else { + self.load_segment(0); + } + } + #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))] + { + // Unreachable: scalar_until is usize::MAX on this arch. + self.scalar_until = usize::MAX; + } + } + } +} diff --git a/fast/src/pretok/mod.rs b/fast/src/pretok/mod.rs new file mode 100644 index 0000000..248a245 --- /dev/null +++ b/fast/src/pretok/mod.rs @@ -0,0 +1,398 @@ +// Vendored from gigatoken (https://github.com/marcelroed/gigatoken), +// rev 542367a3efed134883fb4f1140b49c04e6fad3a3, MIT license. +// See src/pretok/mod.rs for what was trimmed and why. + +//! The gigatoken O200k pretokenizer, vendored. +//! +//! Why vendored and not a git dependency: gigatoken's manifest uses +//! `[profile.profiling] rustflags`, which stable cargo refuses to parse even +//! in dependencies (the crate itself builds on a pinned nightly), and every +//! rev that has the O200k scheme also has that key. This is the subset the +//! Iterator interface needs: the O200k scheme, its shared mask-scanner +//! infrastructure, the unicode class tables, and the scalar helpers below. +//! Dropped relative to upstream: the other schemes (r50k/cl100k/qwen/...), +//! the SpanBatch/PretokenSpans encode-loop machinery, and upstream tests +//! (they compare against fancy-regex and a local OWT corpus; our digest +//! parity against the reference implementation covers the same ground). +//! +//! `FastO200kPretokenizer` iterates `Pretoken` byte slices exactly matching +//! the o200k_base regex, which is byte-for-byte our GPT_PATTERN. +#![allow(dead_code)] +// The NEON movemask asm writes its accumulator registers in place; rustc's +// dataflow can't see through `asm!` operands. +#![allow(unused_assignments)] +#![allow(clippy::all)] + +pub(crate) mod mask; +pub(crate) mod o200k_family; +pub(crate) mod unicode; +pub(crate) mod o200k; + +pub(crate) use o200k::FastO200kPretokenizer; + +use std::ops::Deref; + +#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)] +pub struct Pretoken<'a>(pub &'a [u8]); + +impl AsRef<[u8]> for Pretoken<'_> { + fn as_ref(&self) -> &[u8] { + self.0 + } +} + +impl<'a> Deref for Pretoken<'a> { + type Target = &'a [u8]; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +// ----------------------------------------------------------------------- +// Branchless byte predicates +// ----------------------------------------------------------------------- + +#[inline(always)] +pub(crate) fn is_letter(b: u8) -> bool { + (b | 0x20).wrapping_sub(b'a') < 26 +} + +#[inline(always)] +pub(crate) fn is_digit(b: u8) -> bool { + b.wrapping_sub(b'0') < 10 +} + +#[inline(always)] +pub(crate) fn is_ascii_ws(b: u8) -> bool { + b == b' ' || b.wrapping_sub(9) < 5 +} + +#[inline(always)] +pub(crate) unsafe fn decode_non_ascii(bytes: &[u8]) -> char { + unsafe { + std::str::from_utf8_unchecked(bytes) + .chars() + .next() + .unwrap_unchecked() + } +} + +/// Decode one non-ASCII scalar. Requires only `pos < bytes.len()` and +/// `bytes[pos] >= 0x80`; arbitrary (invalid) bytes are tolerated and +/// decode deterministically. Returns the codepoint and the number of +/// bytes consumed. +/// +/// Invalid input is garbage-in/defined-garbage-out, with two hard +/// guarantees the walkers rely on: +/// +/// - Never reads past `bytes.len()`, and the returned length never +/// overruns it: a multi-byte lead whose sequence is cut off by the +/// buffer end (a truncated tail) consumes exactly the bytes that +/// remain and yields [`CP_INVALID`]. (Pre-fix this read up to 3 bytes +/// past the slice and returned an end past `len` — walker panic on the +/// Iterator path, out-of-bounds span on the SpanBatch path.) +/// - The codepoint is always `<= 0x10FFFF`, so the packed class-table +/// lookups (`unicode::class_of` / `ds_class_of`, indexed unchecked) +/// stay in bounds: invalid leads 0xF5..=0xFF take the 4-byte branch +/// and can assemble "codepoints" up to 0x1FFFFF, which are clamped to +/// [`CP_INVALID`]. (Pre-fix the table lookup read up to ~246 KB past +/// the table — heap memory whose contents depend on other threads' +/// allocations, which is what made >65 KB invalid-UTF-8 pretokens +/// split nondeterministically between the walker paths.) +/// +/// The clamp target U+10FFFF is unassigned (a noncharacter) — class +/// `Other` in every scheme's classifier — so truncated or +/// beyond-Unicode garbage classifies like any other unassigned +/// codepoint, identically on every path. +#[inline(always)] +pub(crate) unsafe fn decode_cp(bytes: &[u8], pos: usize) -> (u32, usize) { + if pos + 4 > bytes.len() { + // Within 3 bytes of the buffer end: the only region where a + // sequence can be truncated. Cold: interior calls (the hot ones) + // never take it, and the branch predicts not-taken. + return decode_cp_near_end(bytes, pos); + } + // SAFETY: pos + 4 <= len just checked. + unsafe { decode_cp_inbounds(bytes, pos) } +} + +/// [`decode_cp`] without the buffer-end guard, for callers that already +/// guarantee `pos + 4 <= bytes.len()` structurally (the mask-scanner +/// batch helpers, whose `scan + 70 <= len` batch guard covers every call +/// site's worst case) — keeps the tail check out of the batch +/// classification path. Identical results to [`decode_cp`] on any input +/// where both are callable, including the [`CP_INVALID`] clamp for +/// beyond-Unicode garbage from invalid 4-byte leads. +#[inline(always)] +pub(crate) unsafe fn decode_cp_inbounds(bytes: &[u8], pos: usize) -> (u32, usize) { + unsafe { + let b0 = *bytes.get_unchecked(pos) as u32; + let b1 = (*bytes.get_unchecked(pos + 1) & 0x3F) as u32; + if b0 < 0xE0 { + return (((b0 & 0x1F) << 6) | b1, 2); + } + let b2 = (*bytes.get_unchecked(pos + 2) & 0x3F) as u32; + if b0 < 0xF0 { + return (((b0 & 0x0F) << 12) | (b1 << 6) | b2, 3); + } + let b3 = (*bytes.get_unchecked(pos + 3) & 0x3F) as u32; + ( + (((b0 & 0x07) << 18) | (b1 << 12) | (b2 << 6) | b3).min(CP_INVALID), + 4, + ) + } +} + +/// The codepoint reported for byte sequences that cannot be decoded +/// within bounds (truncated tails) or that assemble past the Unicode +/// range (invalid 4-byte-lead garbage). U+10FFFF: the largest scalar +/// value, an unassigned noncharacter, class `Other` in [`unicode::class_of`], +/// `class_of_marks_join`, and `ds_class_of` alike. +pub(crate) const CP_INVALID: u32 = 0x10FFFF; + +/// [`decode_cp`]'s slow path for `pos + 4 > bytes.len()`: decodes with +/// per-byte bounds, identical results to the fast path for complete +/// sequences; a sequence truncated by the buffer end consumes exactly +/// the remaining bytes and yields [`CP_INVALID`]. +#[cold] +#[inline(never)] +fn decode_cp_near_end(bytes: &[u8], pos: usize) -> (u32, usize) { + let len = bytes.len(); + let b0 = bytes[pos] as u32; + let need = if b0 < 0xE0 { + 2 + } else if b0 < 0xF0 { + 3 + } else { + 4 + }; + if pos + need > len { + // Truncated tail: consume the rest of the buffer as one + // unclassifiable char so every walker path terminates the final + // pretoken at `len` the same way. + return (CP_INVALID, len - pos); + } + let b1 = (bytes[pos + 1] & 0x3F) as u32; + if need == 2 { + return (((b0 & 0x1F) << 6) | b1, 2); + } + let b2 = (bytes[pos + 2] & 0x3F) as u32; + if need == 3 { + return (((b0 & 0x0F) << 12) | (b1 << 6) | b2, 3); + } + let b3 = (bytes[pos + 3] & 0x3F) as u32; + ( + (((b0 & 0x07) << 18) | (b1 << 12) | (b2 << 6) | b3).min(CP_INVALID), + 4, + ) +} + +/// `[\r\n]*`: advance past a run of CR/LF bytes (trailing newlines after a +/// punctuation run in the cl100k-family schemes). +#[inline(always)] +pub(crate) fn scan_newlines(bytes: &[u8], mut pos: usize) -> usize { + while pos < bytes.len() { + let b = unsafe { *bytes.get_unchecked(pos) }; + if b == b'\r' || b == b'\n' { + pos += 1; + } else { + break; + } + } + pos +} + +/// If the char at `pos` is a letter (`\p{L}` under the 4-way `CharClass` +/// classifier), return the offset just past it. +#[inline(always)] +pub(crate) fn letter_end_at(bytes: &[u8], pos: usize) -> Option { + let &b = bytes.get(pos)?; + if is_letter(b) { + return Some(pos + 1); + } + if b >= 0x80 { + let (cp, l) = unsafe { decode_cp(bytes, pos) }; + if unicode::class_of(cp) == unicode::CharClass::Letter { + return Some(pos + l); + } + } + None +} + +// ----------------------------------------------------------------------- +// SWAR +// ----------------------------------------------------------------------- + +pub(crate) const HI: u64 = 0x8080_8080_8080_8080; + +/// Returns the high bit set in each lane that is NOT an ASCII letter, +/// computed directly (rather than as the complement of a letter mask) so +/// the scan loop can branch on `!= 0` and reuse the value for `trailing_zeros`. +#[inline(always)] +pub(crate) fn swar64_letter_nonmask(word: u64) -> u64 { + let lowered = word | 0x2020_2020_2020_2020; + let ge_a = (lowered | HI).wrapping_sub(0x6161_6161_6161_6161); + let le_z = 0xFAFA_FAFA_FAFA_FAFA_u64.wrapping_sub(lowered); + !(ge_a & le_z) & HI +} + +/// SWAR letter scan: advances `pos` past ASCII letters. +/// Returns the updated pos. +#[inline(always)] +pub(crate) fn swar_scan_letters(bytes: &[u8], mut pos: usize) -> usize { + let len = bytes.len(); + // SWAR: 8 bytes at a time + while pos + 8 <= len { + let word = unsafe { (bytes.as_ptr().add(pos) as *const u64).read_unaligned() }; + if word & HI != 0 { + break; + } + let nonletter = swar64_letter_nonmask(word); + if nonletter != 0 { + return pos + nonletter.to_le().trailing_zeros() as usize / 8; + } + pos += 8; + } + // Scalar tail + while pos < len { + let b = unsafe { *bytes.get_unchecked(pos) }; + if is_letter(b) { + pos += 1; + } else { + break; + } + } + pos +} + +/// NEON letter scan: 16 bytes per iteration. Non-ASCII bytes (>= 0x80) fail +/// the `<= 'z'` check after case-folding, so they stop the run exactly like +/// non-letters; the caller's unicode continuation handles them. +/// +/// NOT used by `scan_letters_from`: measured 0.83x of the SWAR scan on OWT. +/// The `vshrn`-based movemask needs a vector→GPR transfer whose latency sits +/// on the serial per-token chain, and typical letter runs (~4-6 bytes) fit in +/// one SWAR iteration anyway. Kept as a reference / benchmark baseline. +#[cfg(target_arch = "aarch64")] +#[allow(dead_code)] +#[inline(always)] +pub(crate) fn neon_scan_letters(bytes: &[u8], mut pos: usize) -> usize { + use std::arch::aarch64::*; + let len = bytes.len(); + while pos + 16 <= len { + unsafe { + let v = vld1q_u8(bytes.as_ptr().add(pos)); + let lowered = vorrq_u8(v, vdupq_n_u8(0x20)); + let ge_a = vcgeq_u8(lowered, vdupq_n_u8(b'a')); + let le_z = vcleq_u8(lowered, vdupq_n_u8(b'z')); + let nonletter = vmvnq_u8(vandq_u8(ge_a, le_z)); + // Narrowing movemask: 4 bits per lane, first set nibble = first + // non-letter lane. + let mask = vget_lane_u64::<0>(vreinterpret_u64_u8(vshrn_n_u16::<4>( + vreinterpretq_u16_u8(nonletter), + ))); + if mask != 0 { + return pos + (mask.trailing_zeros() >> 2) as usize; + } + } + pos += 16; + } + while pos < len { + let b = unsafe { *bytes.get_unchecked(pos) }; + if is_letter(b) { + pos += 1; + } else { + break; + } + } + pos +} + +// ----------------------------------------------------------------------- +// Shared run scans (`\p{L}+`, `\p{N}+`, `\p{N}{1,3}`, `[^\s\p{L}\p{N}]+`) +// ----------------------------------------------------------------------- + +/// `\p{N}{1,3}`: extend a number run that already matched `consumed` chars +/// to at most 3 chars total. Shared by the cl100k and olmo3 schemes. +#[inline(always)] +pub(crate) fn scan_numbers_max3(bytes: &[u8], mut pos: usize, mut consumed: u32) -> usize { + let len = bytes.len(); + while consumed < 3 && pos < len { + let b = unsafe { *bytes.get_unchecked(pos) }; + if is_digit(b) { + pos += 1; + consumed += 1; + continue; + } + if b >= 0x80 { + let (cp, l) = unsafe { decode_cp(bytes, pos) }; + if unicode::class_of(cp) == unicode::CharClass::Number { + pos += l; + consumed += 1; + continue; + } + } + break; + } + pos +} + +#[inline(always)] +pub(crate) fn scan_letters_from(bytes: &[u8], pos: usize) -> usize { + let len = bytes.len(); + let mut p = pos; + loop { + p = swar_scan_letters(bytes, p); + if p < len && unsafe { *bytes.get_unchecked(p) } >= 0x80 { + let (cp, l) = unsafe { decode_cp(bytes, p) }; + if unicode::class_of(cp) == unicode::CharClass::Letter { + p += l; + continue; + } + } + return p; + } +} + +#[inline(always)] +pub(crate) fn scan_digits_from(bytes: &[u8], pos: usize) -> usize { + let len = bytes.len(); + let mut p = pos; + loop { + while p < len && is_digit(unsafe { *bytes.get_unchecked(p) }) { + p += 1; + } + if p < len && unsafe { *bytes.get_unchecked(p) } >= 0x80 { + let (cp, l) = unsafe { decode_cp(bytes, p) }; + if unicode::class_of(cp) == unicode::CharClass::Number { + p += l; + continue; + } + } + return p; + } +} + +#[inline(always)] +pub(crate) fn scan_other_from(bytes: &[u8], pos: usize) -> usize { + let len = bytes.len(); + let mut p = pos; + loop { + while p < len { + let b = unsafe { *bytes.get_unchecked(p) }; + if b >= 0x80 { break; } + if is_letter(b) || is_digit(b) || is_ascii_ws(b) { return p; } + p += 1; + } + if p < len { + let (cp, l) = unsafe { decode_cp(bytes, p) }; + if unicode::class_of(cp) == unicode::CharClass::Other { + p += l; + continue; + } + } + return p; + } +} + diff --git a/fast/src/pretok/o200k.rs b/fast/src/pretok/o200k.rs new file mode 100644 index 0000000..927f388 --- /dev/null +++ b/fast/src/pretok/o200k.rs @@ -0,0 +1,73 @@ +// Vendored from gigatoken (https://github.com/marcelroed/gigatoken), +// rev 542367a3efed134883fb4f1140b49c04e6fad3a3, MIT license. +// See src/pretok/mod.rs for what was trimmed and why. + +//! Fast pretokenizer for the o200k_base regex (GPT-4o, gpt-oss): +//! `[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n/]*|\s*[\r\n]+|\s+(?!\S)|\s+` +//! +//! See `o200k_family` for the shared scalar walker and mask-scanner +//! boundary algebra (`CONTRACTIONS = true`, `DIGITS3 = true`). + +use super::mask::{MaskScheme, MaskState}; +use super::o200k_family; +use super::Pretoken; + +pub(crate) struct O200kScheme; + +impl MaskScheme for O200kScheme { + #[inline(always)] + fn advance(bytes: &[u8], pos: usize) -> usize { + o200k_family::advance_pos::(bytes, pos) + } + + #[cfg(target_arch = "aarch64")] + #[inline(always)] + fn batch_masks(bytes: &[u8], scan: usize) -> (u64, u64) { + o200k_family::batch_masks::(bytes, scan) + } + + #[cfg(target_arch = "x86_64")] + #[inline(always)] + unsafe fn batch_masks_x86(bytes: &[u8], scan: usize) -> (u64, u64) { + // SAFETY: the caller detected the tier (trait contract). + unsafe { o200k_family::batch_masks_x86::(bytes, scan) } + } +} + +/// With SIMD support (aarch64 NEON, or x86_64 AVX-512/AVX2 detected at +/// runtime), iteration runs the shared o200k-family mask scanner (see +/// `o200k_family::batch_masks`); elsewhere every token takes the scalar +/// `advance_pos`. +pub struct FastO200kPretokenizer<'a> { + bytes: &'a [u8], + state: MaskState, +} + +impl<'a> FastO200kPretokenizer<'a> { + #[inline] + pub fn new(bytes: &'a [u8]) -> Self { + Self::with_pos(bytes, 0) + } + + /// Resume iteration at a byte offset previously returned by [`Self::pos`]. + #[inline] + pub fn with_pos(bytes: &'a [u8], pos: usize) -> Self { + Self { bytes, state: MaskState::new(pos) } + } + + /// Current position as a byte offset into the input. + #[inline] + pub fn pos(&self) -> usize { + self.state.pos + } +} + +impl<'a> Iterator for FastO200kPretokenizer<'a> { + type Item = Pretoken<'a>; + + #[inline] + fn next(&mut self) -> Option> { + let (start, end) = self.state.next_span::(self.bytes)?; + Some(Pretoken(&self.bytes[start..end])) + } +} diff --git a/fast/src/pretok/o200k_family.rs b/fast/src/pretok/o200k_family.rs new file mode 100644 index 0000000..e47c96c --- /dev/null +++ b/fast/src/pretok/o200k_family.rs @@ -0,0 +1,1495 @@ +// Vendored from gigatoken (https://github.com/marcelroed/gigatoken), +// rev 542367a3efed134883fb4f1140b49c04e6fad3a3, MIT license. +// See src/pretok/mod.rs for what was trimmed and why. + +//! Shared implementation for the o200k regex family: o200k_base (GPT-4o, +//! gpt-oss), the Nemotron-3 variant, and the Kimi (moonshotai K2) variant. +//! Their patterns share the shape +//! +//! `HAN-RUN?| +//! [^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+SUF?| +//! [^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*SUF?| +//! \p{N}{1,3} or \p{N}| ?[^\s\p{L}\p{N}]+TAIL|\s*[\r\n]+|\s+(?!\S)|\s+` +//! +//! where `SUF = (?i:'s|'t|'re|'ve|'m|'ll|'d)` exists in o200k and Kimi +//! (`CONTRACTIONS`), the digit group is `\p{N}{1,3}` for o200k/Kimi vs +//! `\p{N}` for Nemotron (`DIGITS3`), the absorbed punct tail `TAIL` is +//! `[\r\n/]*` for o200k/Nemotron vs `[\r\n]*` for Kimi (`SLASH`), and Kimi +//! alone (`HAN`) prepends a `[\p{Han}]+` alternative while intersecting +//! both letter brackets with `[^\p{Han}]` — Han chars form their own runs +//! and never join letter runs, though a Han numeral (Nl) still counts +//! toward a `\p{N}{1,3}` group entered on a non-Han digit and a Han symbol +//! (So/Mc) still continues a `[^\s\p{L}\p{N}]+` punct run (the Han +//! alternative only wins where a token starts; see +//! [`crate::pretokenize::unicode::KimiCharClass`]). +//! +//! Differences from the cl100k family: +//! - Letter runs are case-structured. Under leftmost-greedy backtracking +//! the two letter alternatives reduce to a phase automaton (see +//! [`CaseState`]): ULC phase until the first strict-lower (Ll), then +//! LLC phase, where a strict-upper (Lu/Lt) ends the token; a run ending +//! while still in ULC phase backtracks to its last caseless/mark char +//! ("camelCase" -> `camel|Case`, "HTTPResponse" one token, +//! "AxxB" -> `Axx|B` for caseless x). For pure-ASCII text there are no +//! caseless letters, so the rule IS pairwise: a boundary sits before +//! `[A-Z]` exactly when the previous char is `[a-z]` — what the ASCII +//! mask algebra uses; caseless-before-upper needs the phase and +//! lookahead, so the extended path defers those (rare) chars. +//! - Contractions are attached suffixes of the letter alternatives, not a +//! standalone alternative: "don't" is ONE token, and the char after a +//! consumed suffix always starts a new token ("can'ts" -> `can't|s`). +//! A contraction applies only when the apostrophe directly follows a +//! letter-run char; elsewhere `'` is ordinary punctuation (which may +//! still prefix a letter run: "3'ts" -> `3|'ts`). +//! - Punctuation runs absorb a `[\r\n/]*` tail. Since `/` is itself in +//! the punct class, the absorbed tail always begins with a newline: +//! ".\n//" is one token. +//! - Marks (`\p{M}`) are dual-class: they join letter runs (they sit in +//! both letter brackets) AND continue `[^\s\p{L}\p{N}]+` punct runs. +//! Their effective class is run-contextual, so the mask scanner routes +//! mark chars (rare) through the scalar path as bad zones. +//! +//! The boundary algebra below mirrors `cl100k_family` — see that module +//! and pretokenizer_optimization_log.md step 16 for the base rules. + +#[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))] +use super::mask::{self, AsciiMasks}; +use super::{decode_cp, is_ascii_ws, is_digit, is_letter, scan_numbers_max3}; +use super::unicode::{O200kCharClass, kimi_class_of, o200k_class_of}; + +#[inline(always)] +fn is_upper_ascii(b: u8) -> bool { + b.wrapping_sub(b'A') < 26 +} + +/// Is this byte in the absorbed-tail class (`[\r\n/]` or `[\r\n]`)? +#[inline(always)] +fn is_tail_byte(b: u8) -> bool { + b == b'\r' || b == b'\n' || (SLASH && b == b'/') +} + +/// Classify `cp` for the scheme: its effective o200k class plus whether it +/// is a `[\p{Han}]+` run member (always false off the Kimi scheme; one +/// table load either way). +#[inline(always)] +fn family_class_of(cp: u32) -> (O200kCharClass, bool) { + if HAN { + let k = kimi_class_of(cp); + (k.base(), k.is_han()) + } else { + (o200k_class_of(cp), false) + } +} + +// ----------------------------------------------------------------------- +// Scalar ground truth +// ----------------------------------------------------------------------- + +/// Scan state of a case-structured letter run, mirroring the two-bracket +/// alternatives under leftmost-greedy backtracking. `U`: still inside +/// `[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*` (no strict-lower seen yet); +/// `last_cl_end` is the end offset of the last caseless/mark char in the +/// phase (0 = none) — the backtrack split point if the run ends on a +/// strict-upper tail. `L`: inside `[\p{Ll}\p{Lm}\p{Lo}\p{M}]+`, where a +/// strict-upper char ends the token unconditionally. +#[derive(Clone, Copy)] +enum CaseState { + U { last_cl_end: usize }, + L, +} + +/// Initial [`CaseState`] for a run starting on an ASCII letter (ASCII has +/// no caseless letters). +#[inline(always)] +fn ascii_letter_state(b: u8) -> CaseState { + if is_upper_ascii(b) { + CaseState::U { last_cl_end: 0 } + } else { + CaseState::L + } +} + +/// If the char at `pos` is a letter-run member (`\p{L}` or `\p{M}`, minus +/// Han under the Kimi scheme), return (offset past it, initial scan state). +#[inline(always)] +fn letter_run_first(bytes: &[u8], pos: usize) -> Option<(usize, CaseState)> { + let &b = bytes.get(pos)?; + if is_letter(b) { + return Some((pos + 1, ascii_letter_state(b))); + } + if b >= 0x80 { + let (cp, l) = unsafe { decode_cp(bytes, pos) }; + match family_class_of::(cp) { + (_, true) => {} + (O200kCharClass::Upper, _) => return Some((pos + l, CaseState::U { last_cl_end: 0 })), + (O200kCharClass::Lower, _) => return Some((pos + l, CaseState::L)), + (O200kCharClass::Caseless | O200kCharClass::Mark, _) => { + return Some((pos + l, CaseState::U { last_cl_end: pos + l })); + } + _ => {} + } + } + None +} + +/// Letter-run continuation with the o200k casing rules: scan run members +/// (letters and marks) from `pos` with the phase automaton described on +/// [`CaseState`]. In phase U a strict-upper char continues the token and +/// a strict-lower switches to phase L; in phase L a strict-upper ends +/// the token. A run that ends while still in phase U backtracks to the +/// last caseless/mark char, splitting off the trailing strict-upper run +/// (which the next `advance` then consumes whole): "AxxB" -> `Axx|B` +/// for caseless x, "HTTPResponse" and "Z\u{5BF}\u{416}dz" one token. +#[inline(always)] +fn scan_case_run(bytes: &[u8], mut pos: usize, mut st: CaseState) -> usize { + let len = bytes.len(); + loop { + while pos < len { + let b = unsafe { *bytes.get_unchecked(pos) }; + if is_upper_ascii(b) { + if matches!(st, CaseState::L) { + return pos; + } + pos += 1; + } else if is_letter(b) { + st = CaseState::L; + pos += 1; + } else { + break; + } + } + if pos < len && unsafe { *bytes.get_unchecked(pos) } >= 0x80 { + let (cp, l) = unsafe { decode_cp(bytes, pos) }; + match family_class_of::(cp) { + (_, true) => break, + (O200kCharClass::Upper, _) => { + if matches!(st, CaseState::L) { + return pos; + } + pos += l; + } + (O200kCharClass::Lower, _) => { + st = CaseState::L; + pos += l; + } + (O200kCharClass::Caseless | O200kCharClass::Mark, _) => { + pos += l; + if let CaseState::U { ref mut last_cl_end } = st { + *last_cl_end = pos; + } + } + _ => break, + } + continue; + } + break; + } + // End of the letter run. A phase-U run ending on a strict-upper tail + // backtracks to the last caseless/mark char (`[LLC]+` must consume at + // least one char, and only a caseless/mark can be given back). + match st { + CaseState::U { last_cl_end } if last_cl_end != 0 => last_cl_end, + _ => pos, + } +} + +/// Attach a `(?i:'s|'t|'re|'ve|'m|'ll|'d)?` suffix to a letter token +/// ending at `end`, when the scheme has contractions. +#[inline(always)] +fn try_suffix(bytes: &[u8], end: usize) -> usize { + if !CONTRACTIONS || bytes.get(end) != Some(&b'\'') { + return end; + } + match bytes.get(end + 1).map(u8::to_ascii_lowercase) { + Some(b's' | b'd' | b'm' | b't') => end + 2, + Some(b'l') if bytes.get(end + 2).map(u8::to_ascii_lowercase) == Some(b'l') => end + 3, + Some(b'v') if bytes.get(end + 2).map(u8::to_ascii_lowercase) == Some(b'e') => end + 3, + Some(b'r') if bytes.get(end + 2).map(u8::to_ascii_lowercase) == Some(b'e') => end + 3, + // U+017F LATIN SMALL LETTER LONG S case-folds to 's' under `(?i)` + Some(0xC5) if bytes.get(end + 2) == Some(&0xBF) => end + 3, + _ => end, + } +} + +/// `[^\s\p{L}\p{N}]+` from `pos` (punctuation, symbols, marks, controls — +/// everything except letters, numbers, and whitespace; a Han symbol's +/// effective class is Other, so it continues the run under Kimi too). +#[inline(always)] +fn scan_punct_from(bytes: &[u8], pos: usize) -> usize { + let len = bytes.len(); + let mut p = pos; + loop { + while p < len { + let b = unsafe { *bytes.get_unchecked(p) }; + if b >= 0x80 { + break; + } + if is_letter(b) || is_digit(b) || is_ascii_ws(b) { + return p; + } + p += 1; + } + if p < len { + let (cp, l) = unsafe { decode_cp(bytes, p) }; + if matches!( + family_class_of::(cp).0, + O200kCharClass::Other | O200kCharClass::Mark + ) { + p += l; + continue; + } + } + return p; + } +} + +/// The `[\r\n/]*` (or Kimi `[\r\n]*`) tail absorbed after a punct run. +#[inline(always)] +fn scan_tail(bytes: &[u8], mut pos: usize) -> usize { + while pos < bytes.len() && is_tail_byte::(unsafe { *bytes.get_unchecked(pos) }) { + pos += 1; + } + pos +} + +/// `[\p{Han}]+` from `pos` (chars of any Han class — letters, numerals, +/// symbols; the leading alternative consumes the maximal Han run). +#[inline(always)] +fn scan_han_run(bytes: &[u8], mut pos: usize) -> usize { + while pos < bytes.len() { + let b = unsafe { *bytes.get_unchecked(pos) }; + if b < 0x80 { + return pos; + } + let (cp, l) = unsafe { decode_cp(bytes, pos) }; + if !kimi_class_of(cp).is_han() { + return pos; + } + pos += l; + } + pos +} + +/// Whitespace-led token starting at `start`: `\s*[\r\n]+` | `\s+(?!\S)` | +/// `\s+`, in that priority. Precondition: the letter-prefix and +/// space+punct alternatives were ruled out. +#[inline(always)] +fn ws_token_end(bytes: &[u8], start: usize) -> usize { + let len = bytes.len(); + let mut p = start; + let mut last_nl_end = 0usize; // 0 = run contains no \r\n + let mut last_char_start = start; + while p < len { + let b = unsafe { *bytes.get_unchecked(p) }; + if b == b'\r' || b == b'\n' { + last_char_start = p; + p += 1; + last_nl_end = p; + } else if is_ascii_ws(b) { + last_char_start = p; + p += 1; + } else if b >= 0x80 { + let (cp, l) = unsafe { decode_cp(bytes, p) }; + if o200k_class_of(cp) == O200kCharClass::Whitespace { + last_char_start = p; + p += l; + } else { + break; + } + } else { + break; + } + } + if last_nl_end != 0 { + return last_nl_end; // `\s*[\r\n]+`: through the last newline + } + if p >= len { + return p; // `\s+(?!\S)`: lookahead succeeds at EOS + } + if last_char_start > start { + return last_char_start; // `\s+(?!\S)`: all but the last ws char + } + p // `\s+`: single whitespace char before content +} + +/// `\p{N}{1,3}` or `\p{N}` starting at a digit char ending at `first_end`. +#[inline(always)] +fn digit_token_end(bytes: &[u8], first_end: usize) -> usize { + if DIGITS3 { + scan_numbers_max3(bytes, first_end, 1) + } else { + first_end + } +} + +/// Advance past one token starting at `pos`. Returns the new position. +/// `pos` must be < `bytes.len()`. +#[inline(always)] +pub(crate) fn advance_pos< + const CONTRACTIONS: bool, + const DIGITS3: bool, + const SLASH: bool, + const HAN: bool, +>( + bytes: &[u8], + pos: usize, +) -> usize { + let b0 = unsafe { *bytes.get_unchecked(pos) }; + + // Hot path 1: ASCII letter run (empty prefix) + if is_letter(b0) { + let e = scan_case_run::(bytes, pos + 1, ascii_letter_state(b0)); + return try_suffix::(bytes, e); + } + + // Hot path 2: space prefix + if b0 == b' ' { + let Some(&b1) = bytes.get(pos + 1) else { + return pos + 1; // trailing lone space (`\s+(?!\S)` at EOS) + }; + if is_letter(b1) { + let e = scan_case_run::(bytes, pos + 2, ascii_letter_state(b1)); + return try_suffix::(bytes, e); + } + if b1 < 0x80 { + if is_digit(b1) { + return pos + 1; // numbers never absorb the space + } + if is_ascii_ws(b1) { + return ws_token_end(bytes, pos); + } + // ` ?[^\s\p{L}\p{N}]+` + tail + let p = scan_punct_from::(bytes, pos + 2); + return scan_tail::(bytes, p); + } + let (cp, l) = unsafe { decode_cp(bytes, pos + 1) }; + let p1 = pos + 1 + l; + return match family_class_of::(cp) { + // A Han letter can neither join a letter run nor (being \p{L}) + // extend ` ?[^\s\p{L}\p{N}]+`: the space is a lone `\s+` token + // and the Han run starts after it. (Han symbols fall to the + // Other arm — they are ordinary punct-run members here — and + // Han numerals to the Number arm.) + (O200kCharClass::Caseless, true) => pos + 1, + (O200kCharClass::Upper, _) => try_suffix::( + bytes, + scan_case_run::(bytes, p1, CaseState::U { last_cl_end: 0 }), + ), + (O200kCharClass::Lower, _) => { + try_suffix::(bytes, scan_case_run::(bytes, p1, CaseState::L)) + } + (O200kCharClass::Caseless | O200kCharClass::Mark, _) => try_suffix::( + bytes, + scan_case_run::(bytes, p1, CaseState::U { last_cl_end: p1 }), + ), + (O200kCharClass::Whitespace, _) => ws_token_end(bytes, pos), + (O200kCharClass::Number, _) => pos + 1, + (O200kCharClass::Other, _) => { + scan_tail::(bytes, scan_punct_from::(bytes, p1)) + } + }; + } + + // Non-ASCII first char + if b0 >= 0x80 { + let (cp, l) = unsafe { decode_cp(bytes, pos) }; + let p0 = pos + l; + let (class, han) = family_class_of::(cp); + // `[\p{Han}]+` is the leading alternative: any Han char (whatever + // its base class) starting a token starts a Han run. + if HAN && han { + return scan_han_run(bytes, p0); + } + return match class { + O200kCharClass::Upper => try_suffix::( + bytes, + scan_case_run::(bytes, p0, CaseState::U { last_cl_end: 0 }), + ), + O200kCharClass::Lower => { + try_suffix::(bytes, scan_case_run::(bytes, p0, CaseState::L)) + } + O200kCharClass::Caseless | O200kCharClass::Mark => try_suffix::( + bytes, + scan_case_run::(bytes, p0, CaseState::U { last_cl_end: p0 }), + ), + O200kCharClass::Number => digit_token_end::(bytes, p0), + // Any non-letter/number char except \r\n may prefix a run + class => { + if let Some((e, st)) = letter_run_first::(bytes, p0) { + return try_suffix::(bytes, scan_case_run::(bytes, e, st)); + } + if class == O200kCharClass::Whitespace { + ws_token_end(bytes, pos) + } else { + scan_tail::(bytes, scan_punct_from::(bytes, p0)) + } + } + }; + } + + // ASCII digit + if is_digit(b0) { + return digit_token_end::(bytes, pos + 1); + } + + // \r and \n are excluded from the letter-run prefix + if b0 == b'\r' || b0 == b'\n' { + return ws_token_end(bytes, pos); + } + + // Other ASCII whitespace (\t, \x0b, \x0c) may prefix a letter run + if is_ascii_ws(b0) { + if let Some((e, st)) = letter_run_first::(bytes, pos + 1) { + return try_suffix::(bytes, scan_case_run::(bytes, e, st)); + } + return ws_token_end(bytes, pos); + } + + // ASCII punctuation/symbol/control (including `'`: o200k has no + // standalone contraction alternative, so a leading apostrophe is + // ordinary punctuation / a letter-run prefix: "'sound" is one token) + if let Some((e, st)) = letter_run_first::(bytes, pos + 1) { + return try_suffix::(bytes, scan_case_run::(bytes, e, st)); + } + scan_tail::(bytes, scan_punct_from::(bytes, pos + 1)) +} + +// ----------------------------------------------------------------------- +// Mask-scanner boundary algebra +// ----------------------------------------------------------------------- + +/// Smear `seed` upward (toward higher bits) through contiguous set bits of +/// `within`, in log steps. +#[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))] +#[inline(always)] +fn smear_up(seed: u64, within: u64) -> u64 { + let mut a = seed; + let mut m = within; + let mut sh = 1u32; + while sh < 64 { + a |= (a << sh) & m; + m &= m << sh; + sh <<= 1; + } + a +} + +/// Per-byte class masks for a batch's unicode chars under the o200k +/// classifier — the o200k analogue of [`mask::UniClasses`], with the +/// case-split letter masks the scheme needs. Every byte of a classified +/// char carries the char's class, so byte-adjacency == char-adjacency. +#[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))] +#[derive(Clone, Copy, Default)] +struct OUni { + /// All letter-run bytes (upper + lower + caseless; marks excluded — + /// they are contextual and deferred via `mk`). + l: u64, + /// Strict-upper (Lu/Lt) bytes. + u: u64, + /// Caseless-letter (Lm/Lo) bytes. + cl: u64, + /// Number / punct-run / whitespace bytes. + n: u64, + o: u64, + ws: u64, + /// Whitespace lead bits by char length (for `(?!\S)` shift tests). + w2: u64, + w3: u64, + /// Lead bits of all classified chars by length (two-chars-back rule). + lead2: u64, + lead3: u64, + lead4: u64, + /// Continuation bytes of classified chars. + cont: u64, + /// Bytes only the scalar path can decide (±1 bad smear): number chars + /// (char-counted grouping), ws straddling the batch end, stray + /// continuation bytes. + resid: u64, + /// Mark bytes (±4 bad smear). A mark's run-contextual class can + /// affect boundaries up to two CHARS after it, which multi-byte + /// followers can push past the 4-byte smear; those stragglers are + /// wrongly-cleared bits (extending the scalar walk) or wrongly-set + /// bits interior to a token starting inside the zone, both killed by + /// MaskState's resume masking after the scalar overrun — the same + /// invariant the resid zones rely on. Kimi routes Han symbols (So/Mc) + /// here too: punct-run members mid-run, Han-run chars at a token + /// start, so their class is run-contextual exactly like a mark's. + mk: u64, + /// Han-letter bytes (Kimi only): their own run class, with the + /// boundary rule `han_leads & !prev-is-han`. Han numerals go through + /// `resid` (their `\p{N}{1,3}`-vs-Han-run role is contextual) and Han + /// symbols through `mk`. + han: u64, + /// Lead bits of the `han` chars. + han_leads: u64, +} + +/// Boundary carries from the chars before the batch (o200k variant of the +/// cl100k family's `Carries`, plus the case and absorbed-tail bits). +#[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))] +#[derive(Clone, Copy, Default)] +struct OCarries { + /// P1 is a letter / strict-upper / caseless / space (0x20) / + /// non-newline non-space ws / punct / any ws / digit. + pl: u64, + pu: u64, + pcl: u64, + ps: u64, + pwt: u64, + po: u64, + pws: u64, + pd: u64, + /// P1 is a Han letter (Kimi only). + phan: u64, + /// P2 is punct-or-space, for a char lead at bit 0. + c2_os: u64, + /// Same test positioned at the first lead AFTER a straddling-in P1. + b2b_in: u64, + /// P1 is an absorbed `[\r\n/]*` tail byte whose token may continue + /// into this batch (seeds the tail smear at bit 0). + p_abs: bool, + /// The tail walkback could not resolve (pathological run): the batch's + /// leading tail-class run must be a bad zone. + force_bad_lead: bool, +} + +/// Was the tail-class byte at `scan - 1` absorbed by a punct run's +/// `[\r\n/]*` (Kimi: `[\r\n]*`) tail (as opposed to being a fresh +/// punct-run `/` or a ws-run newline)? Walks the tail-class run back +/// (bounded) and classifies the byte before it. `None`: unresolved +/// (over-long run, or a preceding mark or Han symbol whose own class is +/// run-contextual). +#[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))] +fn prev_tail_absorbed( + bytes: &[u8], + scan: usize, +) -> Option { + debug_assert!(scan >= 1 && is_tail_byte::(bytes[scan - 1])); + let mut r = scan - 1; + let mut steps = 0; + while r > 0 && is_tail_byte::(bytes[r - 1]) { + r -= 1; + steps += 1; + if steps > 8 { + return None; + } + } + // T-run = bytes[r..scan]. The `[\r\n/]*` tail is greedy, so once + // absorption triggers — at the first newline that directly follows a + // punct-run char (an in-run slash, or the pre-run char for a + // run-leading newline) — everything to the run's end is absorbed. + // Before the trigger, newlines are ws-run members and slashes are + // ordinary punct-run bytes. + let run = &bytes[r..scan]; + let mut trigger = usize::MAX; + let mut seen_slash = false; + for (j, &b) in run.iter().enumerate() { + if b == b'/' { + seen_slash = true; + continue; + } + if seen_slash { + trigger = j; + break; + } + if j == 0 { + // Run-leading newline: the pre-run char decides. + if r == 0 { + continue; + } + let pb = bytes[r - 1]; + let pred_punct = if pb < 0x80 { + if !is_letter(pb) && !is_digit(pb) && !is_ascii_ws(pb) { + Some(true) + } else { + Some(false) + } + } else { + let mut k = r - 1; + while k > 0 && bytes[k] & 0xC0 == 0x80 { + k -= 1; + } + let (cp, _) = unsafe { decode_cp(bytes, k) }; + match family_class_of::(cp) { + (O200kCharClass::Other, false) => Some(true), + // A mark continues whatever run precedes it; a Han + // symbol is a punct-run member or a Han-run char + // depending on where its token started. + (O200kCharClass::Mark, _) | (O200kCharClass::Other, true) => None, + _ => Some(false), + } + }; + match pred_punct { + Some(true) => { + trigger = 0; + break; + } + Some(false) => {} + None => return None, + } + } + } + Some(scan - 1 - r >= trigger) +} + +/// Two-back "punct or space" test (`c2_os`) for the ASCII byte at `idx`. +/// A slash may be an absorbed tail byte — a token end, neither punct-run +/// member nor space — so it resolves through the walkback (only when the +/// scheme absorbs slashes). `None`: unresolved (callers set +/// `force_bad_lead`). +#[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))] +#[inline(always)] +fn c2_os_ascii(bytes: &[u8], idx: usize) -> Option { + let b2 = bytes[idx]; + if SLASH && b2 == b'/' { + return prev_tail_absorbed::(bytes, idx + 1).map(|abs| u64::from(!abs)); + } + Some(u64::from( + b2 == b' ' || (!is_letter(b2) && !is_digit(b2) && !is_ascii_ws(b2)), + )) +} + +/// Pure-ASCII carries. Requires `scan > 0`, `bytes[scan-1] < 0x80` (and +/// `bytes[scan-2] < 0x80` when present), and `bytes[scan-1]` NOT a +/// tail-class byte (those route through [`prev_tail_absorbed`] first). +#[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))] +#[inline(always)] +fn ascii_carries(bytes: &[u8], scan: usize) -> OCarries { + let b = bytes[scan - 1]; + debug_assert!(!is_tail_byte::(b)); + let bit = |c: bool| u64::from(c); + let (l, d, w) = (is_letter(b), is_digit(b), is_ascii_ws(b)); + let (c2_os, c2_unresolved) = if scan >= 2 { + match c2_os_ascii::(bytes, scan - 2) { + Some(v) => (v, false), + None => (0, true), + } + } else { + (0, false) + }; + OCarries { + force_bad_lead: c2_unresolved, + pl: bit(l), + pu: bit(is_upper_ascii(b)), + pcl: 0, + ps: bit(b == b' '), + pwt: bit(w && b != b' '), // \r\n excluded by the debug_assert + po: bit(!l && !d && !w), + pws: bit(w), + pd: bit(d), + phan: 0, + c2_os, + b2b_in: 0, + p_abs: false, + } +} + +/// Carries when the byte before the batch is tail-class (`[\r\n/]`): +/// resolves absorbed-tail vs fresh-run via [`prev_tail_absorbed`]. An +/// absorbed tail ended the previous token, so every "P1 is X" carry is +/// zero and only the tail-continuation seed survives. +#[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))] +#[inline(never)] +fn tail_carries(bytes: &[u8], scan: usize) -> OCarries { + match prev_tail_absorbed::(bytes, scan) { + Some(true) => OCarries { p_abs: true, ..OCarries::default() }, + Some(false) => { + let b = bytes[scan - 1]; + let bit = |c: bool| u64::from(c); + // A fresh '/' is an ordinary punct byte; \r\n are newlines + // (ws class, po = 0). c2_os as in ascii_carries when P2 is + // ASCII; a non-ASCII P2 next to a tail byte is rare — defer. + if scan >= 2 && bytes[scan - 2] >= 0x80 { + return OCarries { + force_bad_lead: true, + ..OCarries::default() + }; + } + let c2_os = if scan >= 2 { + let b2 = bytes[scan - 2]; + bit(b2 == b' ' || (!is_letter(b2) && !is_digit(b2) && !is_ascii_ws(b2))) + } else { + 0 + }; + OCarries { + po: bit(b == b'/'), + pws: bit(b != b'/'), + c2_os, + ..OCarries::default() + } + } + None => OCarries { + force_bad_lead: true, + ..OCarries::default() + }, + } +} + +/// Classify every unicode char whose lead bit is in `m` for +/// `bytes[scan..scan+64]` with the o200k classifier — the o200k analogue +/// of [`mask::classify_uni_chars`] (NUMBERS = false, LEADS = true), with +/// case-split letter masks and marks deferred via `mk`. +/// +/// # Safety +/// +/// `scan + 70 <= bytes.len()` (the batch classifiers' lookahead guard). +#[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))] +#[inline(always)] +unsafe fn classify_uni_o200k(bytes: &[u8], scan: usize, mut m: u64) -> OUni { + use super::decode_cp_inbounds; + let mut u = OUni::default(); + while m != 0 { + let i = m.trailing_zeros() as usize; + m &= m - 1; + let b = bytes[scan + i]; + if b < 0xC2 { + u.resid |= 1 << i; // stray continuation byte (invalid UTF-8) + continue; + } + let l = if b < 0xE0 { + 2 + } else if b < 0xF0 { + 3 + } else { + 4 + }; + let chm = ((1u64 << l) - 1) << i; // in-batch bytes (excess drops) + let lead = 1u64 << i; + // SAFETY: scan + 70 <= len (this fn's contract), i <= 63, so + // scan + i + 4 <= len even for a 4-byte lead at bit 63. + let (cp, _) = unsafe { decode_cp_inbounds(bytes, scan + i) }; + match family_class_of::(cp) { + // Han letters: their own run class (see OUni::han). Han + // numerals fall to the Number arm (whose resid already defers + // them: their Han-run-vs-digit-group role is contextual) and + // Han symbols to the Mark arm (same run-contextual deferral). + (O200kCharClass::Caseless, true) => { + u.han |= chm; + u.han_leads |= lead; + } + (O200kCharClass::Upper, _) => { + u.l |= chm; + u.u |= chm; + } + (O200kCharClass::Lower, _) => u.l |= chm, + (O200kCharClass::Caseless, _) => { + u.l |= chm; + u.cl |= chm; + } + (O200kCharClass::Mark, _) | (O200kCharClass::Other, true) => { + // Contextual (letter-run joiner AND punct-run member, or + // for Han symbols punct-run member AND Han-run char): + // punct-class for the neighbors' mask algebra, deferred + // (±4) for everything the context could change. + u.o |= chm; + u.mk |= chm; + } + (O200kCharClass::Number, _) => { + u.n |= chm; + u.resid |= chm; // char-counted grouping: scalar + } + (O200kCharClass::Other, _) => u.o |= chm, + (O200kCharClass::Whitespace, _) => { + u.ws |= chm; + if i + l > 64 || l == 4 { + u.resid |= chm; // straddling-out ws stays a bad zone + } else if l == 2 { + u.w2 |= lead; + } else { + u.w3 |= lead; + } + } + } + match l { + 2 => u.lead2 |= lead, + 3 => u.lead3 |= lead, + _ => u.lead4 |= lead, + } + u.cont |= chm & !lead; + m &= !chm; + } + u +} + +/// ASCII class masks the o200k algebra needs on top of +/// [`mask::AsciiMasks`]: strict-uppercase letters and slashes. +#[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))] +#[derive(Clone, Copy, Default)] +struct OAsciiExtra { + up: u64, + sl: u64, +} + +/// Slow(er) path for batches with non-ASCII in or just before them — the +/// o200k analogue of `cl100k_family::family_extended_masks`: carries walk +/// back through multi-byte chars with the o200k classifier, unicode chars +/// join the effective class masks, then the shared algebra applies. +#[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))] +#[cfg_attr( + target_arch = "x86_64", + target_feature(enable = "bmi1,bmi2,lzcnt,popcnt") +)] +#[inline(never)] +fn o200k_extended_masks< + const CONTRACTIONS: bool, + const DIGITS3: bool, + const SLASH: bool, + const HAN: bool, +>( + bytes: &[u8], + scan: usize, + am: AsciiMasks, + ax: OAsciiExtra, +) -> (u64, u64) { + use super::decode_cp_inbounds; + + /// The char containing byte `pos - 1`: its scheme class (plus Han + /// flag), lead index, and end (exclusive) — [`mask::char_through`] + /// with the scheme classifier. Same safety contract (`pos > 0`, + /// `pos + 3 <= len` when `bytes[pos-1]` is non-ASCII; the batch guard + /// covers callers). + #[inline(always)] + unsafe fn char_through_o200k( + bytes: &[u8], + pos: usize, + ) -> (O200kCharClass, bool, usize, usize) { + let b = bytes[pos - 1]; + if b < 0x80 { + let cls = if is_upper_ascii(b) { + O200kCharClass::Upper + } else if is_letter(b) { + O200kCharClass::Lower + } else if is_digit(b) { + O200kCharClass::Number + } else if is_ascii_ws(b) { + O200kCharClass::Whitespace + } else { + O200kCharClass::Other + }; + return (cls, false, pos - 1, pos); + } + let mut j = pos - 1; + while j > 0 && bytes[j] & 0xC0 == 0x80 { + j -= 1; + } + // SAFETY: j < pos and pos + 3 <= len (contract), so j + 4 <= len. + let (cp, l) = unsafe { decode_cp_inbounds(bytes, j) }; + let (cls, han) = family_class_of::(cp); + (cls, han, j, j + l) + } + + let mut cl = OUni::default(); + let cr = if scan == 0 { + OCarries::default() + } else if bytes[scan - 1] < 0x80 && is_tail_byte::(bytes[scan - 1]) { + tail_carries::(bytes, scan) + } else if bytes[scan - 1] < 0x80 && (scan < 2 || bytes[scan - 2] < 0x80) { + ascii_carries::(bytes, scan) + } else { + // A multi-byte char within two bytes of the batch start. + // SAFETY: scan > 0, and the batch guard covers pos + 3 <= len. + let (c1, h1, j1, e1) = unsafe { char_through_o200k::(bytes, scan) }; + let chm = if e1 > scan { (1u64 << (e1 - scan)) - 1 } else { 0 }; + cl.cont = chm; + let (c2v, c2_defer) = if j1 == 0 { + (0, false) + } else if SLASH && bytes[j1 - 1] == b'/' { + // An absorbed tail slash is a token end, not a punct-run char. + match prev_tail_absorbed::(bytes, j1) { + Some(abs) => (u64::from(!abs), false), + None => (0, true), + } + } else { + // SAFETY: j1 > 0, j1 < scan keeps the decode in the guard. + let (c2c, h2, _, _) = unsafe { char_through_o200k::(bytes, j1) }; + ( + u64::from( + bytes[j1 - 1] == b' ' + || (matches!(c2c, O200kCharClass::Other | O200kCharClass::Mark) && !h2), + ), + // A mark P2 makes the two-back test run-contextual; so + // does a Han symbol (punct-run member or Han-run char). + c2c == O200kCharClass::Mark || (h2 && c2c == O200kCharClass::Other), + ) + }; + let mut c = OCarries::default(); + if e1 > scan { + c.b2b_in = c2v << (e1 - scan); + } else { + c.c2_os = c2v; + } + if c2_defer { + c.force_bad_lead = true; + } + c.pd = u64::from(c1 == O200kCharClass::Number); + match (c1, h1) { + // A Han letter: p_han carry; its in-batch bytes join the han + // mask so in-batch followers see char adjacency. + (O200kCharClass::Caseless, true) => { + cl.han |= chm; + c.phan = 1; + } + (O200kCharClass::Upper, _) => { + cl.l |= chm; + cl.u |= chm; + c.pl = 1; + c.pu = 1; + } + (O200kCharClass::Lower, _) => { + cl.l |= chm; + c.pl = 1; + } + (O200kCharClass::Caseless, _) => { + cl.l |= chm; + cl.cl |= chm; + c.pl = 1; + c.pcl = 1; + } + (O200kCharClass::Mark, _) | (O200kCharClass::Other, true) => { + // Contextual (marks; Kimi Han symbols): defer the batch + // front to the scalar path. + cl.o |= chm; + cl.mk |= chm | 1; // bit 0 seeds the ±4 smear even when + // the char sits entirely before the batch + c.po = 1; + } + (O200kCharClass::Number, h) => { + cl.n |= chm; + // A digit char straddling INTO the batch: the leading + // ASCII digit run's `\p{N}{1,3}` phase started before the + // batch, and the `pd` seed below can't see it (bit 0 is a + // continuation byte, not an ASCII digit). Defer via resid + // so the bad<<1 seed catches the run. + cl.resid |= chm; + if h { + // A Han numeral P1: a following Han char's boundary + // depends on whether P1 sat in a digit group or a Han + // run — defer the batch front even when P1 lies + // entirely before the batch. + cl.resid |= 1; + } + } + (O200kCharClass::Other, _) => { + cl.o |= chm; + c.po = 1; + } + (O200kCharClass::Whitespace, _) => { + cl.ws |= chm; + if e1 > scan { + cl.resid |= chm; + } + let pb = bytes[scan - 1]; + c.ps = u64::from(pb == b' '); + let nl = pb == b'\r' || pb == b'\n'; + c.pwt = u64::from(pb < 0x80 && pb != b' ' && !nl || pb >= 0x80); + c.pws = 1; + } + } + c + }; + + let mut uni = if am.hi != 0 { + // SAFETY: the batch guard is exactly `classify_uni_o200k`'s + // contract. + unsafe { classify_uni_o200k::(bytes, scan, am.hi & !cl.cont) } + } else { + OUni::default() + }; + uni.l |= cl.l; + uni.u |= cl.u; + uni.cl |= cl.cl; + uni.n |= cl.n; + uni.o |= cl.o; + uni.ws |= cl.ws; + uni.cont |= cl.cont; + uni.resid |= cl.resid; + uni.mk |= cl.mk; + uni.han |= cl.han; + + o200k_algebra::(bytes, scan, am, ax, cr, uni) +} + +/// The o200k family's shared u64 boundary algebra over per-byte class +/// masks — `cl100k_family::family_algebra` with the o200k rules: casing +/// boundaries inside letter runs, suffix contractions, `[\r\n/]*` punct +/// tails. `uni` is all-zero on the pure-ASCII path. +#[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))] +#[inline(always)] +fn o200k_algebra< + const CONTRACTIONS: bool, + const DIGITS3: bool, + const SLASH: bool, + const HAN: bool, +>( + bytes: &[u8], + scan: usize, + am: AsciiMasks, + ax: OAsciiExtra, + cr: OCarries, + uni: OUni, +) -> (u64, u64) { + let OCarries { + pl, + pu, + pcl, + ps, + pwt, + po, + pws, + pd, + phan, + c2_os, + b2b_in, + p_abs, + force_bad_lead, + } = cr; + let contm = uni.cont; + + // Effective per-byte classes. + let lb = am.l | uni.l; + let ub = ax.up | uni.u; + let clb = uni.cl; + let sb = am.s; + let wtb = am.wt | uni.ws; + let ob = !(am.l | am.d | am.s | am.wt | am.n | am.hi) | uni.o; + let ws_all = sb | wtb | am.n; + + // --- Absorbed `[\r\n/]*` (Kimi `[\r\n]*`) tails ------------------------- + // Seed: a newline right after a punct byte (a slash after a punct run + // is already a run member, so tails always begin with a newline), or + // a tail continuing from before the batch. Smear through the tail + // class. + let tcls = if SLASH { am.n | ax.sl } else { am.n }; + let abs_seed = (am.n & ((ob << 1) | po)) | (u64::from(p_abs) & tcls); + let abs_t = if abs_seed == 0 { 0 } else { smear_up(abs_seed, tcls) }; + // Absorbed bytes are no longer punct-run members for any boundary rule. + let ob_eff = ob & !abs_t; + + // --- Letters ------------------------------------------------------------- + let len1 = !(contm | uni.lead2 | uni.lead3 | uni.lead4); + let c_test = ((ob_eff | sb) << 1) | po | ps; // bit 0: byte scan-1 in O|S + let b2back = ((c_test & len1) << 1) + | ((c_test & uni.lead2) << 2) + | ((c_test & uni.lead3) << 3) + | ((c_test & uni.lead4) << 4) + | c2_os + | b2b_in; + let p_l = (lb << 1) | pl; + let p_u = (ub << 1) | pu; + let p_cl = (clb << 1) | pcl; + let p_s = (sb << 1) | ps; + let p_wt = (wtb << 1) | pwt; + let p_o = (ob_eff << 1) | po; + let absorb = p_o & !b2back; + // Casing boundary: a strict-upper char after a strict-lower one. (For + // ASCII text this is the whole rule; see the module docs.) + let p_sl = p_l & !p_u & !p_cl; + let b_su = ub & !contm & p_sl; + let b_letters = + (lb & !contm & !p_l & !p_s & !p_wt & !absorb) | b_su; + + // --- Digits: `\p{N}{1,3}` or `\p{N}` ------------------------------------- + let b_digits = if DIGITS3 && am.d & (am.d >> 1) != 0 { + mask::digit_run_splits3(am.d) + } else { + am.d + }; + + // --- Punct: ` ?[^\s\p{L}\p{N}]+[\r\n/]*` ---------------------------------- + let b_punct = ob_eff & !contm & !p_o & !p_s; + + // --- Bad zones ------------------------------------------------------------ + let resid = uni.resid; + let mut bad = resid | resid << 1 | resid >> 1; + // Marks are run-contextual: they, and anything whose boundary rules + // can see them (up to two chars back — 4 bytes of lookahead for the + // following leads), go to the scalar path. + let mk = uni.mk; + if mk != 0 { + bad |= mk | mk << 1 | mk << 2 | mk << 3 | mk << 4 | mk >> 1; + } + // A strict-upper char after a caseless letter: phase- and + // lookahead-dependent (see the module docs) — scalar. + bad |= ub & !contm & ((clb << 1) | pcl); + if force_bad_lead { + // Unresolved carries: the leading tail-class run (plus the byte + // after it) can't be trusted. + bad |= smear_up(tcls & 1, tcls) << 1 | 0b11; + } + + // --- Whitespace ----------------------------------------------------------- + let ws_eff = ws_all & !abs_t; + + // Byte-64 lookahead: is the char at the next batch's first byte + // non-ws? (See cl100k_family for the full reasoning.) + let nb64 = bytes[scan + 64]; // in bounds: scan + 70 <= len + let nn64 = if nb64 < 0x80 { + !is_ascii_ws(nb64) + } else { + // SAFETY: the batch guard puts the decode at scan + 64 in bounds. + bad >> 63 == 0 && unsafe { mask::nn_at_full(bytes, scan + 64) } + }; + let nn64m = u64::from(nn64).wrapping_neg(); + + // An absorbed tail touching the batch end continues iff byte 64 is + // tail-class; the next batch's `tail_carries` walkback re-derives the + // context either way, so nothing defers here. A ws run touching the + // batch end still defers when byte 64 is ws (its last newline may lie + // beyond this batch). + let nonws = !ws_eff; + if ws_eff >> 63 != 0 && !nn64 { + if nonws == 0 { + return (0, u64::MAX); // whole batch one ws run + } + let h = 63 - nonws.leading_zeros(); // highest non-ws bit (< 63) + bad |= u64::MAX << (h + 1); + } + + // A digit run whose `\p{N}{1,3}` phase did not start inside this batch + // (continuation from before it, or after a bad zone) defers. + if DIGITS3 { + let seed = (am.d & (bad << 1)) | (am.d & pd); + if seed != 0 { + bad |= smear_up(seed, am.d); + } + } + + // Base rule (NL-free runs; NL runs are overridden below). + let ws_leads1 = (am.s | am.wt | am.n) & ws_eff; + let ws_leads = (ws_leads1 | uni.w2 | uni.w3) & !abs_t; + let p_ws = (ws_eff << 1) | pws; + let edge_last = (ws_leads1 & (1 << 63)) | (uni.w2 & (1 << 62)) | (uni.w3 & (1 << 61)); + let split_ok = (ws_leads1 & (nonws >> 1)) + | (uni.w2 & (nonws >> 2)) + | (uni.w3 & (nonws >> 3)) + | (edge_last & nn64m); + let mut b_ws = ws_leads & (!p_ws | split_ok); + + // Override every run containing a (non-absorbed) newline: one token + // through the run's last newline, then tail rules. + let mut runs_n = am.n & ws_eff & !bad; + while runs_n != 0 { + let f = runs_n.trailing_zeros(); + let below_gap = nonws & ((1u64 << f) - 1); + let a = if below_gap == 0 { 0 } else { 64 - below_gap.leading_zeros() }; + let e = (nonws & (u64::MAX << f)).trailing_zeros(); + let run_mask = (u64::MAX << a) & !u64::MAX.unbounded_shl(e); + b_ws &= !run_mask; + b_ws |= 1u64 << a; + let q = 63 - (am.n & run_mask).leading_zeros(); // last NL in run + if (q + 1) < e { + b_ws |= 1u64 << (q + 1); + let tail = run_mask & (u64::MAX << (q + 1)); + let tail_leads = ws_leads & tail; + b_ws |= 1u64 << (63 - tail_leads.leading_zeros()); + } + runs_n &= !run_mask; + } + + // --- Han runs (Kimi): `[\p{Han}]+` -------------------------------------- + // A boundary at every Han-letter lead whose previous char is not a Han + // letter. Han numerals/symbols sit in resid/mk bad zones, so any lead + // adjacent to those resolves on the scalar path. + let b_han = if HAN { + uni.han_leads & !((uni.han << 1) | phan) + } else { + 0 + }; + + let mut boundary = b_letters | b_digits | b_punct | b_ws | b_han; + + // --- Contractions: suffix `(?i:'s|'t|'re|'ve|'m|'ll|'d)?` ---------------- + // An apostrophe right after a letter-run char merges the suffix into + // that token and forces a boundary right after it. ('ſ is non-ASCII: + // an apostrophe before any non-ASCII char defers.) + if CONTRACTIONS { + let mut cand = am.ap & boundary & p_l & !bad; + let mut last_forced = usize::MAX; + while cand != 0 { + let i = cand.trailing_zeros() as usize; + cand &= cand - 1; + if i <= 2 { + // The preceding letter could itself end an earlier + // contraction that started before the batch — scalar. + bad |= 0b111u64 << i; + continue; + } + if i >= 61 { + bad |= u64::MAX << i; + break; + } + if i == last_forced { + // "x'll'd": the letter before this apostrophe is a + // consumed suffix's last char; a new (prefix) match + // starts here instead. + continue; + } + // The letter before this apostrophe may itself be a consumed + // suffix's last char resolved where last_forced can't see it + // (a scalar-walked zone like 'ſ, or a fixup before the + // batch): locally ambiguous, defer. + let p = scan + i; + let prev_suffix_possible = (bytes[p - 2] == b'\'' + && matches!(bytes[p - 1] | 0x20, b's' | b'd' | b'm' | b't')) + || (p >= 3 + && bytes[p - 3] == b'\'' + && (matches!( + (bytes[p - 2] | 0x20, bytes[p - 1] | 0x20), + (b'l', b'l') | (b'v', b'e') | (b'r', b'e') + ) || (bytes[p - 2] == 0xC5 && bytes[p - 1] == 0xBF))); + if prev_suffix_possible { + bad |= 0b111u64 << i; + continue; + } + let b1 = bytes[scan + i + 1]; + if b1 >= 0x80 { + bad |= 0b111u64 << i; + continue; + } + let k = match b1 | 0x20 { + b's' | b'd' | b'm' | b't' => 2, + b'l' if bytes[scan + i + 2] | 0x20 == b'l' => 3, + b'v' if bytes[scan + i + 2] | 0x20 == b'e' => 3, + b'r' if bytes[scan + i + 2] | 0x20 == b'e' => 3, + _ => 0, + }; + if k != 0 { + boundary &= !(1u64 << i); + boundary &= !(((1u64 << (k - 1)) - 1) << (i + 1)); + boundary |= 1u64 << (i + k); + last_forced = i + k; + } + } + } + + (boundary & !bad, bad) +} + +// ----------------------------------------------------------------------- +// Batch classifiers (per-arch front-ends) +// ----------------------------------------------------------------------- + +/// Carries for a batch known to have only ASCII in and just before it: +/// tail-class prev bytes route through the walkback, everything else +/// through the branchless ASCII carries. +#[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))] +#[inline(always)] +fn ascii_batch_carries(bytes: &[u8], scan: usize) -> OCarries { + if scan == 0 { + OCarries::default() + } else if is_tail_byte::(bytes[scan - 1]) { + tail_carries::(bytes, scan) + } else { + ascii_carries::(bytes, scan) + } +} + +/// `(usable, bad)` for `bytes[scan..scan+64]` under the o200k-family +/// rules — same contract as the cl100k family's `batch_masks`. +/// +/// NEON front-end: classifies the ASCII classes (letter, upper, digit, +/// space, whitespace, newline) with movemasks; apostrophe and slash sit +/// behind horizontal any-tests. Batches with non-ASCII in or just before +/// them take [`o200k_extended_masks`]. +#[cfg(target_arch = "aarch64")] +#[inline] +pub(crate) fn batch_masks< + const CONTRACTIONS: bool, + const DIGITS3: bool, + const SLASH: bool, + const HAN: bool, +>( + bytes: &[u8], + scan: usize, +) -> (u64, u64) { + use std::arch::aarch64::*; + let len = bytes.len(); + if scan + 70 > len { + // Not enough lookahead for the batch-edge char classification. + return (0, u64::MAX); + } + unsafe { + let p = bytes.as_ptr().add(scan); + let zero = vdupq_n_u8(0); + let mut lv = [zero; 4]; + let mut uv = [zero; 4]; + let mut dv = [zero; 4]; + let mut sv = [zero; 4]; + let mut wsv = [zero; 4]; + let mut nv = [zero; 4]; + let mut hiv = [zero; 4]; + let mut apv = [zero; 4]; + let mut slv = [zero; 4]; + for i in 0..4 { + let v = vld1q_u8(p.add(16 * i)); + let lowered = vorrq_u8(v, vdupq_n_u8(0x20)); + lv[i] = vcleq_u8(vsubq_u8(lowered, vdupq_n_u8(b'a')), vdupq_n_u8(25)); + uv[i] = vcleq_u8(vsubq_u8(v, vdupq_n_u8(b'A')), vdupq_n_u8(25)); + dv[i] = vcleq_u8(vsubq_u8(v, vdupq_n_u8(b'0')), vdupq_n_u8(9)); + sv[i] = vceqq_u8(v, vdupq_n_u8(b' ')); + wsv[i] = vorrq_u8( + sv[i], + vcleq_u8(vsubq_u8(v, vdupq_n_u8(9)), vdupq_n_u8(4)), + ); + nv[i] = vorrq_u8( + vceqq_u8(v, vdupq_n_u8(b'\r')), + vceqq_u8(v, vdupq_n_u8(b'\n')), + ); + hiv[i] = vcltzq_s8(vreinterpretq_s8_u8(v)); + apv[i] = vceqq_u8(v, vdupq_n_u8(b'\'')); + slv[i] = vceqq_u8(v, vdupq_n_u8(b'/')); + } + let l64 = mask::movemask64(lv[0], lv[1], lv[2], lv[3]); + let u64_ = mask::movemask64(uv[0], uv[1], uv[2], uv[3]); + let d64 = mask::movemask64(dv[0], dv[1], dv[2], dv[3]); + let s64 = mask::movemask64(sv[0], sv[1], sv[2], sv[3]); + let wsa = mask::movemask64(wsv[0], wsv[1], wsv[2], wsv[3]); + let n64 = mask::movemask64(nv[0], nv[1], nv[2], nv[3]); + + // Apostrophes only matter for the contraction fixup, slashes only + // for the tail smear: movemask each lazily. + let ap64 = if CONTRACTIONS { + let ap_any = vorrq_u8(vorrq_u8(apv[0], apv[1]), vorrq_u8(apv[2], apv[3])); + if vmaxvq_u8(ap_any) != 0 { + mask::movemask64(apv[0], apv[1], apv[2], apv[3]) + } else { + 0 + } + } else { + 0 + }; + let sl64 = if SLASH { + let sl_any = vorrq_u8(vorrq_u8(slv[0], slv[1]), vorrq_u8(slv[2], slv[3])); + if vmaxvq_u8(sl_any) != 0 { + mask::movemask64(slv[0], slv[1], slv[2], slv[3]) + } else { + 0 + } + } else { + 0 + }; + + let am = AsciiMasks { + l: l64, + d: d64, + s: s64, + wt: wsa & !s64 & !n64, + n: n64, + hi: 0, + ap: ap64, + }; + let ax = OAsciiExtra { up: u64_, sl: sl64 }; + + let hi_any = vorrq_u8(vorrq_u8(hiv[0], hiv[1]), vorrq_u8(hiv[2], hiv[3])); + if vmaxvq_u8(hi_any) != 0 + || (scan >= 1 && bytes[scan - 1] >= 0x80) + || (scan >= 2 && bytes[scan - 2] >= 0x80) + { + let mut am = am; + am.hi = mask::movemask64(hiv[0], hiv[1], hiv[2], hiv[3]); + return o200k_extended_masks::(bytes, scan, am, ax); + } + + let cr = ascii_batch_carries::(bytes, scan); + o200k_algebra::(bytes, scan, am, ax, cr, OUni::default()) + } +} + +/// x86-64 front-end: same contract as the NEON `batch_masks` above, +/// monomorphized on the SIMD tier (see `MaskScheme::batch_masks_x86`, +/// whose provided `batch_masks` supplies the runtime-dispatched form). +/// `#[inline(always)]` (with no `target_feature` of its own) so the body +/// fuses into whichever feature region calls it — LLVM's cost model +/// declined to inline the previous `#[target_feature]` form into the +/// tier-monomorphized fill wrappers and left a call per 64-byte batch. +/// +/// # Safety +/// +/// The selected tier must have been runtime-detected +/// ([`mask::avx512_scanner_available`] / +/// [`mask::avx2_scanner_available`]). +#[cfg(target_arch = "x86_64")] +#[inline(always)] +pub(crate) unsafe fn batch_masks_x86< + const AVX512: bool, + const CONTRACTIONS: bool, + const DIGITS3: bool, + const SLASH: bool, + const HAN: bool, +>( + bytes: &[u8], + scan: usize, +) -> (u64, u64) { + let len = bytes.len(); + if scan + 70 > len { + return (0, u64::MAX); + } + let (am, ax) = if AVX512 { + // SAFETY: the caller detected the AVX-512 tier (fn contract). + unsafe { (mask::ascii_masks_avx512(bytes, scan), ascii_extra_avx512(bytes, scan)) } + } else { + // SAFETY: the caller detected the AVX2 tier (fn contract). + unsafe { (mask::ascii_masks_avx2(bytes, scan), ascii_extra_avx2(bytes, scan)) } + }; + if am.hi != 0 + || (scan >= 1 && bytes[scan - 1] >= 0x80) + || (scan >= 2 && bytes[scan - 2] >= 0x80) + { + // SAFETY: both detected tiers include the BMI1/BMI2/LZCNT/POPCNT + // bit features `o200k_extended_masks` re-declares (fn contract). + return unsafe { + o200k_extended_masks::(bytes, scan, am, ax) + }; + } + let cr = ascii_batch_carries::(bytes, scan); + o200k_algebra::(bytes, scan, am, ax, cr, OUni::default()) +} + +/// The strict-uppercase and slash masks for `bytes[scan..scan+64]`, +/// AVX-512 tier. +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx512f,avx512bw,avx512vl,bmi1,bmi2,lzcnt,popcnt")] +#[inline] +fn ascii_extra_avx512(bytes: &[u8], scan: usize) -> OAsciiExtra { + use std::arch::x86_64::*; + unsafe { + let v = _mm512_loadu_si512(bytes.as_ptr().add(scan) as *const _); + let up = _mm512_cmple_epu8_mask( + _mm512_sub_epi8(v, _mm512_set1_epi8(b'A' as i8)), + _mm512_set1_epi8(25), + ); + let sl = _mm512_cmpeq_epi8_mask(v, _mm512_set1_epi8(b'/' as i8)); + OAsciiExtra { up, sl } + } +} + +/// The strict-uppercase and slash masks for `bytes[scan..scan+64]`, +/// AVX2 tier. `#[inline(never)]` for the same vector-domain reason as +/// [`mask::ascii_masks_avx2`]. +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx2,bmi1,bmi2,lzcnt,popcnt")] +#[inline(never)] +fn ascii_extra_avx2(bytes: &[u8], scan: usize) -> OAsciiExtra { + use std::arch::x86_64::*; + unsafe { + let le = |v: __m256i, lim: __m256i| -> __m256i { + _mm256_cmpeq_epi8(_mm256_min_epu8(v, lim), v) + }; + let mm = |m0: __m256i, m1: __m256i| -> u64 { + (_mm256_movemask_epi8(m0) as u32 as u64) + | ((_mm256_movemask_epi8(m1) as u32 as u64) << 32) + }; + let p = bytes.as_ptr().add(scan); + let v0 = _mm256_loadu_si256(p as *const _); + let v1 = _mm256_loadu_si256(p.add(32) as *const _); + let ca = _mm256_set1_epi8(b'A' as i8); + let c25 = _mm256_set1_epi8(25); + let up = mm( + le(_mm256_sub_epi8(v0, ca), c25), + le(_mm256_sub_epi8(v1, ca), c25), + ); + let slc = _mm256_set1_epi8(b'/' as i8); + let sl = mm(_mm256_cmpeq_epi8(v0, slc), _mm256_cmpeq_epi8(v1, slc)); + OAsciiExtra { up, sl } + } +} diff --git a/fast/src/pretok/unicode.rs b/fast/src/pretok/unicode.rs new file mode 100644 index 0000000..032e041 --- /dev/null +++ b/fast/src/pretok/unicode.rs @@ -0,0 +1,448 @@ +// Vendored from gigatoken (https://github.com/marcelroed/gigatoken), +// rev 542367a3efed134883fb4f1140b49c04e6fad3a3, MIT license. +// See src/pretok/mod.rs for what was trimmed and why. + +use icu_properties::props::{EnumeratedProperty, GeneralCategory, GeneralCategoryGroup, WhiteSpace}; +use icu_properties::CodePointSetData; + +#[inline] +pub(crate) fn get_general_category(c: char) -> GeneralCategory { + GeneralCategory::for_char(c) +} + +#[inline] +pub(crate) fn is_gc_letter(gc: GeneralCategory) -> bool { + GeneralCategoryGroup::Letter.contains(gc) +} + +#[inline] +pub(crate) fn is_gc_number(gc: GeneralCategory) -> bool { + GeneralCategoryGroup::Number.contains(gc) +} + +/// Unicode White_Space property — matches the same characters as `\s` in regex. +/// This includes GeneralCategory::Separator (Zs/Zl/Zp) PLUS control characters +/// like U+0009 (TAB), U+000A (LF), U+000D (CR), U+0085 (NEL), etc. +#[inline] +pub(crate) fn is_whitespace(c: char) -> bool { + // The set is a static compiled-data lookup, but cache the borrowed handle + // to avoid repeated constructor overhead. + static WS: std::sync::LazyLock> = + std::sync::LazyLock::new(CodePointSetData::new::); + WS.contains(c) +} + +#[inline] +pub(crate) fn is_letter(c: char) -> bool { + is_gc_letter(get_general_category(c)) +} + +#[inline] +pub(crate) fn is_number(c: char) -> bool { + is_gc_number(get_general_category(c)) +} + +#[inline] +pub(crate) fn is_other_complete(c: char) -> bool { + if c.is_ascii() { + return !c.is_ascii_alphanumeric() && !c.is_ascii_whitespace(); + } + let gc = get_general_category(c); + !is_gc_letter(gc) && !is_gc_number(gc) && !is_whitespace(c) +} + +// --------------------------------------------------------------------------- +// Packed codepoint → class table (hot-path classification) +// --------------------------------------------------------------------------- + +/// Character class as used by the pretokenization regexes: `\p{L}`, `\p{N}`, +/// `\s` (White_Space), and everything else. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(u8)] +pub(crate) enum CharClass { + Letter = 0, + Number = 1, + Whitespace = 2, + Other = 3, +} + +/// 2-bit class per codepoint, 4 codepoints per byte (~272 KiB total). +/// A single L1 load replaces the ICU GeneralCategory trie walk plus the +/// White_Space set binary search that the `is_*` predicates above pay per +/// call. Only the cache lines for scripts actually present in the input +/// stay resident. +static CLASS_TABLE: std::sync::LazyLock> = + std::sync::LazyLock::new(build_class_table); + +fn build_class_table() -> Box<[u8]> { + use icu_properties::CodePointMapData; + const N: usize = 0x110000; + let mut classes = vec![CharClass::Other as u8; N]; + let gc = CodePointMapData::::new(); + for (group, class) in [ + (GeneralCategoryGroup::Letter, CharClass::Letter), + (GeneralCategoryGroup::Number, CharClass::Number), + ] { + for range in gc.iter_ranges_for_group(group) { + classes[*range.start() as usize..=*range.end() as usize].fill(class as u8); + } + } + // White_Space is disjoint from GC Letter/Number, so fill order is moot. + for range in CodePointSetData::new::().iter_ranges() { + classes[*range.start() as usize..=*range.end() as usize].fill(CharClass::Whitespace as u8); + } + classes + .as_chunks::<4>().0.iter() + .map(|c| c[0] | (c[1] << 2) | (c[2] << 4) | (c[3] << 6)) + .collect() +} + +/// Pre-resolved handle to the packed class table. The static is a +/// `LazyLock>`, so every bare [`class_of`] call pays the +/// lazy-init state check plus a dependent load of the Box pointer before +/// the table load itself; per-char classify loops resolve the handle once +/// and index the slice directly. +#[derive(Clone, Copy)] +pub(crate) struct ClassTable(&'static [u8]); + +impl ClassTable { + #[inline] + pub(crate) fn get() -> Self { + Self(&CLASS_TABLE) + } + + /// [`class_of`] without the per-call static resolution. `cp` must be + /// a valid scalar value (guaranteed when decoded from valid UTF-8). + #[inline(always)] + pub(crate) fn class_of(self, cp: u32) -> CharClass { + debug_assert!(cp < 0x110000); + // SAFETY: `self.0` is CLASS_TABLE (the only constructor), sized + // 0x110000 / 4; cp >> 2 is in range for any scalar value. + let byte = unsafe { *self.0.get_unchecked((cp >> 2) as usize) }; + match (byte >> ((cp & 3) << 1)) & 3 { + 0 => CharClass::Letter, + 1 => CharClass::Number, + 2 => CharClass::Whitespace, + _ => CharClass::Other, + } + } +} + +/// Classify a codepoint with one table load. `cp` must be a valid scalar +/// value (guaranteed when decoded from valid UTF-8). +#[inline(always)] +pub(crate) fn class_of(cp: u32) -> CharClass { + ClassTable::get().class_of(cp) +} + +// --------------------------------------------------------------------------- +// DeepSeek character classes (finer split of `Other`) +// --------------------------------------------------------------------------- + +/// Character class as used by the DeepSeek V3 main regex, which additionally +/// distinguishes `\p{M}` (joins letter runs) and `\p{P}`/`\p{S}` (punctuation +/// runs) from the remaining `Other` codepoints (controls, format chars, +/// unassigned), which the regex leaves unmatched. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(u8)] +pub(crate) enum DsCharClass { + Letter = 0, + Number = 1, + Whitespace = 2, + Mark = 3, + PunctSym = 4, + Other = 5, +} + +/// Four-class view for schemes whose regex joins `\p{M}` into letter +/// runs and excludes it from punctuation runs (Qwen3.5's +/// `[\p{L}\p{M}]+` / `[^\s\p{L}\p{M}\p{N}]+`): marks classify as +/// letters, everything else as in [`class_of`]. +#[inline(always)] +pub(crate) fn class_of_marks_join(cp: u32) -> CharClass { + DsClassTable::get().class_of_marks_join(cp) +} + +/// 4-bit class per codepoint, 2 codepoints per byte (~544 KiB total). +static DS_CLASS_TABLE: std::sync::LazyLock> = + std::sync::LazyLock::new(build_ds_class_table); + +fn build_ds_class_table() -> Box<[u8]> { + use icu_properties::CodePointMapData; + const N: usize = 0x110000; + let mut classes = vec![DsCharClass::Other as u8; N]; + let gc = CodePointMapData::::new(); + for (group, class) in [ + (GeneralCategoryGroup::Letter, DsCharClass::Letter), + (GeneralCategoryGroup::Number, DsCharClass::Number), + (GeneralCategoryGroup::Mark, DsCharClass::Mark), + (GeneralCategoryGroup::Punctuation, DsCharClass::PunctSym), + (GeneralCategoryGroup::Symbol, DsCharClass::PunctSym), + ] { + for range in gc.iter_ranges_for_group(group) { + classes[*range.start() as usize..=*range.end() as usize].fill(class as u8); + } + } + // White_Space is disjoint from the groups above except Zs/Zl/Zp (which + // are in none of them), so fill order is moot. + for range in CodePointSetData::new::().iter_ranges() { + classes[*range.start() as usize..=*range.end() as usize] + .fill(DsCharClass::Whitespace as u8); + } + classes + .as_chunks::<2>().0.iter() + .map(|c| c[0] | (c[1] << 4)) + .collect() +} + +/// Pre-resolved handle to the packed DeepSeek class table — same +/// LazyLock-hoist rationale as [`ClassTable`]. +#[derive(Clone, Copy)] +pub(crate) struct DsClassTable(&'static [u8]); + +impl DsClassTable { + #[inline] + pub(crate) fn get() -> Self { + Self(&DS_CLASS_TABLE) + } + + /// [`ds_class_of`] without the per-call static resolution. `cp` must + /// be a valid scalar value (guaranteed when decoded from valid UTF-8). + #[inline(always)] + pub(crate) fn ds_class_of(self, cp: u32) -> DsCharClass { + debug_assert!(cp < 0x110000); + // SAFETY: `self.0` is DS_CLASS_TABLE (the only constructor), sized + // 0x110000 / 2; cp >> 1 is in range for any scalar value. + let byte = unsafe { *self.0.get_unchecked((cp >> 1) as usize) }; + match (byte >> ((cp & 1) << 2)) & 0xF { + 0 => DsCharClass::Letter, + 1 => DsCharClass::Number, + 2 => DsCharClass::Whitespace, + 3 => DsCharClass::Mark, + 4 => DsCharClass::PunctSym, + _ => DsCharClass::Other, + } + } + + /// [`class_of_marks_join`] without the per-call static resolution. + #[inline(always)] + pub(crate) fn class_of_marks_join(self, cp: u32) -> CharClass { + match self.ds_class_of(cp) { + DsCharClass::Letter | DsCharClass::Mark => CharClass::Letter, + DsCharClass::Number => CharClass::Number, + DsCharClass::Whitespace => CharClass::Whitespace, + DsCharClass::PunctSym | DsCharClass::Other => CharClass::Other, + } + } +} + +/// Classify a codepoint for the DeepSeek scheme with one table load. `cp` +/// must be a valid scalar value (guaranteed when decoded from valid UTF-8). +#[inline(always)] +pub(crate) fn ds_class_of(cp: u32) -> DsCharClass { + DsClassTable::get().ds_class_of(cp) +} + +// --------------------------------------------------------------------------- +// o200k character classes (case-aware split of Letter) +// --------------------------------------------------------------------------- + +/// Character class as used by the o200k regex family (gpt-oss, Nemotron-3), +/// whose letter runs are case-structured: +/// `[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+` etc. +/// `Upper` is Lu|Lt (the strict-uppercase classes that appear only in the +/// first bracket), `Lower` is Ll (only in the second), and `Caseless` is +/// Lm|Lo (in both). Marks (`\p{M}`) are their own class: they join letter +/// runs like `Caseless` but, being outside `\p{L}`, also continue +/// `[^\s\p{L}\p{N}]+` punctuation runs. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(u8)] +pub(crate) enum O200kCharClass { + Upper = 0, + Lower = 1, + Caseless = 2, + Mark = 3, + Number = 4, + Whitespace = 5, + Other = 6, +} + +/// 4-bit class per codepoint, 2 codepoints per byte (~544 KiB total). +static O200K_CLASS_TABLE: std::sync::LazyLock> = + std::sync::LazyLock::new(build_o200k_class_table); + +fn build_o200k_class_table() -> Box<[u8]> { + pack_nibbles(&o200k_classes_unpacked()) +} + +/// One `O200kCharClass` byte per codepoint (the unpacked form both the +/// o200k and Kimi table builders start from). +fn o200k_classes_unpacked() -> Vec { + use icu_properties::CodePointMapData; + const N: usize = 0x110000; + let mut classes = vec![O200kCharClass::Other as u8; N]; + let gc = CodePointMapData::::new(); + for (category, class) in [ + (GeneralCategory::UppercaseLetter, O200kCharClass::Upper), + (GeneralCategory::TitlecaseLetter, O200kCharClass::Upper), + (GeneralCategory::LowercaseLetter, O200kCharClass::Lower), + (GeneralCategory::ModifierLetter, O200kCharClass::Caseless), + (GeneralCategory::OtherLetter, O200kCharClass::Caseless), + ] { + for range in gc.iter_ranges_for_value(category) { + classes[*range.start() as usize..=*range.end() as usize].fill(class as u8); + } + } + for (group, class) in [ + (GeneralCategoryGroup::Mark, O200kCharClass::Mark), + (GeneralCategoryGroup::Number, O200kCharClass::Number), + ] { + for range in gc.iter_ranges_for_group(group) { + classes[*range.start() as usize..=*range.end() as usize].fill(class as u8); + } + } + // White_Space is disjoint from Letter/Mark/Number, so fill order is moot. + for range in CodePointSetData::new::().iter_ranges() { + classes[*range.start() as usize..=*range.end() as usize] + .fill(O200kCharClass::Whitespace as u8); + } + classes +} + +/// Pack one-byte-per-codepoint classes into 4-bit nibbles, 2 per byte. +fn pack_nibbles(classes: &[u8]) -> Box<[u8]> { + classes + .as_chunks::<2>().0.iter() + .map(|c| c[0] | (c[1] << 4)) + .collect() +} + +/// Classify a codepoint for the o200k scheme family with one table load. +/// `cp` must be a valid scalar value (guaranteed when decoded from valid +/// UTF-8). +#[inline(always)] +pub(crate) fn o200k_class_of(cp: u32) -> O200kCharClass { + debug_assert!(cp < 0x110000); + let byte = unsafe { *O200K_CLASS_TABLE.get_unchecked((cp >> 1) as usize) }; + match (byte >> ((cp & 1) << 2)) & 0xF { + 0 => O200kCharClass::Upper, + 1 => O200kCharClass::Lower, + 2 => O200kCharClass::Caseless, + 3 => O200kCharClass::Mark, + 4 => O200kCharClass::Number, + 5 => O200kCharClass::Whitespace, + _ => O200kCharClass::Other, + } +} + +// --------------------------------------------------------------------------- +// Kimi character classes (o200k classes with Script=Han split out) +// --------------------------------------------------------------------------- + +/// Character class for the Kimi (moonshotai K2 family) regex: the o200k +/// classes with `\p{Han}` split out. The pattern gives Han runs their own +/// leading alternative (`[\p{Han}]+`) and excludes Han from both letter +/// brackets (`[...&&[^\p{Han}]]`), but the general-category rules are +/// otherwise Han-blind: `\p{N}{1,3}` still counts a Han numeral and +/// `[^\s\p{L}\p{N}]+` still spans a Han symbol mid-run. Each Han variant +/// therefore records which base class the char behaves as outside a Han +/// run ([`Self::base`]). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(u8)] +pub(crate) enum KimiCharClass { + Upper = 0, + Lower = 1, + Caseless = 2, + Mark = 3, + Number = 4, + Whitespace = 5, + Other = 6, + /// Han letters (Lo, plus Lm like U+3005 々): the bulk of `\p{Han}`. + /// Never letter-run members; a maximal run of Han-class chars starting + /// a token is one `[\p{Han}]+` token. + Han = 7, + /// Han numerals (Nl: U+3007 〇, Suzhou numerals): `\p{N}` mid-number, + /// Han-run members otherwise. + HanNumber = 8, + /// Han symbols and marks (So: Kangxi radicals; Mc: U+16FF0/1 reading + /// marks, which `&&[^\p{Han}]` evicts from the letter brackets): + /// punct-run members mid-run, Han-run members at a token start. + HanOther = 9, +} + +impl KimiCharClass { + /// The o200k class the char behaves as in non-Han-run contexts. + #[inline(always)] + pub(crate) fn base(self) -> O200kCharClass { + match self { + KimiCharClass::Upper => O200kCharClass::Upper, + KimiCharClass::Lower => O200kCharClass::Lower, + KimiCharClass::Caseless | KimiCharClass::Han => O200kCharClass::Caseless, + KimiCharClass::Mark => O200kCharClass::Mark, + KimiCharClass::Number | KimiCharClass::HanNumber => O200kCharClass::Number, + KimiCharClass::Whitespace => O200kCharClass::Whitespace, + KimiCharClass::Other | KimiCharClass::HanOther => O200kCharClass::Other, + } + } + + /// Is the char in `\p{Han}` (a `[\p{Han}]+` run member)? + #[inline(always)] + pub(crate) fn is_han(self) -> bool { + self as u8 >= KimiCharClass::Han as u8 + } +} + +/// 4-bit class per codepoint, 2 codepoints per byte (~544 KiB total). +static KIMI_CLASS_TABLE: std::sync::LazyLock> = + std::sync::LazyLock::new(build_kimi_class_table); + +fn build_kimi_class_table() -> Box<[u8]> { + use icu_properties::CodePointMapData; + use icu_properties::props::Script; + let mut classes = o200k_classes_unpacked(); + let script = CodePointMapData::