hermes-tokenizer 1.8.101

//! Pretokenization: split documents into pretokens following a tokenizer's
//! pretokenization regex.
//!
//! The production implementations live in `fast` (one submodule per scheme:
//! `fast::r50k` for GPT-2, `fast::cl100k` for GPT-4, ...), selected via
//! [`PretokenizerType`]. Superseded designs (state machine, combinator,
//! SIMD prototypes) live in [`reference`] as benchmark baselines and test
//! oracles — nothing there runs in the encode path.
//!
//! The main entry points are:
//! - `pretokenize_as_iter`: iterate pretokens of a `&[u8]` (r50k scheme)
//! - `PretokenizerType::pretokenize`: iterate pretokens of any scheme
//! - `Pretokenize` trait: `doc.pretokens()` on any `&[u8]`

// Fast scanners carry target-specific helpers and iterator adapters that are
// intentionally dormant on some architectures and tokenizer families.
#![allow(dead_code, unused_assignments, unused_imports)]

pub(crate) use crate::pretokenize::pretoken::Pretoken;

pub mod fast;
mod options;
mod pretoken;
mod unicode;

pub use fast::{
    FastCl100kPretokenizer, FastDeepSeekV3Pretokenizer, FastOlmo3Pretokenizer,
    FastQwen2Pretokenizer, FastQwen35Pretokenizer, FastR50kPretokenizer,
};
pub use options::{FastPretokenizerDispatch, PretokenizerType};

/// Default document separator used in common training corpora.
pub const DEFAULT_SEPARATOR: &[u8] = b"<|endoftext|>";

/// Iterate the pretokens of `bytes` using the production (r50k) pretokenizer.
#[inline]
pub fn pretokenize_as_iter(bytes: &[u8]) -> FastR50kPretokenizer<'_> {
    FastR50kPretokenizer::new(bytes)
}

// ---------------------------------------------------------------------------
// Batched pretoken pulling (the encode loop's input interface)
// ---------------------------------------------------------------------------

/// Chunk size of [`PretokenSpans::fill_spans_keyed`] — the live entries of
/// one [`SpanBatch`] fill.
pub const PRETOKEN_CHUNK: usize = 256;

/// Both 64-bit halves of the per-length pack mask, in scalar ALU ops. A
/// u128 `MAX >> s` lowers to a multi-instruction sequence and the 16-entry
/// table this replaces put a dependent L1 load (2.43% of process) on the
/// `n → key → store` chain; per-half variable shifts are single 1-cycle
/// ops, so the halves cost two independent 3-deep chains and no load port.
/// The length tag (`n << 120`) touches only the high half and is OR'd in
/// by the caller. `const`: the phase-B emission loop's `PACK_MASK_TABLE`
/// (see `fast::mask`) is built from this at compile time, so the ALU and
/// table forms cannot drift apart.
#[inline(always)]
pub(crate) const fn pack_mask_halves(n: usize) -> (u64, u64) {
    debug_assert!(n >= 1 && n <= 15);
    let s = (n * 8) as u32;
    let lo = if n < 8 {
        u64::MAX >> (64u32.wrapping_sub(s) & 63)
    } else {
        u64::MAX
    };
    let hi = if n > 8 {
        u64::MAX >> (128u32.wrapping_sub(s) & 63)
    } else {
        0
    };
    (lo, hi)
}

// The key packers below (`pack_pretoken_key`, `fill_spans_keyed_with_buf`,
// phase B of the two-phase walker) read span bytes as native-endian words
// and mask, and the emit loop stores packed token lanes as one native-endian
// word — all little-endian layouts. A big-endian build would silently
// produce wrong keys and swapped tokens, so refuse to compile instead.
#[cfg(target_endian = "big")]
compile_error!("hermes-tokenizer requires little-endian token packing");

/// Pack a pretoken of ≤ 15 bytes into a `u128` cache key: bytes in the low
/// 15 lanes, length in the top byte (so keys of different lengths never
/// collide, and a real key is never 0). Returns `None` for longer
/// pretokens, which use the slice-keyed fallback map.
///
/// The common path is a single unaligned 16-byte load followed by a mask,
/// avoiding both a variable-length `memcpy` and per-byte branching. The
/// load is only taken when it cannot cross a page boundary, so it can
/// never touch an unmapped page; the rare near-boundary case falls back to
/// a plain copy. Both paths produce the identical key.
#[inline(always)]
pub(crate) fn pack_pretoken_key(bytes: &[u8]) -> Option<u128> {
    let n = bytes.len();
    if n > 15 {
        return None;
    }
    if n == 0 {
        // Empty pretokens (possible through the public API, never from a
        // pretokenizer) pack to key 0, which the short table reserves as
        // its empty sentinel — the encode loop routes key 0 to the long
        // map. Also keeps the read below from touching a zero-length
        // slice's dangling pointer.
        return Some(0);
    }
    let p = bytes.as_ptr();
    let low = if (p as usize) & 4095 <= 4096 - 16 {
        // SAFETY: the offset within the (≥ 4096-byte) page is ≤ 4096 - 16,
        // so a 16-byte read stays inside the page holding `p`, which is
        // mapped because `p` points to at least one valid byte.
        let v = unsafe { (p as *const u128).read_unaligned() };
        let (mask_lo, mask_hi) = pack_mask_halves(n);
        ((v as u64 & mask_lo) as u128) | ((((v >> 64) as u64 & mask_hi) as u128) << 64)
    } else {
        // Rare: `p` is within 16 bytes of a page boundary. Gather with a
        // plain copy (≤ 15 bytes) — correctness over speed on this cold
        // path. Lanes past `n` stay zero, so no mask is needed.
        let mut lanes = [0u8; 16];
        lanes[..n].copy_from_slice(bytes);
        u128::from_le_bytes(lanes)
    };
    Some(low | ((n as u128) << 120))
}

/// Hash of a packed pretoken key. Quality is noncritical for correctness
/// (the table compares full keys), but every consumer — the fill loops
/// (`fill_spans_keyed_with{,_buf}`, `fill_spans_two_phase`),
/// `ShortPretokenCache::grow`'s rehash, and the vocab-seeding paths
/// (`seeded_pretoken_cache`, `add_special_token`, `fork_sized`) — must
/// compute the same function of the key: the arms below produce different
/// values and may never mix in one process image. On aarch64 the arm is
/// picked at compile time; on x86_64 it is picked per process by
/// [`crc_hash_selected`], an immutable pure function of the CPU — this
/// entry point branches on that bit (cheap at the cold/slow sites that
/// use it, including the test-only [`fill_spans_keyed_with`]), and the
/// two hot fill loops instead embed one arm per monomorphization,
/// dispatched once per fill on the same bit (see [`fill_span_hash`]), so
/// the same key always hashes the same way. All
/// arms map key 0 to hash 0, which the fill loops' long-pretoken route
/// stores.
#[inline(always)]
pub(crate) fn pretoken_key_hash(key: u128) -> u64 {
    // Note: `crc` is in the default feature set for aarch64-apple-darwin
    // but NOT for aarch64-unknown-linux-gnu — generic aarch64 Linux builds
    // need `-C target-feature=+crc` (e.g. via RUSTFLAGS) to get this fast
    // hash; without it they silently take the multiply fold below.
    #[cfg(all(target_arch = "aarch64", target_feature = "crc"))]
    {
        // Hardware CRC32: two 3-cycle ops replace the 5-op multiply fold.
        // Linear over GF(2), so the low bits (the table index) see every
        // key bit; 32 bits suffice for any table under 2^32 slots.
        use core::arch::aarch64::__crc32d;
        // SAFETY: gated on the `crc` target feature at compile time.
        unsafe { __crc32d(__crc32d(0, key as u64), (key >> 64) as u64) as u64 }
    }
    #[cfg(target_arch = "x86_64")]
    {
        if crc_hash_selected() {
            // SAFETY: `crc_hash_selected` verified SSE4.2 support (or the
            // build enables it statically, folding this branch away).
            unsafe { pretoken_key_hash_crc32c(key) }
        } else {
            pretoken_key_hash_fold(key)
        }
    }
    #[cfg(not(any(
        all(target_arch = "aarch64", target_feature = "crc"),
        target_arch = "x86_64"
    )))]
    {
        pretoken_key_hash_fold(key)
    }
}

/// The multiply-fold arm of [`pretoken_key_hash`]: one folded multiply,
/// the cheapest mix whose low bits still see every key bit. Every target
/// can execute it; it is the process's hash wherever no hardware CRC arm
/// applies. Maps key 0 to hash 0 (0 · M = 0).
#[allow(dead_code)] // no cfg arm references it under aarch64 + crc
#[inline(always)]
fn pretoken_key_hash_fold(key: u128) -> u64 {
    let lo = key as u64;
    let hi = (key >> 64) as u64;
    let mut h = (lo ^ hi.rotate_right(25)).wrapping_mul(0x9E37_79B9_7F4A_7C15);
    h ^= h >> 32;
    h
}

/// The hardware CRC32C (SSE4.2) arm of [`pretoken_key_hash`]: same shape
/// and rationale as the aarch64 CRC32 arm — linear over GF(2) so the low
/// bits (the table index) see every key bit, 3-cycle latency and one µop
/// per `crc32` on Zen 2 (two chained ops vs the 5-op multiply fold), and
/// `_mm_crc32_u64(0, 0) == 0` preserves the key 0 -> hash 0 property the
/// fill loops' long-pretoken route stores.
///
/// `sse4.2` is NOT in baseline x86-64, so distributed wheels cannot gate
/// this arm at compile time; it is instead selected per process by
/// [`crc_hash_selected`] and reached only through call sites guarded by
/// that bit. Builds with `sse4.2` statically enabled (`-C
/// target-cpu=znver2`, any x86-64-v2+ setting) fold the guards away and
/// keep the pure-CRC codegen of a compile-time arm.
///
/// # Safety
///
/// The CPU must support SSE4.2: callers reach this only after
/// [`crc_hash_selected`] returned true (directly, or structurally via a
/// `fill_span_hash::<true>` monomorphization — see its contract), or from
/// a build with `sse4.2` statically enabled.
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "sse4.2")]
#[inline]
unsafe fn pretoken_key_hash_crc32c(key: u128) -> u64 {
    use core::arch::x86_64::_mm_crc32_u64;
    // SSE4.2 is enabled on this function; the caller (per the contract
    // above) only reaches it on a CPU that has it.
    _mm_crc32_u64(_mm_crc32_u64(0, key as u64), (key >> 64) as u64)
}

/// Does this process hash pretoken keys with CRC32C (x86_64)? A pure
/// function of the CPU, so one immutable answer for the process lifetime:
/// [`pretoken_key_hash`] and every fill-loop dispatch branch on this same
/// bit, which is what keeps the one-hash-per-process invariant airtight
/// without a per-key runtime check in the hot loops. std caches the
/// CPUID result, so after the first call this is a relaxed atomic load +
/// bit test; builds with `sse4.2` statically enabled const-fold it to
/// `true`.
#[cfg(target_arch = "x86_64")]
#[inline(always)]
pub(crate) fn crc_hash_selected() -> bool {
    std::arch::is_x86_feature_detected!("sse4.2")
}

/// Per-span hash for the monomorphized fill-loop bodies:
/// [`pretoken_key_hash`] with the x86_64 per-process selection hoisted
/// out of the per-span path. On x86_64 the two instantiations embed one
/// arm each, and each is reachable only under the matching value of
/// [`crc_hash_selected`]:
///
/// - `X86_CRC = true` bodies are called exclusively from the
///   `#[target_feature(enable = "sse4.2")]` fill wrappers, which their
///   dispatchers enter only when `crc_hash_selected()` is true;
/// - `X86_CRC = false` bodies are called exclusively from the dispatchers'
///   other arm, i.e. only when `crc_hash_selected()` is false (with
///   `sse4.2` statically enabled that arm is statically dead).
///
/// So every instantiation agrees with what [`pretoken_key_hash`] returns
/// for the same key in the same process. Off x86_64 the parameter is
/// ignored and this IS [`pretoken_key_hash`].
#[inline(always)]
pub(crate) fn fill_span_hash<const X86_CRC: bool>(key: u128) -> u64 {
    #[cfg(target_arch = "x86_64")]
    {
        if X86_CRC {
            // Reachability contract: only the sse4.2-gated fill wrappers
            // instantiate `X86_CRC = true`, so the selection bit must hold.
            debug_assert!(crc_hash_selected());
            // SAFETY: the `true` instantiation is only reachable from the
            // sse4.2-gated fill wrappers (see the contract above), so the
            // CPU has SSE4.2.
            unsafe { pretoken_key_hash_crc32c(key) }
        } else {
            pretoken_key_hash_fold(key)
        }
    }
    #[cfg(not(target_arch = "x86_64"))]
    {
        pretoken_key_hash(key)
    }
}

/// One batch slot: a pretoken span with its packed cache key and key hash,
/// as one 32-byte record (half a cache line, never straddling one thanks
/// to the alignment).
///
/// AoS instead of the previous three parallel arrays for dataflow on both
/// sides: the fill loops store one record with two `stp`s on a single
/// store stream (the parallel arrays cost 4 store µops across 3 streams —
/// the split u128 key store alone was two), and the probe loop's per-`i`
/// `(key, hash)` read touches one cache line instead of three.
///
/// `meta` carries the field the consumer needs next, keyed on `key`:
/// - `key != 0` (short pretoken, ≤ 15 bytes): `meta` is the full 64-bit
///   key hash. The span length rides in the key's top byte, so `ptr` +
///   `key >> 120` reconstructs the span on the (rare) slow path.
/// - `key == 0` (long pretoken, or an empty span through the public
///   adapter): `meta` is the span length in bytes. The hash is not stored:
///   the long route never probes the short table, and
///   `pretoken_key_hash(0) == 0` is what the old layout recorded anyway.
///   Prefetching `meta` as if it were a hash touches an arbitrary
///   (masked, in-bounds) table line — harmless, long pretokens are rare.
///
/// Fields are `pub(crate)`: only the in-crate fill loops may write entries
/// (safe external writes of an arbitrary `ptr`/`key` would let safe code
/// drive [`SpanBatch::span`]'s `from_raw_parts` with garbage — see the
/// [`PretokenSpans`] safety contract).
#[derive(Clone, Copy)]
#[repr(C, align(32))]
pub(crate) struct BatchEntry {
    pub(crate) key: u128,
    pub(crate) ptr: *const u8,
    pub(crate) meta: u64,
}

const _: () = assert!(std::mem::size_of::<BatchEntry>() == 32);

impl BatchEntry {
    /// Span length, independent of the short/long route.
    #[inline(always)]
    pub(crate) fn span_len(&self) -> usize {
        if self.key != 0 {
            (self.key >> 120) as usize
        } else {
            self.meta as usize
        }
    }
}

/// Readable slack entries past a full chunk, so the emit loop's
/// prefetch-ahead `entries[i + D].meta` load needs no index clamp (a
/// per-pretoken `add + cmp + csel` in the hottest loop). Slack entries
/// are never written by a fill; prefetching a stale or zero `meta`
/// requests an arbitrary masked (in-bounds) table line — harmless.
pub(crate) const SPAN_BATCH_SLACK: usize = 16;

/// One chunk of pretoken spans with their packed cache keys (0 = longer
/// than 15 bytes, routed to the slice-keyed fallback map) and key hashes,
/// filled by [`PretokenSpans::fill_spans_keyed`]. See [`BatchEntry`] for
/// the record layout. Fills only ever write the first [`PRETOKEN_CHUNK`]
/// entries; the tail is prefetch slack (see [`SPAN_BATCH_SLACK`]).
pub struct SpanBatch<'a> {
    /// `pub(crate)`: writable only by the in-crate fill loops, which uphold
    /// the [`PretokenSpans`] safety contract on every entry they write.
    pub(crate) entries: [BatchEntry; PRETOKEN_CHUNK + SPAN_BATCH_SLACK],
    /// The entries' `ptr`s borrow the spans' backing storage.
    _spans: std::marker::PhantomData<&'a [u8]>,
}

impl<'a> SpanBatch<'a> {
    pub fn new() -> Self {
        SpanBatch {
            entries: [BatchEntry {
                key: 0,
                ptr: std::ptr::null(),
                meta: 0,
            }; PRETOKEN_CHUNK + SPAN_BATCH_SLACK],
            _spans: std::marker::PhantomData,
        }
    }

    /// Reconstruct entry `i`'s span.
    ///
    /// # Safety
    /// Entry `i` must have been written by the most recent fill (`i` below
    /// its returned count), so `ptr` still points at a live span of `'a`.
    #[inline(always)]
    pub unsafe fn span(&self, i: usize) -> &'a [u8] {
        let e = &self.entries[i];
        // SAFETY: per the contract, `ptr` points at `span_len()` live bytes.
        unsafe { std::slice::from_raw_parts(e.ptr, e.span_len()) }
    }
}

impl Default for SpanBatch<'_> {
    fn default() -> Self {
        Self::new()
    }
}

/// A source of pretoken spans, pulled a chunk at a time with their cache
/// keys derived on the way out.
///
/// `Tokenizer::memoized_encode` consumes pretokens through this instead of
/// `Iterator` for codegen reasons: pulling happens in a dedicated
/// out-of-line loop, so the pretokenizer state is register-allocated
/// across the whole chunk (inlined into the register-starved encode loop
/// it lives in stack slots, ~9 cycles/pretoken of spill traffic on Zen 2).
/// Key packing, hashing, and the cache-line prefetch ride along in the
/// same loop because the span walker is a serial dependency chain (IPC
/// ~1.7 standalone): the independent per-span key math fills its idle
/// issue slots nearly for free, where a separate pass paid for it in full.
///
/// # Safety
///
/// The consumer trusts every fill unconditionally: after
/// [`Self::fill_spans_keyed`] returns `n`, the emit loop calls
/// [`SpanBatch::span`] (a raw `from_raw_parts`) on any entry `i < n`.
/// Implementations must therefore uphold, for every call returning `n`:
///
/// - `batch.entries[0..n]` were all written by THIS call (no stale entries
///   counted), and
/// - each written entry holds a valid `(ptr, len)` into caller-live input
///   bytes of lifetime `'a`: `ptr` non-null and readable for
///   [`BatchEntry::span_len`] bytes, with `key`/`meta` derived from exactly
///   those bytes via `pack_pretoken_key`/`pretoken_key_hash` semantics.
///
/// The entry fields are `pub(crate)`, so implementations outside this
/// crate cannot write entries at all and can only soundly return 0; the
/// in-crate fill helpers (`fill_spans_keyed_with{,_buf}`,
/// `fill_spans_two_phase`) uphold the contract.
pub unsafe trait PretokenSpans<'a> {
    /// Fill `batch` from the front with the next pretoken spans, calling
    /// `prefetch(hash)` for each. Returns how many were written; a short
    /// count (including 0) means the input is exhausted. (Recomputing the
    /// hash at the consumer instead of storing it here measured 4% slower
    /// end to end: the extra multiply sits on the probe loop's critical
    /// path, while this loop has store slots to spare.)
    fn fill_spans_keyed(&mut self, batch: &mut SpanBatch<'a>, prefetch: &impl Fn(u64)) -> usize;
}

/// Shared body of the iterator-backed [`PretokenSpans`] implementations
/// (spans with no single backing buffer): pull spans from `next` and derive
/// each one's key, hash, and prefetch on the way out. `#[inline(always)]`
/// so each `#[inline(never)]` implementation fuses it with its span walker
/// into a single out-of-line loop (see the trait docs for why that fusion
/// matters). Sources that walk one backing slice use
/// [`fill_spans_keyed_with_buf`] instead.
///
/// No production caller walks this path — the concrete pretokenizers and
/// the dispatch enum all fill through [`fill_spans_keyed_with_buf`] or
/// `fast::fill_spans_two_phase` — so unlike those two, it takes no
/// CRC-monomorphized wrapper: [`pretoken_key_hash`]'s per-span dispatch
/// branch (same process-immutable [`crc_hash_selected`] bit, so hash
/// values agree with every other site) is irrelevant off the hot path.
#[inline(always)]
pub(crate) fn fill_spans_keyed_with<'a>(
    mut next: impl FnMut() -> Option<&'a [u8]>,
    batch: &mut SpanBatch<'a>,
    prefetch: &impl Fn(u64),
) -> usize {
    let mut n = 0;
    while n < PRETOKEN_CHUNK {
        let Some(span) = next() else { break };
        let (key, h) = match pack_pretoken_key(span) {
            Some(key) => (key, pretoken_key_hash(key)),
            None => (0, 0),
        };
        prefetch(h);
        // Long (and empty) spans record their length; short spans their
        // hash (see `BatchEntry::meta`).
        let meta = if key != 0 { h } else { span.len() as u64 };
        batch.entries[n] = BatchEntry {
            key,
            ptr: span.as_ptr(),
            meta,
        };
        n += 1;
    }
    n
}

/// [`fill_spans_keyed_with`] for span sources that walk a single backing
/// slice, with `next` yielding `(start, end)` byte offsets into `bytes`.
/// Knowing the buffer removes the two data-dependent branches of
/// [`pack_pretoken_key`] from the per-span path:
///
/// - the `> 15` long-pretoken route becomes a select — for a long span the
///   16-byte load sits entirely inside the span, so it needs no guard, and
///   the select (not the mask clamp) provides the key-0 routing;
/// - the per-span page-boundary check (mispredict-prone: ~0.4% of spans on
///   4 KiB pages) becomes one buffer-end bound, hoisted per fill and false
///   only for short spans starting in the last 15 bytes of `bytes` — once
///   per input, so the branch predicts ~perfectly.
///
/// Empty spans cannot occur (`next` contract: `start < end`), so the key-0
/// route is exactly "longer than 15 bytes", as in the fallible packer.
#[inline(always)]
pub(crate) fn fill_spans_keyed_with_buf<'a>(
    bytes: &'a [u8],
    next: impl FnMut() -> Option<(usize, usize)>,
    batch: &mut SpanBatch<'a>,
    prefetch: &impl Fn(u64),
) -> usize {
    // Hash-arm dispatch, once per fill — see [`fill_spans_keyed_with`].
    #[cfg(target_arch = "x86_64")]
    if crc_hash_selected() {
        // SAFETY: `crc_hash_selected` verified SSE4.2 support.
        return unsafe { fill_spans_keyed_with_buf_crc(bytes, next, batch, prefetch) };
    }
    fill_spans_keyed_with_buf_impl::<false>(bytes, next, batch, prefetch)
}

/// The SSE4.2 (CRC-hash) monomorphization of [`fill_spans_keyed_with_buf`].
///
/// # Safety
///
/// The CPU must support SSE4.2 ([`crc_hash_selected`] must have returned
/// true).
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "sse4.2")]
unsafe fn fill_spans_keyed_with_buf_crc<'a>(
    bytes: &'a [u8],
    next: impl FnMut() -> Option<(usize, usize)>,
    batch: &mut SpanBatch<'a>,
    prefetch: &impl Fn(u64),
) -> usize {
    fill_spans_keyed_with_buf_impl::<true>(bytes, next, batch, prefetch)
}

/// [`fill_spans_keyed_with_buf`]'s loop body, monomorphized on the hash
/// arm (`X86_CRC` — see [`fill_span_hash`]'s reachability contract).
#[inline(always)]
fn fill_spans_keyed_with_buf_impl<'a, const X86_CRC: bool>(
    bytes: &'a [u8],
    mut next: impl FnMut() -> Option<(usize, usize)>,
    batch: &mut SpanBatch<'a>,
    prefetch: &impl Fn(u64),
) -> usize {
    // A 16-byte load at `start` is in bounds iff `start < tail_lim`.
    let tail_lim = bytes.len().saturating_sub(15);
    let mut n = 0;
    while n < PRETOKEN_CHUNK {
        let Some((start, end)) = next() else { break };
        debug_assert!(start < end && end <= bytes.len());
        // SAFETY: `next` returns in-bounds span boundaries.
        let span = unsafe { bytes.get_unchecked(start..end) };
        let len = end - start;
        let long = len > 15;
        // `|` not `||`: one combined test, taken for all but the tail spans.
        let key = if long | (start < tail_lim) {
            // SAFETY: `start < tail_lim` puts `start + 16` within `bytes`;
            // a long span (len ≥ 16) contains its own first 16 bytes.
            let v = unsafe { (bytes.as_ptr().add(start) as *const u128).read_unaligned() };
            let m = len.min(15);
            let (mask_lo, mask_hi) = pack_mask_halves(m);
            let lo = (v as u64) & mask_lo;
            let hi = ((v >> 64) as u64 & mask_hi) | ((m as u64) << 56);
            let packed = (lo as u128) | ((hi as u128) << 64);
            if long { 0 } else { packed }
        } else {
            // Short span starting in the buffer's last 15 bytes: gather
            // with a plain copy, once per input. Lanes past `len` stay
            // zero, so no mask is needed.
            let mut lanes = [0u8; 16];
            lanes[..len].copy_from_slice(span);
            u128::from_le_bytes(lanes) | ((len as u128) << 120)
        };
        let h = fill_span_hash::<X86_CRC>(key);
        prefetch(h);
        // Long spans record their length instead of the (unused) hash —
        // see `BatchEntry::meta`.
        let meta = if long { len as u64 } else { h };
        batch.entries[n] = BatchEntry {
            key,
            ptr: span.as_ptr(),
            meta,
        };
        n += 1;
    }
    n
}

/// Adapter giving any pretoken iterator (reference pretokenizers, tests,
/// custom sources) the [`PretokenSpans`] interface. The `Fast*`
/// pretokenizers implement the trait directly over their walker state
/// instead (see `fast::fill_spans_keyed_mask`): routing them through
/// `Iterator::next` left the (large, `#[inline(always)]`) `next_span`
/// un-inlined behind a real call — measured cost in
/// `fast::fill_spans_keyed_mask`'s docs.
pub struct SpanIter<I>(pub I);

// SAFETY: delegates to `fill_spans_keyed_with`, which writes exactly the
// first `n` entries from the iterator's live `'a` spans.
unsafe impl<'a, I: Iterator<Item = Pretoken<'a>>> PretokenSpans<'a> for SpanIter<I> {
    #[inline(never)]
    fn fill_spans_keyed(&mut self, batch: &mut SpanBatch<'a>, prefetch: &impl Fn(u64)) -> usize {
        fill_spans_keyed_with(|| self.0.next().map(|p| p.0), batch, prefetch)
    }
}

// ---------------------------------------------------------------------------
// Pretokenize trait — Layer 3
// ---------------------------------------------------------------------------

/// Anything that can be split into a stream of pretokens.
pub trait Pretokenize {
    fn pretokens(&self) -> FastR50kPretokenizer<'_>;
}

impl Pretokenize for [u8] {
    fn pretokens(&self) -> FastR50kPretokenizer<'_> {
        pretokenize_as_iter(self)
    }
}

// ---------------------------------------------------------------------------
// Pretoken-safe document splitting
// ---------------------------------------------------------------------------

/// Split `bytes` into ranges of roughly `target` bytes whose boundaries are
/// pretoken boundaries under every supported pretokenization scheme, so
/// encoding the ranges independently and concatenating the token streams is
/// identical to encoding `bytes` in one pass.
///
/// A boundary sits on a space that is preceded by an ASCII alphanumeric and
/// followed by an ASCII letter ("…word word…"). No scheme's pretoken can
/// cross such a point: whitespace only attaches to adjacent pretokens as a
/// single *leading* space of a following word (` ?\p{L}+` and friends), and
/// the only trailing attachments are `[\r\n]*`, which cannot contain a
/// space. Letter/digit runs cannot contain a space either, and the
/// all-whitespace rules (`\s+(?!\S)`, `\s*[\r\n]+`, …) never see a run that
/// crosses the boundary because the preceding byte is alphanumeric. The
/// three ASCII bytes also cannot sit inside a multi-byte UTF-8 character.
///
/// `added_tokens` are the byte sequences matched atomically *before*
/// pretokenization (see `Tokenizer::encode_with_added_tokens`), each paired
/// with its `rstrip` flag; a candidate boundary is rejected when an
/// occurrence of one straddles it, since the halves would otherwise be
/// BPE-encoded as plain text. Only tokens that contain a space can ever
/// straddle a boundary (every boundary sits on a space byte), so for typical
/// vocabularies the check costs nothing. If no occurrence crosses a
/// boundary, greedy leftmost-longest matching restarted there reproduces the
/// single-pass matches: the matcher's only state is its scan position, and
/// no match can carry it across the boundary.
///
/// An `rstrip` token absorbs the whitespace after its match, so a boundary
/// is also rejected when such an occurrence ends exactly at it — the space
/// opening the next chunk would encode as plain text instead of being
/// absorbed. (`lstrip` needs no counterpart: a boundary space preceding a
/// match lands at the start of the next chunk and is trimmed identically
/// there, and a boundary can never sit inside a longer whitespace run
/// because the byte before it must be alphanumeric.)
pub fn safe_split_ranges(
    bytes: &[u8],
    target: usize,
    added_tokens: &[(&[u8], bool)],
) -> Vec<std::ops::Range<usize>> {
    let blockers: Vec<memchr::memmem::Finder> = added_tokens
        .iter()
        .filter(|(t, _)| t.contains(&b' '))
        .map(|(t, _)| memchr::memmem::Finder::new(t))
        .collect();
    let max_blocker = blockers.iter().map(|f| f.needle().len()).max().unwrap_or(0);
    // rstrip tokens can only end at a boundary when their last byte is the
    // alphanumeric byte before the boundary space.
    let end_blockers: Vec<&[u8]> = added_tokens
        .iter()
        .filter(|(t, rstrip)| *rstrip && t.last().is_some_and(u8::is_ascii_alphanumeric))
        .map(|&(t, _)| t)
        .collect();
    // Whether an added-token occurrence spans the cut between `p - 1` and
    // `p`. Such an occurrence must start within `max_blocker - 1` bytes
    // before `p`, so searching a window of that radius is exhaustive.
    let cuts_added_token = |p: usize| -> bool {
        let lo = p.saturating_sub(max_blocker.saturating_sub(1));
        let hi = (p + max_blocker.saturating_sub(1)).min(bytes.len());
        blockers.iter().any(|f| {
            f.find_iter(&bytes[lo..hi])
                .any(|s| lo + s < p && lo + s + f.needle().len() > p)
        })
    };
    // Whether an rstrip added-token occurrence ends exactly at `p`.
    let ends_rstrip_token = |p: usize| end_blockers.iter().any(|t| bytes[..p].ends_with(t));
    let len = bytes.len();
    let target = target.max(1);
    let mut out = Vec::new();
    let mut start = 0;
    'chunks: while start < len {
        let mut probe = start + target;
        while probe + 1 < len {
            if bytes[probe] == b' '
                && bytes[probe - 1].is_ascii_alphanumeric()
                && bytes[probe + 1].is_ascii_alphabetic()
                && !(max_blocker > 0 && cuts_added_token(probe))
                && !(!end_blockers.is_empty() && ends_rstrip_token(probe))
            {
                out.push(start..probe);
                start = probe;
                continue 'chunks;
            }
            probe += 1;
        }
        out.push(start..len);
        break;
    }
    if out.is_empty() {
        out.push(0..0);
    }
    out
}

// ---------------------------------------------------------------------------
// Parallel pretokenization with document splitting
// ---------------------------------------------------------------------------