From 4964fe480ccee648626295c6d2ef5cfee6348af5 Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Thu, 17 Sep 2026 13:26:58 -0700 Subject: [PATCH 1/3] perf(tokenizer): reuse PCRE2 scratch for each scan Signed-off-by: jthomson04 --- src/pre_tokenizers/split.rs | 81 ++++++++++++++++++++++++++++++++++++- 1 file changed, 80 insertions(+), 1 deletion(-) diff --git a/src/pre_tokenizers/split.rs b/src/pre_tokenizers/split.rs index f50f271..6886eea 100644 --- a/src/pre_tokenizers/split.rs +++ b/src/pre_tokenizers/split.rs @@ -994,8 +994,13 @@ fn find_matches_pcre2( let mut matches = Vec::with_capacity(input.len() / 3); let bytes = input.as_bytes(); let mut pos = 0; + if bytes.is_empty() { + return Ok(matches); + } + // Reuse one workspace while retaining find_at's exact stopping behavior. + let mut locations = regex.regex.capture_locations(); while pos < bytes.len() { - match regex.regex.find_at(bytes, pos) { + match regex.regex.captures_read_at(&mut locations, bytes, pos) { Ok(Some(m)) => { if m.start() == m.end() { pos = m.end() + 1; @@ -1026,6 +1031,80 @@ mod tests { use super::*; + #[test] + fn scratch_reuse_matches_stock_offsets_and_errors() { + fn stock( + input: &str, + base: usize, + regex: &Pcre2Regex, + ) -> Result, String> { + let mut out = Vec::new(); + let mut pos = 0; + while pos < input.len() { + match regex.regex.find_at(input.as_bytes(), pos) { + Ok(Some(m)) if m.start() == m.end() => pos = m.end() + 1, + Ok(Some(m)) => { + out.push((base + m.start(), base + m.end())); + pos = m.end(); + } + Ok(None) => break, + Err(e) => return Err(format!("PCRE2: {e}")), + } + } + Ok(out) + } + let mut inputs = vec![String::new()]; + let mut level = vec![String::new()]; + for _ in 0..4 { + level = level + .iter() + .flat_map(|s| ["a", "b", " ", "\n", "Γ©", "πŸ™‚"].map(|c| format!("{s}{c}"))) + .collect(); + inputs.extend(level.clone()); + } + inputs.push("aaaaaaaaaaaaaaaa!".into()); + inputs.push("hello δΈ–η•Œ 123\n".repeat(MIN_CHUNK_SIZE)); + for pattern in [ + r"a|(?=b)|$(?R)", + r"\w+|\s+|[^\w\s]", + r"a*", + r"", + r"^|$", + r"\b", + r"(?=a)", + r"(?<=a)b", + r"a|(?=b)|$", + r"\A.+\z", + r"\s+(?!\S)", + r"never_matches", + r"(*LIMIT_MATCH=1)^(a+)+$", + r"a\K", + r"(?=Γ©)|.", + ] { + let regex = Pcre2Regex { + regex: pcre2::bytes::RegexBuilder::new() + .utf(true) + .ucp(true) + .jit_if_available(true) + .build(pattern) + .unwrap(), + source: pattern.into(), + max_jit_stack_size: None, + }; + for input in &inputs { + let got = find_matches_pcre2(input, 17, ®ex).map_err(|e| match e { + Error::Unsupported(s) => s, + e => e.to_string(), + }); + assert_eq!( + got, + stock(input, 17, ®ex), + "pattern={pattern:?}, input={input:?}" + ); + } + } + } + // ── Behavior tests ────────────────────────────────── #[test] From 5541d6eff9acc9e1b02611998b5141257ca54fbc Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Thu, 17 Sep 2026 15:07:34 -0700 Subject: [PATCH 2/3] fix(tokenizer): preserve pooled scratch for custom JIT stacks Signed-off-by: jthomson04 --- src/pre_tokenizers/split.rs | 119 +++++++++++++++++++++--------------- 1 file changed, 70 insertions(+), 49 deletions(-) diff --git a/src/pre_tokenizers/split.rs b/src/pre_tokenizers/split.rs index 6886eea..0751745 100644 --- a/src/pre_tokenizers/split.rs +++ b/src/pre_tokenizers/split.rs @@ -997,10 +997,18 @@ fn find_matches_pcre2( if bytes.is_empty() { return Ok(matches); } - // Reuse one workspace while retaining find_at's exact stopping behavior. - let mut locations = regex.regex.capture_locations(); + // Reuse one workspace within this scan for the default JIT stack. Custom + // stacks use find_at's pool so they are not allocated and freed per scan. + let mut locations = regex + .max_jit_stack_size + .is_none() + .then(|| regex.regex.capture_locations()); while pos < bytes.len() { - match regex.regex.captures_read_at(&mut locations, bytes, pos) { + let found = match &mut locations { + Some(locations) => regex.regex.captures_read_at(locations, bytes, pos), + None => regex.regex.find_at(bytes, pos), + }; + match found { Ok(Some(m)) => { if m.start() == m.end() { pos = m.end() + 1; @@ -1081,26 +1089,26 @@ mod tests { r"a\K", r"(?=Γ©)|.", ] { - let regex = Pcre2Regex { - regex: pcre2::bytes::RegexBuilder::new() - .utf(true) - .ucp(true) - .jit_if_available(true) - .build(pattern) - .unwrap(), - source: pattern.into(), - max_jit_stack_size: None, - }; - for input in &inputs { - let got = find_matches_pcre2(input, 17, ®ex).map_err(|e| match e { - Error::Unsupported(s) => s, - e => e.to_string(), - }); - assert_eq!( - got, - stock(input, 17, ®ex), - "pattern={pattern:?}, input={input:?}" - ); + for max_jit_stack_size in [None, Some(1 << 20)] { + let limits = Pcre2Limits { + max_jit_stack_size, + ..Default::default() + }; + let regexes = try_compile_pcre2_regexes(pattern, 1, limits) + .unwrap() + .unwrap(); + let regex = ®exes[0]; + for input in &inputs { + let got = find_matches_pcre2(input, 17, regex).map_err(|e| match e { + Error::Unsupported(s) => s, + e => panic!("unexpected error variant: {e:?}"), + }); + assert_eq!( + got, + stock(input, 17, regex), + "pattern={pattern:?}, stack={max_jit_stack_size:?}, input={input:?}" + ); + } } } } @@ -1742,32 +1750,45 @@ mod tests { } assert!(input.len() >= 2 * chunk, "input must trigger parallel path"); - let split = - Split::from_config(&serde_json::json!({"Regex": "[a-z]+"}), "Isolated", false).unwrap(); - - let pieces = split.split(&input).unwrap(); - // There should be exactly 3 pieces: digits, long 'a' run, digits. - // The 'a' run must be ONE piece, not split across chunks. - let long_piece = pieces.iter().find(|p| p.starts_with('a')).unwrap(); - assert_eq!( - long_piece.len(), - match_len, - "long match should be {match_len} bytes, got {}; \ - match was split across chunk boundaries", - long_piece.len(), - ); - // Also verify via byte offsets - let a_pieces: Vec<&str> = pieces - .iter() - .copied() - .filter(|p| p.starts_with('a')) - .collect(); - assert_eq!( - a_pieces.len(), - 1, - "should be exactly one 'a' piece, got {a_pieces:?}" - ); - assert_eq!(&input[match_start..match_end], *long_piece,); + for max_jit_stack_size in [None, Some(1 << 20)] { + let limits = Pcre2Limits { + max_jit_stack_size, + ..Default::default() + }; + let mut split = Split::from_config_with_limits( + &json!({"Regex": "[a-z]+"}), + "Isolated", + false, + limits, + ) + .unwrap(); + // Force two matchers even when the test runs with one available CPU. + split.pcre2_regexes = try_compile_pcre2_regexes("[a-z]+", 2, limits).unwrap(); + + let pieces = split.split(&input).unwrap(); + // There should be exactly 3 pieces: digits, long 'a' run, digits. + // The 'a' run must be ONE piece, not split across chunks. + let long_piece = pieces.iter().find(|p| p.starts_with('a')).unwrap(); + assert_eq!( + long_piece.len(), + match_len, + "long match should be {match_len} bytes, got {}; \ + match was split across chunk boundaries", + long_piece.len(), + ); + // Also verify via byte offsets + let a_pieces: Vec<&str> = pieces + .iter() + .copied() + .filter(|p| p.starts_with('a')) + .collect(); + assert_eq!( + a_pieces.len(), + 1, + "should be exactly one 'a' piece, got {a_pieces:?}" + ); + assert_eq!(&input[match_start..match_end], *long_piece,); + } } // ── Incremental cache: lookahead at divergence boundary ── From cf923b7e5ddd7fe57c22211ee11f5699ed29c312 Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Thu, 17 Sep 2026 16:00:46 -0700 Subject: [PATCH 3/3] test(tokenizer): require PCRE2 in boundary repair test Signed-off-by: jthomson04 --- src/pre_tokenizers/split.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/pre_tokenizers/split.rs b/src/pre_tokenizers/split.rs index 0751745..7db5ac9 100644 --- a/src/pre_tokenizers/split.rs +++ b/src/pre_tokenizers/split.rs @@ -1723,6 +1723,8 @@ mod tests { #[test] fn long_match_crosses_parallel_boundary() { + const PATTERN: &str = "[a-z]+"; + // Build an input large enough for parallel matching (>= 2 * MIN_CHUNK_SIZE = 16KB). // Place a long run of lowercase letters that exceeds CHUNK_OVERLAP (1KB) // and spans the authority zone boundary. @@ -1756,14 +1758,18 @@ mod tests { ..Default::default() }; let mut split = Split::from_config_with_limits( - &json!({"Regex": "[a-z]+"}), + &json!({"Regex": PATTERN}), "Isolated", false, limits, ) .unwrap(); // Force two matchers even when the test runs with one available CPU. - split.pcre2_regexes = try_compile_pcre2_regexes("[a-z]+", 2, limits).unwrap(); + split.pcre2_regexes = Some( + try_compile_pcre2_regexes(PATTERN, 2, limits) + .unwrap() + .expect("PCRE2 unavailable; boundary repair path not covered"), + ); let pieces = split.split(&input).unwrap(); // There should be exactly 3 pieces: digits, long 'a' run, digits.