diff --git a/crates/sparse-ngrams/Cargo.toml b/crates/sparse-ngrams/Cargo.toml index d1491ef..2da4bc9 100644 --- a/crates/sparse-ngrams/Cargo.toml +++ b/crates/sparse-ngrams/Cargo.toml @@ -1,12 +1,15 @@ [package] name = "sparse-ngrams" -version = "0.1.0" -edition = "2021" +version = "0.2.0" +edition = "2024" description = "Fast sparse n-gram extraction from byte slices." repository = "https://github.com/github/rust-gems" license = "MIT" -keywords = ["ngram", "algorithm", "search", "index"] -categories = ["algorithms", "data-structures", "text-processing"] +keywords = ["sparse-ngram", "ngram", "algorithm", "search", "code-search", "regular-expression"] +categories = ["algorithms", "data-structures", "text-processing", "search-engine", "data-retrieval"] +# `data/` holds the reference bigram ranking used only for offline model training/verification; +# it is not needed to build or use the crate, so keep it out of the published package. +exclude = ["data/"] [lib] bench = false @@ -16,6 +19,7 @@ bench = false serde = ["dep:serde"] [dependencies] +casefold = "0.1" serde = { version = "1", features = ["derive"], optional = true } [[bench]] diff --git a/crates/sparse-ngrams/README.md b/crates/sparse-ngrams/README.md index 3d9b465..9358b14 100644 --- a/crates/sparse-ngrams/README.md +++ b/crates/sparse-ngrams/README.md @@ -46,13 +46,49 @@ invoked once per n-gram in emission order: use sparse_ngrams::{collect_sparse_grams_deque, NGram}; let mut count = 0; -collect_sparse_grams_deque(b"hello world", |gram: NGram| { +collect_sparse_grams_deque(b"hello world", |gram: NGram, _idx| { count += 1; // ... insert `gram` into an index, hash it, etc. }); assert!(count > 0); ``` +### Query-time extraction + +`collect_sparse_grams` emits *every* candidate gram, which is what you want when building an index. +At query time you instead want the *minimum* set of grams that still covers the query string, fed +incrementally as the user types. `QueryGrams` is a streaming state machine for exactly that: it +accepts one character (or already index-folded byte) at a time, emits grams as soon as they are +determined, and can be `flush`ed to drain the tail. + +The consumer receives `(gram, end, follow)`: the n-gram, the position of the character just after +it, and that following byte when it has already been fed (`None` at the current stream end). + +```rust +use sparse_ngrams::{QueryGrams, NGram}; + +let mut q = QueryGrams::default(); +let mut grams = Vec::new(); +// Feed the query one character at a time (each is index-folded internally). +for c in "hello world".chars() { + q.append_char(c, |gram: NGram, _end: u32, _follow: Option| { + grams.push(gram); + }); +} +// Drain the remaining tail grams. +q.flush(|gram: NGram, _end, _follow| { + grams.push(gram); +}); + +assert!(!grams.is_empty()); +assert!(grams.iter().all(|g| g.len() >= 2 && g.len() <= 8)); +``` + +`QueryGrams` is `Clone` and hashes/compares by its canonical `state()`, so it can be used as an +automaton state (e.g. cloned across branches while traversing a trie). Callers that buffer +already-folded bytes can feed them with `append_byte` (using the re-exported `index_fold_char` to +fold identically), and `consume_first` drains a single leading gram to shrink retained state. + ## Performance Throughput on an Apple M4 Max (the ~15 KB `benchmarks/fixtures/sample_code.txt` corpus): diff --git a/crates/sparse-ngrams/benchmarks/performance.rs b/crates/sparse-ngrams/benchmarks/performance.rs index 1f8825d..9828ebb 100644 --- a/crates/sparse-ngrams/benchmarks/performance.rs +++ b/crates/sparse-ngrams/benchmarks/performance.rs @@ -1,8 +1,8 @@ use std::hint::black_box; -use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; use sparse_ngrams::{ - collect_sparse_grams_deque, collect_sparse_grams_scan, max_sparse_grams, NGram, + NGram, collect_sparse_grams_deque, collect_sparse_grams_scan, max_sparse_grams, }; fn bench_collect(c: &mut Criterion) { @@ -28,7 +28,7 @@ fn bench_collect(c: &mut Criterion) { group.bench_with_input(BenchmarkId::new("deque", name), input, |b, input| { b.iter(|| { let mut w = 0usize; - collect_sparse_grams_deque(black_box(input), |gram| { + collect_sparse_grams_deque(black_box(input), |gram, _idx| { buf[w] = gram; w += 1; }); @@ -38,7 +38,7 @@ fn bench_collect(c: &mut Criterion) { group.bench_with_input(BenchmarkId::new("scan", name), input, |b, input| { b.iter(|| { let mut w = 0usize; - collect_sparse_grams_scan(black_box(input), |gram| { + collect_sparse_grams_scan(black_box(input), |gram, _idx| { buf[w] = gram; w += 1; }); diff --git a/crates/sparse-ngrams/src/extract.rs b/crates/sparse-ngrams/src/extract.rs index d95de81..90fabd3 100644 --- a/crates/sparse-ngrams/src/extract.rs +++ b/crates/sparse-ngrams/src/extract.rs @@ -1,8 +1,8 @@ //! Core sparse n-gram extraction algorithm. +use crate::MAX_SPARSE_GRAM_SIZE; use crate::ngram::NGram; use crate::table::{bigram_h, bigram_priority_rolling}; -use crate::MAX_SPARSE_GRAM_SIZE; /// Returns the maximum number of sparse n-grams that can be produced from /// `content_len` bytes of input. Use this to reserve capacity for the output (e.g. a `Vec` the @@ -38,7 +38,7 @@ pub fn collect_sparse_grams(content: &[u8]) -> Vec { let mut out = Vec::with_capacity(max_sparse_grams(content.len())); let spare = out.spare_capacity_mut(); let mut w = 0; - collect_sparse_grams_deque(content, |gram| { + collect_sparse_grams_deque(content, |gram, _idx| { spare[w].write(gram); w += 1; }); @@ -51,9 +51,11 @@ pub fn collect_sparse_grams(content: &[u8]) -> Vec { /// Monotone-deque extraction, with the deque held in fixed ring buffers. Calls `emit` once for /// every sparse n-gram, in emission order (all bigrams, plus algorithmically selected longer -/// grams). `emit` decides what to do with each gram — push it into a `Vec`, write it into a -/// pre-sized slice, feed it straight into an index, etc. — so no output buffer needs to be sized -/// or allocated up front. +/// grams), as `(gram, idx)` where `idx` is the position of the character just after `gram` in +/// `content`. +/// `emit` decides what to do with each gram — push it into a `Vec`, write it into a pre-sized +/// slice, feed it straight into an index, etc. — so no output buffer needs to be sized or +/// allocated up front. /// /// The boundary candidates form a monotone run, kept in fixed `[_; MAX_SPARSE_GRAM_SIZE]` ring /// buffers addressed by a single running depth `tail` — there is no head. Rather than dropping the @@ -71,10 +73,10 @@ pub fn collect_sparse_grams(content: &[u8]) -> Vec { /// use sparse_ngrams::{collect_sparse_grams_deque, NGram}; /// /// let mut count = 0; -/// collect_sparse_grams_deque(b"hello world", |_gram: NGram| count += 1); +/// collect_sparse_grams_deque(b"hello world", |_gram: NGram, _idx| count += 1); /// assert!(count > 0); /// ``` -pub fn collect_sparse_grams_deque(content: &[u8], mut emit: impl FnMut(NGram)) { +pub fn collect_sparse_grams_deque(content: &[u8], mut emit: impl FnMut(NGram, u32)) { let n = content.len(); if n < 2 { return; @@ -108,7 +110,7 @@ pub fn collect_sparse_grams_deque(content: &[u8], mut emit: impl FnMut(NGram)) { h = h_b; // The bigram (length 2) is always emitted. - emit(window_to_gram(window, 2)); + emit(window_to_gram(window, 2), idx + 1); // Walk back over the candidates from the tail, emitting one gram each. Stop at the first // candidate too far back to form a gram of at most `MAX_SPARSE_GRAM_SIZE` bytes (this @@ -123,10 +125,10 @@ pub fn collect_sparse_grams_deque(content: &[u8], mut emit: impl FnMut(NGram)) { if idx.wrapping_sub(begin) + 1 >= MAX_SPARSE_GRAM_SIZE as u32 { break; } - emit(window_to_gram( - window, - (idx.wrapping_sub(begin) + 2) as usize, - )); + emit( + window_to_gram(window, (idx.wrapping_sub(begin) + 2) as usize), + idx + 1, + ); let bval = val_buf[slot]; if bval < value { break; @@ -145,12 +147,13 @@ pub fn collect_sparse_grams_deque(content: &[u8], mut emit: impl FnMut(NGram)) { } /// Queue-free scan-based extraction. Calls `emit` once for every sparse n-gram, in the same order -/// as [`collect_sparse_grams_deque`]. +/// as [`collect_sparse_grams_deque`], as `(gram, idx)` where `idx` is the position of the character +/// just after `gram` in `content`. /// /// Produces identical output (same order) as [`collect_sparse_grams_deque`]: the boundary /// candidates are exactly the positions where a backward scan hits a new suffix minimum, so a /// fixed-size ring of recent priorities replaces the monotone deque. -pub fn collect_sparse_grams_scan(content: &[u8], mut emit: impl FnMut(NGram)) { +pub fn collect_sparse_grams_scan(content: &[u8], mut emit: impl FnMut(NGram, u32)) { let n = content.len(); if n < 2 { return; @@ -168,7 +171,7 @@ pub fn collect_sparse_grams_scan(content: &[u8], mut emit: impl FnMut(NGram)) { priorities[idx as usize & MASK] = v1; // The bigram (length 2) is always emitted. - emit(window_to_gram(window, 2)); + emit(window_to_gram(window, 2), idx + 1); // Scan backwards, tracking the minimum interior priority seen so far. Each new strict // minimum is a boundary candidate; emit its gram while the right boundary `v1` is strictly @@ -184,7 +187,7 @@ pub fn collect_sparse_grams_scan(content: &[u8], mut emit: impl FnMut(NGram)) { break; } let len = d as usize + 2; - emit(window_to_gram(window, len)); + emit(window_to_gram(window, len), idx + 1); running_min = v_p; } } @@ -197,9 +200,9 @@ mod tests { use crate::table::bigram_priority; use std::collections::HashSet; - fn collect_to_vec(run: impl FnOnce(&mut dyn FnMut(NGram))) -> Vec { + fn collect_to_vec(run: impl FnOnce(&mut dyn FnMut(NGram, u32))) -> Vec { let mut out = Vec::new(); - run(&mut |gram| out.push(gram)); + run(&mut |gram, _idx| out.push(gram)); out } diff --git a/crates/sparse-ngrams/src/lib.rs b/crates/sparse-ngrams/src/lib.rs index f6af7ef..5fa5921 100644 --- a/crates/sparse-ngrams/src/lib.rs +++ b/crates/sparse-ngrams/src/lib.rs @@ -40,11 +40,17 @@ mod extract; mod ngram; +mod query; mod table; pub use ngram::NGram; +pub use query::QueryGrams; pub use table::bigram_priority; +/// Re-export of the index-folding used internally by [`QueryGrams::append_char`], so callers that +/// buffer already-folded bytes (and feed them via [`QueryGrams::append_byte`]) fold identically. +pub use casefold::index_fold_char; + /// Maximum length (in bytes) of a sparse n-gram. pub const MAX_SPARSE_GRAM_SIZE: usize = 8; diff --git a/crates/sparse-ngrams/src/ngram.rs b/crates/sparse-ngrams/src/ngram.rs index 413e4ec..87fad82 100644 --- a/crates/sparse-ngrams/src/ngram.rs +++ b/crates/sparse-ngrams/src/ngram.rs @@ -131,6 +131,23 @@ impl NGram { Self::pack(len, payload) } + /// Builds an `NGram` from a right-shifted window where the gram's end byte is in the low + /// byte and older bytes are above it. This helper keeps only the low `len` bytes, aligns them + /// into the top bytes, and then delegates to [`from_window`](Self::from_window). + #[inline] + pub(crate) fn from_window_masked(value: u64, len: usize) -> Self { + debug_assert!( + (Self::LEN_BIAS as usize..=MAX_SPARSE_GRAM_SIZE).contains(&len), + "ngram length {len} out of range [{}, {}]", + Self::LEN_BIAS, + MAX_SPARSE_GRAM_SIZE, + ); + let bits = (len * 8) as u32; + let low_mask = u64::MAX >> (u64::BITS - bits); + let aligned = (value & low_mask) << ((MAX_SPARSE_GRAM_SIZE - len) * 8); + Self::from_window(aligned, len) + } + /// Packs a length and payload into the structured value, then stores it *mixed* (via [`mix27`]) /// so the hot sorting/bucketing paths read a well-distributed value directly from the field; /// [`len`](Self::len) and [`Debug`] unmix on demand. diff --git a/crates/sparse-ngrams/src/query.rs b/crates/sparse-ngrams/src/query.rs new file mode 100644 index 0000000..2826c2c --- /dev/null +++ b/crates/sparse-ngrams/src/query.rs @@ -0,0 +1,802 @@ +//! Streaming sparse n-gram extraction state for query-time traversal. +//! +//! Unlike indexing, which emits every candidate gram, query extraction emits the minimum number of +//! grams that cover the input stream and supports incremental character feeding, cloning, and +//! partial/full draining. + +use std::hash::{Hash, Hasher}; + +use casefold::index_fold_char; + +use crate::MAX_SPARSE_GRAM_SIZE; +use crate::ngram::NGram; +use crate::table::{bigram_h, bigram_priority_rolling}; + +/// Bit mask that wraps ring-buffer indices into `[0, MAX_SPARSE_GRAM_SIZE)`. +const RING_MASK: u32 = MAX_SPARSE_GRAM_SIZE as u32 - 1; + +#[derive(Clone, Copy, Debug)] +struct PosState { + index: u32, + value: u32, +} + +#[derive(Clone, Debug, Default)] +struct Queue { + idx_buf: [u32; MAX_SPARSE_GRAM_SIZE], + val_buf: [u32; MAX_SPARSE_GRAM_SIZE], + head: u32, + len: u32, +} + +impl Queue { + fn new() -> Self { + Self { + idx_buf: [0; MAX_SPARSE_GRAM_SIZE], + val_buf: [0; MAX_SPARSE_GRAM_SIZE], + head: 0, + len: 0, + } + } + + fn clear(&mut self) { + self.len = 0; + } + + fn is_empty(&self) -> bool { + self.len == 0 + } + + fn front_idx(&self) -> u32 { + // Callers only read the front boundary when the queue is non-empty; the empty case is + // handled earlier (e.g. via `state`'s early return) and must never reach here. + debug_assert!(!self.is_empty()); + self.idx_buf[self.head as usize] + } + + fn front_value(&self) -> u32 { + // Callers only read the front priority when the queue is non-empty; the empty case is + // handled earlier and must never reach here. + debug_assert!(!self.is_empty()); + self.val_buf[self.head as usize] + } + + fn pop_front_idx(&mut self) -> u32 { + debug_assert!(!self.is_empty()); + let first = self.idx_buf[self.head as usize]; + self.head = (self.head + 1) & RING_MASK; + self.len -= 1; + first + } + + fn push(&mut self, state: PosState) { + // Drop tail candidates whose priority exceeds the new one, keeping the deque monotone + // (nondecreasing priorities front-to-back). + while !self.is_empty() { + let slot = ((self.head + self.len - 1) & RING_MASK) as usize; + if self.val_buf[slot] <= state.value { + break; + } + self.len -= 1; + } + debug_assert!(self.len < MAX_SPARSE_GRAM_SIZE as u32); + let slot = ((self.head + self.len) & RING_MASK) as usize; + self.idx_buf[slot] = state.index; + self.val_buf[slot] = state.value; + self.len += 1; + } +} + +/// Streaming query n-gram state. +/// +/// This state accepts one character at a time, can be cloned while traversing an automaton, and +/// can emit grams incrementally. It uses the same compact bigram-priority model as indexing, but +/// emits the minimum number of grams that cover the input stream. Its state space is fully +/// represented by the content buffer and can be shrunk on demand when only part of the stream +/// needs to be retained. +#[derive(Clone)] +pub struct QueryGrams { + /// Queue of candidate boundaries (strictly increasing indices and nondecreasing priorities). + queue: Queue, + /// Active content packed into a `u64` (newest byte in the low byte). + content: u64, + /// Absolute index one past the last active byte in `content`. + content_end_idx: u32, + /// Rolling `H` value of the first byte of the next bigram to score. + h: u32, +} + +impl Eq for QueryGrams {} + +impl PartialEq for QueryGrams { + fn eq(&self, other: &Self) -> bool { + self.state() == other.state() + } +} + +impl Hash for QueryGrams { + fn hash(&self, state: &mut H) { + self.state().hash(state); + } +} + +impl Default for QueryGrams { + fn default() -> Self { + Self { + content: 0, + queue: Queue::new(), + content_end_idx: 0, + h: 0, + } + } +} + +impl QueryGrams { + /// Returns the canonical, position-independent fingerprint of the active window as + /// `(content_len, content)`. + /// + /// `content_len` is the number of active bytes and `content` is exactly those bytes packed into + /// a `u64` (newest in the low byte), with all higher bytes masked off. The active window starts + /// two bytes before the front boundary (the oldest candidate still able to start a gram) and + /// runs to the newest fed byte: the front boundary is a bigram, so both of its source bytes — + /// which determine its priority and the left half of the next bigram — must be retained for the + /// state to fully predict future grams. Everything before that has already been covered by an + /// emitted gram and can no longer influence future grams, so it is excluded; this makes the + /// fingerprint independent of the absolute stream position and of already-drained history. + /// [`PartialEq`], [`Eq`] and [`Hash`] are all defined in terms of it, so two states with + /// identical active windows compare and hash equal. + /// + /// When the queue is empty there is no front boundary, so the whole packed buffer is active: + /// `(0, 0)` for a fresh state and `(1, last_byte)` once a single trailing byte has been + /// retained (e.g. after `consume_first` collapses). The retained byte is part of the state + /// because it becomes the left half of the next bigram. + #[inline] + pub fn state(&self) -> (u32, u64) { + let begin = if self.queue.is_empty() { + 0 + } else { + // The front boundary is a bigram at index `front_idx`, spanning byte positions + // `front_idx - 2` and `front_idx - 1`; include both so its priority (and the next + // bigram's left byte) are captured. + self.queue.front_idx() - 2 + }; + let content_len = self.content_end_idx - begin; + debug_assert!(content_len < MAX_SPARSE_GRAM_SIZE as u32); + let mask = (1u64 << (content_len * 8)) - 1; + (content_len, self.content & mask) + } + + /// Returns the smallest active boundary priority. + pub fn min_priority(&self) -> u32 { + if self.queue.is_empty() { + 0 + } else { + self.queue.front_value() + } + } + + /// The character following a gram whose reported position is `end` — i.e. the byte at content + /// position `end` — if that character has already been fed into the active window. + /// + /// Returns `None` when `end` is at (or past) the newest fed byte — the following character has + /// not been fed yet, e.g. for the gram ending at the most recently appended byte, or for the + /// single trailing bigram of the whole input. Panics if the position lies outside the packed + /// window. A one-byte-lookahead wrapper can use its buffered byte for the `None` case. + #[inline] + fn follow_byte(&self, end: u32) -> Option { + if end >= self.content_end_idx { + None + } else { + let shift = self.content_end_idx - 1 - end; + debug_assert!(shift < MAX_SPARSE_GRAM_SIZE as u32); + Some((self.content >> (shift * 8)) as u8) + } + } + + fn extract_gram(&mut self, begin_index: u32, end_index: u32, consumer: &mut F) + where + F: FnMut(NGram, u32, Option), + { + debug_assert!(end_index >= begin_index); + debug_assert!(end_index <= self.content_end_idx); + + let len = (end_index - begin_index + 2) as usize; + let dist = self.content_end_idx - end_index; + let shifted = self.content >> (dist * 8); + // Report the position of the character just after the emitted ngram, along with that + // character itself when it has already been fed (see `follow_byte`). + let follow = self.follow_byte(end_index); + consumer(NGram::from_window_masked(shifted, len), end_index, follow); + } + + /// Appends a single character to the n-gram state. + /// + /// The character is index-folded using the `casefold` crate and may trigger one or more grams. + pub fn append_char(&mut self, c: char, consumer: F) + where + F: FnMut(NGram, u32, Option), + { + self.append_byte(index_fold_char(c), consumer); + } + + /// Appends a single already index-folded byte to the n-gram state. + /// + /// This is the byte-level counterpart of [`append_char`](Self::append_char): the caller has + /// already index-folded the character (e.g. buffered it as a byte), so no further folding is + /// applied. May trigger one or more grams. + pub fn append_byte(&mut self, right: u8, mut consumer: F) + where + F: FnMut(NGram, u32, Option), + { + let left = (self.content & 0xFF) as u8; + self.content_end_idx += 1; + self.content = (self.content << 8) | right as u64; + + let idx = self.content_end_idx; + + // Initialize rolling state from the first character; from then on each append consumes + // exactly one new character via the bigram path. + if idx == 1 { + self.h = bigram_h(right); + } else { + let (value, h_b) = bigram_priority_rolling(left, right, self.h); + self.h = h_b; + if !self.queue.is_empty() + && let priority = self.queue.front_value() + && value < priority + { + let mut begin = self.queue.pop_front_idx(); + while !self.queue.is_empty() && self.queue.front_value() == priority { + let end = self.queue.pop_front_idx(); + self.extract_gram(begin, end, &mut consumer); + begin = end; + } + self.queue.clear(); + self.queue.push(PosState { index: idx, value }); + self.extract_gram(begin, idx, &mut consumer); + } else { + self.queue.push(PosState { index: idx, value }); + } + if idx - self.queue.front_idx() + 2 >= MAX_SPARSE_GRAM_SIZE as u32 && self.queue.len > 1 + { + let begin = self.queue.pop_front_idx(); + let end = self.queue.front_idx(); + self.extract_gram(begin, end, &mut consumer); + } + } + } + + /// Flushes all buffered characters and emits remaining grams. + pub fn flush(mut self, mut consumer: F) + where + F: FnMut(NGram, u32, Option), + { + if self.content_end_idx == 2 { + self.extract_gram(2, 2, &mut consumer); + } else { + while self.queue.len > 1 { + let begin = self.queue.pop_front_idx(); + let end = self.queue.front_idx(); + self.extract_gram(begin, end, &mut consumer); + } + } + } + + /// Consumes and emits at most one queued gram (if available), shrinking state. + pub fn consume_first(&mut self, mut consumer: F) + where + F: FnMut(NGram, u32, Option), + { + if self.queue.len > 1 { + // Emit the gram spanning the first boundary to the next one, mirroring `flush`. + let begin = self.queue.pop_front_idx(); + let end = self.queue.front_idx(); + self.extract_gram(begin, end, &mut consumer); + } else if self.content_end_idx == 2 { + // Only a single bigram remains; emit it like `flush` does. + self.extract_gram(2, 2, &mut consumer); + } + + // The last boundary is a dangling endpoint that never becomes its own gram (just like + // `flush` stops at `queue.len > 1`). Once only that boundary (or nothing) is left, collapse + // to the retained last byte so a continuation matches a fresh run starting from that byte. + if self.queue.len <= 1 { + self.queue.clear(); + if self.content_end_idx > 1 { + // `self.h` already holds `bigram_h(last)` by invariant, so it needs no update. + let last = (self.content & 0xFF) as u8; + self.content = last as u64; + self.content_end_idx = 1; + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::collect_sparse_grams_deque; + use std::collections::hash_map::DefaultHasher; + use std::ops::Range; + + // Test intervals are stored as `Range` in bigram-boundary coordinates: a gram covering + // 0-based bytes `s..=s+L-1` is represented as `s+1..s+L-1`, the inclusive range of bigram + // boundaries (right-byte indices) it covers. The cover-chain and DP helpers treat `end` + // inclusively. + + fn query_intervals(input: &str) -> Vec> { + let mut q = QueryGrams::default(); + let mut out = Vec::new(); + for c in input.chars() { + q.append_char(c, |gram, end, _follow| { + let begin = end + 1 - gram.len() as u32; + out.push(begin..end - 1); + }); + } + q.flush(|gram, end, _follow| { + let begin = end + 1 - gram.len() as u32; + out.push(begin..end - 1); + }); + out + } + + // The full candidate set (every sparse gram the index extractor would emit) is exactly the + // pool the query cover must be drawn from, so reuse `collect_sparse_grams_deque` instead of + // re-deriving the boundary rule here. `idx` is the position just after the gram, matching the + // `Range` boundary convention above. + fn candidate_intervals(bytes: &[u8]) -> Vec> { + let mut out = Vec::new(); + collect_sparse_grams_deque(bytes, |gram, idx| { + let begin = idx + 1 - gram.len() as u32; + out.push(begin..idx - 1); + }); + out + } + + fn min_cover_size_query_semantics(m: u32, intervals: &[Range]) -> usize { + let inf = usize::MAX / 4; + let mut dp = vec![0usize; m as usize + 1]; + for pos in (1..m).rev() { + let mut best = inf; + for iv in intervals { + if iv.start <= pos && iv.end > pos { + let next = iv.end.min(m) as usize; + best = best.min(1 + dp[next]); + } + } + dp[pos as usize] = best; + } + dp[1] + } + + fn is_valid_query_cover_chain(m: u32, intervals: &[Range]) -> bool { + if m <= 1 { + return true; + } + let mut produced = 1u32; + for iv in intervals { + if iv.start > produced || iv.end <= produced { + return false; + } + produced = iv.end; + if produced >= m { + return true; + } + } + false + } + + #[derive(Clone, Copy, Debug)] + struct Rng64(u64); + + impl Rng64 { + fn new(seed: u64) -> Self { + Self(seed) + } + + fn next_u64(&mut self) -> u64 { + // xorshift64*: tiny deterministic RNG for tests. + let mut x = self.0; + x ^= x >> 12; + x ^= x << 25; + x ^= x >> 27; + self.0 = x; + x.wrapping_mul(0x2545_F491_4F6C_DD1D) + } + + fn gen_range(&mut self, upper: usize) -> usize { + (self.next_u64() % upper as u64) as usize + } + } + + #[test] + fn append_and_flush_emit_grams() { + let mut q = QueryGrams::default(); + for c in "hello world".chars() { + q.append_char(c, |_gram, _idx, _follow| {}); + } + let mut out = Vec::new(); + q.flush(|gram, idx, _follow| out.push((gram, idx))); + assert!(!out.is_empty()); + assert!( + out.iter() + .all(|(g, _)| (2..=MAX_SPARSE_GRAM_SIZE).contains(&g.len())) + ); + } + + #[test] + fn single_bigram_has_no_follow() { + // A two-character input yields exactly one gram, the trailing bigram, emitted at flush and + // followed by nothing. + let bytes: Vec = "ab".chars().map(crate::index_fold_char).collect(); + let mut q = QueryGrams::default(); + for &b in &bytes { + q.append_byte(b, |_gram, _end, _follow| {}); + } + let mut emitted = Vec::new(); + q.flush(|gram, end, follow| emitted.push((gram.len(), end, follow))); + assert_eq!(emitted, vec![(2usize, 2u32, None)]); + } + + #[test] + fn emitted_follow_is_the_actual_next_char() { + // Whenever a gram is emitted with a follow character, it must be the actual next byte of the + // input at the reported boundary (and a gram ending at the newest fed byte reports `None`). + for input in [ + "ab", + "abc", + "hello", + "querystream", + "aaaaaaaaaa", + "lardeee", + "sssdk", + ] { + let bytes: Vec = input.chars().map(crate::index_fold_char).collect(); + let n = bytes.len() as u32; + let mut q = QueryGrams::default(); + let mut check = |_gram: NGram, end: u32, follow: Option| { + if let Some(b) = follow { + assert!( + end < n, + "follow reported past input for {input:?}: end={end}" + ); + assert_eq!( + b, bytes[end as usize], + "wrong follow byte for {input:?} at end={end}" + ); + } + }; + for &b in &bytes { + q.append_byte(b, &mut check); + } + q.flush(&mut check); + } + } + + #[test] + fn state_tracks_active_window_each_step() { + // `state()` must report `(content_len, packed active window)` after every processed byte, + // where the active window is exactly the last `content_len` bytes of the input so far. This + // pins two things at once: + // * the pre-append snapshot equals the previous post-append snapshot (state only moves on + // append), and + // * no step collapses the window below its true length — the regression where an empty or + // front-coincident queue reported `(0, 0)` (dropping the retained trailing byte / the + // pending front bigram) would show up here as a wrong length at bytes 1, 3, 5, 7, ... + let input = b"hello world"; + // Active-window length after each processed byte (byte 0 = after the first byte). + let expected_len = [1u32, 2, 3, 2, 3, 2, 3, 2, 3, 4, 5]; + + let pack = |bytes: &[u8]| bytes.iter().fold(0u64, |acc, &b| (acc << 8) | b as u64); + + let mut q = QueryGrams::default(); + let mut prev = q.state(); + assert_eq!(prev, (0, 0), "a fresh state must be empty"); + + for (i, &byte) in input.iter().enumerate() { + // Before the append the state must be untouched since the last append. + assert_eq!( + q.state(), + prev, + "state moved without an append before byte {i}" + ); + + q.append_byte(byte, |_, _, _| {}); + + let len = expected_len[i]; + let window = &input[i + 1 - len as usize..=i]; + let after = q.state(); + assert_eq!( + after, + (len, pack(window)), + "unexpected state after byte {i} ({:?})", + char::from(byte), + ); + prev = after; + } + } + + #[test] + fn state_eq_and_hash_ignore_absolute_history() { + let mut a = QueryGrams::default(); + let mut b = QueryGrams::default(); + + for c in "abc".chars() { + a.append_char(c, |_gram, _idx, _follow| {}); + } + for c in "zabc".chars() { + b.append_char(c, |_gram, _idx, _follow| {}); + } + b.consume_first(|_gram, _idx, _follow| {}); + + // Only assert hash consistency when Eq says they are equivalent. + if a == b { + let mut ha = DefaultHasher::new(); + let mut hb = DefaultHasher::new(); + a.hash(&mut ha); + b.hash(&mut hb); + assert_eq!(ha.finish(), hb.finish()); + } + } + + #[test] + fn query_flush_is_minimum_cover_on_small_inputs() { + for input in [ + "abc", + "abcd", + "abcdef", + "hello", + "hello world", + "ababababab", + ] { + let bytes = input.as_bytes(); + if bytes.len() < 3 { + continue; + } + let produced = query_intervals(input); + let candidates = candidate_intervals(bytes); + let m = bytes.len() as u32 - 1; + + assert!( + is_valid_query_cover_chain(m, &produced), + "produced set is not a valid cover chain for {input:?}: {:?}", + produced + ); + let optimum = min_cover_size_query_semantics(m, &candidates); + assert_eq!( + produced.len(), + optimum, + "produced set is not minimum-size query cover for {input:?}; produced={:?}; optimum={optimum}", + produced + ); + } + } + + #[test] + fn query_lardeee_diagnostic() { + let input = "lardeee"; + let bytes = input.as_bytes(); + let produced = query_intervals(input); + let candidates = candidate_intervals(bytes); + let m = bytes.len() as u32 - 1; + + assert!( + is_valid_query_cover_chain(m, &produced), + "produced set is not a valid cover chain for {input:?}: {:?}", + produced + ); + + let optimum = min_cover_size_query_semantics(m, &candidates); + + // Diagnostic check: if this fails, query extraction emitted an interval the oracle does + // not consider legal under its candidate rules. + for iv in &produced { + assert!( + candidates + .iter() + .any(|c| c.start == iv.start && c.end == iv.end), + "produced interval not present in oracle candidates for {input:?}: {:?}; candidates={:?}", + iv, + candidates + ); + } + + // This currently captures the known mismatch that motivated the randomized check. + assert_eq!( + produced.len(), + optimum, + "minimum-cover mismatch for {input:?}; produced={:?}; optimum={optimum}; candidates={:?}", + produced, + candidates + ); + } + + #[test] + fn query_sssdk_diagnostic() { + let input = "sssdk"; + let bytes = input.as_bytes(); + let produced = query_intervals(input); + let candidates = candidate_intervals(bytes); + let m = bytes.len() as u32 - 1; + + assert!( + is_valid_query_cover_chain(m, &produced), + "produced set is not a valid cover chain for {input:?}: {:?}", + produced + ); + + let optimum = min_cover_size_query_semantics(m, &candidates); + + for iv in &produced { + assert!( + candidates + .iter() + .any(|c| c.start == iv.start && c.end == iv.end), + "produced interval not present in oracle candidates for {input:?}: {:?}; candidates={:?}", + iv, + candidates + ); + } + + assert_eq!( + produced.len(), + optimum, + "minimum-cover mismatch for {input:?}; produced={:?}; optimum={optimum}; candidates={:?}", + produced, + candidates + ); + } + + #[test] + fn query_flush_is_minimum_cover_on_randomized_inputs() { + let mut rng = Rng64::new(0xA5A5_0123_89AB_CDEF); + + for _ in 0..2000 { + let len = 3 + rng.gen_range(14); // [3, 16] + let mut bytes = vec![0u8; len]; + for b in &mut bytes { + *b = b'a' + rng.gen_range(26) as u8; // casefold-stable lowercase ASCII + } + let input = std::str::from_utf8(&bytes).expect("ASCII should be valid UTF-8"); + + let produced = query_intervals(input); + let candidates = candidate_intervals(&bytes); + let m = bytes.len() as u32 - 1; + + assert!( + is_valid_query_cover_chain(m, &produced), + "produced set is not a valid cover chain for randomized input {:?}: {:?}", + input, + produced + ); + assert!( + produced.iter().all(|iv| { + let gram_len = (iv.end - iv.start + 2) as usize; + (2..=MAX_SPARSE_GRAM_SIZE).contains(&gram_len) + }), + "produced set contains out-of-range gram length for randomized input {:?}: {:?}", + input, + produced + ); + let optimum = min_cover_size_query_semantics(m, &candidates); + assert_eq!( + produced.len(), + optimum, + "produced set is not minimum-size query cover for randomized input {:?}; produced={:?}; optimum={optimum}", + input, + produced + ); + } + } + + #[test] + fn query_consume_first_on_randomized_inputs() { + let mut rng = Rng64::new(0xC0DE_CAFE_1234_5678); + + for _ in 0..2000 { + let len = 4 + rng.gen_range(13); // [4, 16] + let mut bytes = vec![0u8; len]; + for b in &mut bytes { + *b = b'a' + rng.gen_range(26) as u8; + } + let input = std::str::from_utf8(&bytes).expect("ASCII should be valid UTF-8"); + + let split = 2 + rng.gen_range(len - 2); + let (prefix, suffix) = input.split_at(split); + let mut q = QueryGrams::default(); + + // Capture the full first-half cover: grams emitted eagerly while appending the prefix + // plus grams drained via `consume_first`. + let mut first_half = Vec::new(); + for c in prefix.chars() { + q.append_char(c, |gram, end, _follow| { + let begin = end + 1 - gram.len() as u32; + first_half.push(begin..end - 1); + }); + } + + while !q.queue.is_empty() { + q.consume_first(|gram, end, _follow| { + let begin = end + 1 - gram.len() as u32; + first_half.push(begin..end - 1); + }); + } + + assert!(first_half.iter().all(|iv| { + let gram_len = (iv.end - iv.start + 2) as usize; + (2..=MAX_SPARSE_GRAM_SIZE).contains(&gram_len) + })); + + // The first half must be an optimal (minimum-size) cover of the prefix. The DP optimum + // is only meaningful for inputs of length >= 3 (a 2-char input has no interior + // boundary, so the DP reports 0 while a single bigram is still emitted). + let prefix_bytes = prefix.as_bytes(); + if prefix_bytes.len() >= 3 { + let prefix_m = prefix_bytes.len() as u32 - 1; + let prefix_candidates = candidate_intervals(prefix_bytes); + assert!( + is_valid_query_cover_chain(prefix_m, &first_half), + "first half is not a valid cover chain for prefix {:?}: {:?}", + prefix, + first_half + ); + let prefix_optimum = min_cover_size_query_semantics(prefix_m, &prefix_candidates); + assert_eq!( + first_half.len(), + prefix_optimum, + "first half is not a minimum-size query cover for prefix {:?}; produced={:?}; optimum={prefix_optimum}", + prefix, + first_half + ); + } + + let retained = char::from((q.content & 0xFF) as u8); + let local_input = format!("{retained}{suffix}"); + + let mut remaining = Vec::new(); + for c in suffix.chars() { + q.append_char(c, |gram, end, _follow| { + let begin = end + 1 - gram.len() as u32; + remaining.push(begin..end - 1); + }); + } + q.flush(|gram, end, _follow| { + let begin = end + 1 - gram.len() as u32; + remaining.push(begin..end - 1); + }); + + assert_eq!( + remaining, + query_intervals(&local_input), + "post-drain continuation mismatch for randomized input {:?}; prefix={:?}; suffix={:?}; first_half={:?}; remaining={:?}; local_input={:?}", + input, + prefix, + suffix, + first_half, + remaining, + local_input + ); + + // The second half must be an optimal (minimum-size) cover of the retained tail + suffix. + let local_bytes = local_input.as_bytes(); + if local_bytes.len() >= 3 { + let local_m = local_bytes.len() as u32 - 1; + let local_candidates = candidate_intervals(local_bytes); + assert!( + is_valid_query_cover_chain(local_m, &remaining), + "second half is not a valid cover chain for {:?}: {:?}", + local_input, + remaining + ); + let local_optimum = min_cover_size_query_semantics(local_m, &local_candidates); + assert_eq!( + remaining.len(), + local_optimum, + "second half is not a minimum-size query cover for {:?}; produced={:?}; optimum={local_optimum}", + local_input, + remaining + ); + } + } + } +}