splintr 0.19.1

Fast Rust tokenizer (BPE + SentencePiece + WordPiece) with Python bindings
Documentation
use regexr::{Regex as RegexrRegex, RegexBuilder};

#[cfg(feature = "pcre2")]
use pcre2::bytes::Regex as Pcre2Regex;

use super::error::TokenizerError;

/// Regex backend enum for switching between regexr (default) and PCRE2 (optional)
pub(super) enum RegexBackend {
    Regexr(Box<RegexrRegex>),
    /// A bundled pre-tokenizer expression, split by [`super::scanner`] rather
    /// than by the engine. The expression is not compiled at all on this path,
    /// which also takes it off the tokenizer's load time; the scanner's
    /// differential test compiles it separately and holds the two to the same
    /// spans.
    Scanner(super::scanner::SpanScanner),
    #[cfg(feature = "pcre2")]
    Pcre2(Pcre2Regex),
}

impl RegexBackend {
    /// Find all matches in the given text, returning (start, end) byte offsets
    ///
    /// Allocates the vector it returns. [`RegexBackend::find_into`] is the form
    /// the encode path uses, so that the vector can outlive the call and be
    /// reused; this one remains for callers that genuinely want an owned list.
    pub(super) fn find_iter(&self, text: &str) -> Vec<(usize, usize)> {
        let mut out = Vec::with_capacity(crate::core::pretokenizer::estimated_pieces(text));
        self.find_into(text, &mut out);
        out
    }

    /// [`RegexBackend::find_iter`] appending into a caller-owned buffer.
    pub(super) fn find_into(&self, text: &str, out: &mut Vec<(usize, usize)>) {
        match self {
            RegexBackend::Scanner(scan) => {
                // Sized up front rather than grown from empty. A ~110 byte text
                // produces ~25 spans, so an empty `Vec` reallocated four or five
                // times per call — and on macOS's xzone allocator a `realloc`
                // is a fresh allocation plus a `memmove` plus a free, not an
                // in-place extension the way glibc usually manages. Sampled on
                // an M1, that growth was 25% of single-text encode time.
                //
                // A reused buffer is already large enough and this is then a
                // no-op, which is the case the encode path is in.
                out.reserve(crate::core::pretokenizer::estimated_pieces(text));
                scan(text, out);
            }
            RegexBackend::Regexr(regex) => {
                out.extend(regex.find_iter(text).map(|m| (m.start(), m.end())))
            }
            #[cfg(feature = "pcre2")]
            RegexBackend::Pcre2(regex) => out.extend(
                regex
                    .find_iter(text.as_bytes())
                    .filter_map(|m| m.ok())
                    .map(|m| (m.start(), m.end())),
            ),
        }
    }
}

/// Compile one pre-tokenizer expression on the selected backend.
///
/// The single place that knows how each backend is configured (PCRE2 needs
/// `utf`+`ucp` to give `\p{…}` and `\s` their Unicode meanings), so a chained
/// pre-tokenizer's later passes are compiled exactly like its first one.
pub(super) fn compile_pattern(
    pattern: &str,
    use_pcre2: bool,
    use_jit: bool,
) -> Result<RegexBackend, TokenizerError> {
    #[cfg(feature = "pcre2")]
    if use_pcre2 {
        let mut regex_builder = pcre2::bytes::RegexBuilder::new();
        if use_jit {
            regex_builder.jit_if_available(true);
        }
        regex_builder.utf(true);
        regex_builder.ucp(true);
        return Ok(RegexBackend::Pcre2(regex_builder.build(pattern)?));
    }
    #[cfg(not(feature = "pcre2"))]
    let _ = use_pcre2;

    // Matched on exact expression text, so a caller passing a *different*
    // expression that merely looks similar keeps the engine: each scanner is
    // only proven equivalent to the one string it is paired with here, and that
    // pairing is what the differential test checks.
    //
    // DeepSeek's first two passes stay on the engine. They are single runs
    // rather than alternations, so the engine has little to do and a scanner
    // would win little.
    if let Some(scan) = super::scanner::for_pattern(pattern) {
        return Ok(RegexBackend::Scanner(scan));
    }

    let regex = RegexBuilder::new(pattern).jit(use_jit).build()?;
    Ok(RegexBackend::Regexr(Box::new(regex)))
}

/// One pass of llama.cpp's `unicode_regex_split` (`unicode.cpp:990-1088`).
///
/// Every span produced by the previous pass is re-matched **independently** —
/// the expression sees only that span's text, so `^`, `$` and lookaround treat
/// the span's edges as the edges of the world — and each span is replaced by the
/// ordered sequence of its matches AND the gaps between them
/// (`unicode_regex_split_stl`, `unicode.cpp:486-505`: an unmatched prefix is
/// emitted before every match, and any unmatched tail after the last one).
///
/// So a later expression NEVER re-examines the whole text and never merges
/// anything: it can only cut existing pieces finer. Nothing is exempted from a
/// later pass either — a gap left by pass 1 is an ordinary span that pass 2
/// subdivides like any other. That is why a list of N expressions is not the
/// alternation of those N expressions.
pub(super) fn subdivide(
    re: &RegexBackend,
    text: &str,
    spans: &[(usize, usize)],
) -> Vec<(usize, usize)> {
    let mut out = Vec::with_capacity(spans.len());
    for &(span_start, span_end) in spans {
        let Some(piece) = text.get(span_start..span_end) else {
            continue;
        };
        let mut last = 0;
        for (start, end) in re.find_iter(piece) {
            if start > last {
                out.push((span_start + last, span_start + start));
            }
            // A zero-width match would emit an empty piece upstream too; it
            // carries no bytes, so it is simply not recorded.
            if end > start {
                out.push((span_start + start, span_start + end));
            }
            last = end;
        }
        if last < piece.len() {
            out.push((span_start + last, span_end));
        }
    }
    out
}