diff --git a/bindings/node/Cargo.lock b/bindings/node/Cargo.lock index 9fc267d0c..dfe74a851 100644 --- a/bindings/node/Cargo.lock +++ b/bindings/node/Cargo.lock @@ -54,6 +54,7 @@ name = "atomsplit" version = "0.1.0" dependencies = [ "memchr", + "wide", ] [[package]] diff --git a/bindings/python/Cargo.lock b/bindings/python/Cargo.lock index 4caf4c78c..213d91518 100644 --- a/bindings/python/Cargo.lock +++ b/bindings/python/Cargo.lock @@ -98,6 +98,7 @@ name = "atomsplit" version = "0.1.0" dependencies = [ "memchr", + "wide", ] [[package]] diff --git a/tokenizers/Cargo.lock b/tokenizers/Cargo.lock index 9e7c653e3..b7adeb744 100644 --- a/tokenizers/Cargo.lock +++ b/tokenizers/Cargo.lock @@ -76,6 +76,7 @@ dependencies = [ "memchr", "onig", "pcre2", + "wide", ] [[package]] diff --git a/tokenizers/atomsplit/Cargo.toml b/tokenizers/atomsplit/Cargo.toml index fb8947f3f..9e0bd7022 100644 --- a/tokenizers/atomsplit/Cargo.toml +++ b/tokenizers/atomsplit/Cargo.toml @@ -26,6 +26,7 @@ path = "src/lib.rs" [dependencies] memchr = "2.8.2" # SIMD search: used for single-byte search in CharDelimiterSplit (1.4–23× vs scalar), or multi-byte string pattern search in `literal`. +wide = "1.6.0" # NOTE: classify tables live in src/atom_tables.rs (committed, generated). Regenerate after any atom # scheme change with `cargo run -p bitmap_gen`. No build script / build-dep — atomsplit builds clean. # onig is C (Oniguruma) and fancy-regex pulls it in transitively for benches only — neither diff --git a/tokenizers/atomsplit/src/literal.rs b/tokenizers/atomsplit/src/literal.rs index c826b2dda..ff17fa847 100644 --- a/tokenizers/atomsplit/src/literal.rs +++ b/tokenizers/atomsplit/src/literal.rs @@ -6,51 +6,146 @@ //! The atom classification is unnecessary for pre-tokenizers splitting on an exact character or literal string: //! We can use a simpler byte search looking at 16 bytes at a time. -use memchr::memmem; -use std::fmt; +use std::{fmt, u32}; +use wide::{self, u8x32}; -/// The pattern handed to [`Literal::new`] was empty. +/// The pattern handed to [`Literal::new`] cannot be searched for. #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct EmptyPattern; +pub enum InvalidPattern { + /// An empty pattern would match everywhere. + Empty, + /// Longer than [`Literal::MAX_PATTERN_LEN`] bytes. + TooLong, +} -impl fmt::Display for EmptyPattern { +impl fmt::Display for InvalidPattern { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.write_str("an empty pattern matches everywhere") + match self { + Self::Empty => f.write_str("an empty pattern matches everywhere"), + Self::TooLong => write!( + f, + "patterns longer than {} bytes are not supported", + Literal::MAX_PATTERN_LEN + ), + } } } -impl std::error::Error for EmptyPattern {} +impl std::error::Error for InvalidPattern {} -/// A literal string to split on. -/// -/// The finder is boxed because it is large — a few hundred bytes of prefilter state on x86_64 — and -/// callers store it inside enums whose other variants are tiny. +/// A literal string to split on, at most [`Literal::MAX_PATTERN_LEN`] bytes. #[derive(Debug, Clone)] pub struct Literal { - finder: Box>, + pattern: [u8; Self::MAX_PATTERN_LEN], + pattern_len: u8, } impl Literal { + /// Longest supported pattern: the matcher compares one 32-byte block at a time, so a longer + /// pattern could never match inside a block. + pub const MAX_PATTERN_LEN: usize = 32; + /// # Errors - /// If `pattern` is empty, which would match everywhere. - pub fn new(pattern: &[u8]) -> Result { + /// If `pattern` is empty (it would match everywhere) or longer than + /// [`Literal::MAX_PATTERN_LEN`] bytes. + pub fn new(pattern: &[u8]) -> Result { if pattern.is_empty() { - return Err(EmptyPattern); + return Err(InvalidPattern::Empty); } + if pattern.len() > Self::MAX_PATTERN_LEN { + return Err(InvalidPattern::TooLong); + } + let mut buf = [0u8; Self::MAX_PATTERN_LEN]; + buf[..pattern.len()].copy_from_slice(pattern); Ok(Self { - finder: Box::new(memmem::Finder::new(pattern).into_owned()), + pattern: buf, + pattern_len: pattern.len() as u8, }) } /// The string being searched for. #[must_use] pub fn pattern(&self) -> &[u8] { - self.finder.needle() + &self.pattern[..usize::from(self.pattern_len)] } /// Byte offset of every match, left to right. Matches never overlap, so `"aa"` is found once in /// `"aaa"` — the same matches a regex engine would report. pub fn matches<'t>(&'t self, text: &'t [u8]) -> impl Iterator + 't { - self.finder.find_iter(text) + LiteralMatcher::new(text, self.pattern()) + } +} + +struct LiteralMatcher<'a> { + text: &'a [u8], + needle: &'a [u8], + needle_splats: Vec, + equality_mask: u32, + block_start: usize, + next_block: usize, + min_next_start: usize, +} + +impl<'a> LiteralMatcher<'a> { + pub fn new(text: &'a [u8], needle: &'a [u8]) -> Self { + Self { + text, + needle, + needle_splats: needle.iter().map(|&b| u8x32::splat(b)).collect(), + equality_mask: 0, + block_start: 0, + next_block: 0, + min_next_start: 0, + } + } + + pub fn load_more(&mut self) -> Option { + if self.next_block >= self.text.len() { + // We have read the whole text + return None; + } + self.block_start = self.next_block; + let remaining = &self.text[self.block_start..]; + let len_to_load = remaining.len().min(32); + let mut padded_text = [0u8; 32]; + padded_text[..len_to_load].copy_from_slice(&remaining[..len_to_load]); + let simd_bytes = u8x32::new(padded_text); + let mut mask = u32::MAX; + for (idx, &needle_splat) in self.needle_splats.iter().enumerate() { + mask &= simd_bytes.simd_eq(needle_splat).to_bitmask() >> idx; + } + + mask &= u32::MAX >> (self.needle.len() - 1); + if len_to_load < 32 { + let valid = len_to_load.saturating_sub(self.needle.len() - 1); + mask &= (1u32 << valid) - 1 + } + self.equality_mask = mask; + self.next_block += 32 - (self.needle.len() - 1); + Some(len_to_load) + } + + pub fn next_match(&mut self) -> Option { + loop { + // Load more until a match or text has been entirely read + while self.equality_mask == 0 { + self.load_more()?; + } + let offset = self.equality_mask.trailing_zeros() as usize; + self.equality_mask &= self.equality_mask - 1; + let start = self.block_start + offset; + if start >= self.min_next_start { + self.min_next_start = start + self.needle.len(); + return Some(start); + } + } + } +} + +impl<'a> Iterator for LiteralMatcher<'a> { + type Item = usize; + + fn next(&mut self) -> Option { + self.next_match() } } diff --git a/tokenizers/atomsplit/tests/literal.rs b/tokenizers/atomsplit/tests/literal.rs index 14ceed9a0..ec2b3339f 100644 --- a/tokenizers/atomsplit/tests/literal.rs +++ b/tokenizers/atomsplit/tests/literal.rs @@ -1,5 +1,5 @@ //! Tests for the literal search. Kept out of `src/` so the core stays production-only. -use atomsplit::literal::{EmptyPattern, Literal}; +use atomsplit::literal::{InvalidPattern, Literal}; #[test] fn finds_every_match_and_nothing_else() { @@ -29,22 +29,55 @@ fn matches_do_not_overlap() { assert_eq!(literal.matches(b"aaaa").collect::>(), [0, 2]); } -/// A `Literal` is stored inline in the normalizer and decoder enums, where every other variant is a -/// handful of bytes. `memmem::Finder` itself is a few hundred, so it has to stay behind a pointer. #[test] -fn a_literal_is_pointer_sized() { - assert_eq!( - size_of::(), - size_of::<*const u8>(), - "a Literal must not carry its finder inline" - ); +fn finds_matches_beyond_the_first_simd_load() { + // The matcher reads the text 32 bytes at a time; a match in the second load checks that + // reported offsets count from the start of the text, not the start of the load. + let mut text = vec![b'x'; 40]; + text.extend_from_slice(b"-yy"); + let literal = Literal::new(b"-").unwrap(); + assert_eq!(literal.matches(&text).collect::>(), [40]); +} + +#[test] +fn finds_a_match_straddling_two_simd_loads() { + // "ab" placed so that 'a' is the last byte of the first 32-byte load and 'b' the first byte + // of the second: the match is invisible to either load alone. + let mut text = vec![b'x'; 31]; + text.extend_from_slice(b"ab"); + text.extend_from_slice(&[b'x'; 10]); + let literal = Literal::new(b"ab").unwrap(); + assert_eq!(literal.matches(&text).collect::>(), [31]); +} + +#[test] +fn finds_matches_at_every_load_alignment() { + // 120 bytes of "xab" put a match at every offset of the form 3k + 1, so some match starts at + // every position relative to a 32-byte load, including each straddle of a load boundary. + let text = b"xab".repeat(40); + let literal = Literal::new(b"ab").unwrap(); + let expected: Vec = (0..40).map(|k| 3 * k + 1).collect(); + assert_eq!(literal.matches(&text).collect::>(), expected); } #[test] fn an_empty_pattern_is_rejected() { - assert_eq!(Literal::new(b"").unwrap_err(), EmptyPattern); + assert_eq!(Literal::new(b"").unwrap_err(), InvalidPattern::Empty); assert_eq!( - EmptyPattern.to_string(), + InvalidPattern::Empty.to_string(), "an empty pattern matches everywhere" ); } + +#[test] +fn patterns_longer_than_max_len_are_rejected() { + assert!(Literal::new(&[b'a'; Literal::MAX_PATTERN_LEN]).is_ok()); + assert_eq!( + Literal::new(&[b'a'; Literal::MAX_PATTERN_LEN + 1]).unwrap_err(), + InvalidPattern::TooLong + ); + assert_eq!( + InvalidPattern::TooLong.to_string(), + "patterns longer than 32 bytes are not supported" + ); +} diff --git a/tokenizers/tk-encode/src/tokenizer/pattern.rs b/tokenizers/tk-encode/src/tokenizer/pattern.rs index 5a147f5f4..e93a81a26 100644 --- a/tokenizers/tk-encode/src/tokenizer/pattern.rs +++ b/tokenizers/tk-encode/src/tokenizer/pattern.rs @@ -1,6 +1,6 @@ use crate::utils::SysRegex; use crate::{Offsets, Result}; -use atomsplit::literal::Literal; +use atomsplit::literal::{InvalidPattern, Literal}; use regex::Regex; /// Pattern used to split a NormalizedString @@ -73,7 +73,8 @@ impl Pattern for &str { match Literal::new(self.as_bytes()) { Ok(literal) => (&literal).find_matches(inside), // An empty pattern would match everywhere, so it matches nowhere instead. - Err(_) => Ok(vec![((0, inside.len()), false)]), + Err(InvalidPattern::Empty) => Ok(vec![((0, inside.len()), false)]), + Err(e @ InvalidPattern::TooLong) => Err(e.into()), } } }