From 7432a174e58a3658381a2d9e3693d949e767497f Mon Sep 17 00:00:00 2001 From: Alexander Neubeck Date: Thu, 30 Jul 2026 17:45:43 +0200 Subject: [PATCH] sparse-ngrams: report each query gram's index-folded bytes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `QueryGrams` emitted only the compact `NGram` key. Callers that want to show a gram — tracing, snapshot tests, debugging output — cannot recover the text from it: grams longer than 3 bytes are hashed into the 24-bit payload, so `Debug` degrades to `NGram('0x20baba', len=4)`. The state already holds the bytes: `extract_gram` builds the key from the packed content window. Align that window once and hand the consumer a `&[u8]` view of it, so consumers become `FnMut(NGram, u32, Option, &[u8])`. The slice borrows a stack array live only for the call, so nothing is allocated per gram. `NGram::from_window_masked` goes away: aligning the window in `extract_gram` (which the byte slice needs anyway) makes the masking redundant, since the left shift already drops the high bytes. Breaking change for the four `QueryGrams` emitters; the indexing entrypoints (`collect_sparse_grams*`) are untouched, as their callers already hold the content slice. Version bumped to 0.4.0. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 62aac6cc-77fd-42e4-9d10-27ab508f4062 --- crates/sparse-ngrams/Cargo.toml | 2 +- crates/sparse-ngrams/README.md | 11 ++-- crates/sparse-ngrams/src/ngram.rs | 17 ------- crates/sparse-ngrams/src/query.rs | 84 ++++++++++++++++++++++--------- 4 files changed, 68 insertions(+), 46 deletions(-) diff --git a/crates/sparse-ngrams/Cargo.toml b/crates/sparse-ngrams/Cargo.toml index 2a0bd5c..5fd3603 100644 --- a/crates/sparse-ngrams/Cargo.toml +++ b/crates/sparse-ngrams/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "sparse-ngrams" -version = "0.3.0" +version = "0.4.0" edition = "2024" description = "Fast sparse n-gram extraction from byte slices." repository = "https://github.com/github/rust-gems" diff --git a/crates/sparse-ngrams/README.md b/crates/sparse-ngrams/README.md index 9358b14..9527383 100644 --- a/crates/sparse-ngrams/README.md +++ b/crates/sparse-ngrams/README.md @@ -61,8 +61,10 @@ incrementally as the user types. `QueryGrams` is a streaming state machine for e 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). +The consumer receives `(gram, end, follow, bytes)`: the n-gram, the position of the character just +after it, that following byte when it has already been fed (`None` at the current stream end), and +the gram's index-folded bytes. `bytes` borrows a stack buffer valid only for the duration of the +call — copy it if you need to keep it — so nothing is allocated per gram. ```rust use sparse_ngrams::{QueryGrams, NGram}; @@ -71,12 +73,13 @@ 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| { + q.append_char(c, |gram: NGram, _end: u32, _follow: Option, bytes: &[u8]| { + assert_eq!(gram, NGram::from_bytes(bytes)); grams.push(gram); }); } // Drain the remaining tail grams. -q.flush(|gram: NGram, _end, _follow| { +q.flush(|gram: NGram, _end, _follow, _bytes| { grams.push(gram); }); diff --git a/crates/sparse-ngrams/src/ngram.rs b/crates/sparse-ngrams/src/ngram.rs index 87fad82..413e4ec 100644 --- a/crates/sparse-ngrams/src/ngram.rs +++ b/crates/sparse-ngrams/src/ngram.rs @@ -131,23 +131,6 @@ 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 index f7d3a64..4c7d8b8 100644 --- a/crates/sparse-ngrams/src/query.rs +++ b/crates/sparse-ngrams/src/query.rs @@ -94,6 +94,20 @@ impl Queue { /// 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. +/// +/// # Consumer callback +/// +/// Every emitting method ([`append_char`](Self::append_char), [`append_byte`](Self::append_byte), +/// [`flush`](Self::flush), [`consume_first`](Self::consume_first)) takes a consumer of shape +/// `FnMut(NGram, u32, Option, &[u8])`, called once per emitted gram with: +/// +/// * `gram` — the compact [`NGram`] identifier to look up in an index. +/// * `end` — the position of the character just after the gram in the fed stream. +/// * `follow` — that following (index-folded) byte, when it has already been fed; `None` at the +/// current stream end. +/// * `bytes` — the gram's index-folded bytes, in reading order. They are borrowed from a stack +/// buffer that is only valid for the duration of the call, so a consumer that needs to keep them +/// must copy them. Nothing is allocated per gram, so this stays cheap on the hot path. #[derive(Clone)] pub struct QueryGrams { /// Queue of candidate boundaries (strictly increasing indices and nondecreasing priorities). @@ -202,7 +216,7 @@ impl QueryGrams { fn extract_gram(&mut self, begin_index: u32, end_index: u32, consumer: &mut F) where - F: FnMut(NGram, u32, Option), + F: FnMut(NGram, u32, Option, &[u8]), { debug_assert!(end_index >= begin_index); debug_assert!(end_index <= self.content_end_idx); @@ -210,18 +224,29 @@ impl QueryGrams { let len = (end_index - begin_index + 2) as usize; let dist = self.content_end_idx - end_index; let shifted = self.content >> (dist * 8); + // The gram is the low `len` bytes of `shifted`; shifting them up into the most-significant + // bytes gives both the big-endian layout `NGram::from_window` consumes and, via + // `to_be_bytes`, the gram's bytes in reading order in the first `len` slots. + let aligned = shifted << ((MAX_SPARSE_GRAM_SIZE - len) * 8); + let bytes = aligned.to_be_bytes(); // 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); + consumer( + NGram::from_window(aligned, len), + end_index, + follow, + &bytes[..len], + ); } /// 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. + /// See the [type-level docs](Self#consumer-callback) for the consumer arguments. pub fn append_char(&mut self, c: char, consumer: F) where - F: FnMut(NGram, u32, Option), + F: FnMut(NGram, u32, Option, &[u8]), { self.append_byte(index_fold_char(c), consumer); } @@ -233,7 +258,7 @@ impl QueryGrams { /// applied. May trigger one or more grams. pub fn append_byte(&mut self, right: u8, mut consumer: F) where - F: FnMut(NGram, u32, Option), + F: FnMut(NGram, u32, Option, &[u8]), { let left = (self.content & 0xFF) as u8; self.content_end_idx += 1; @@ -276,7 +301,7 @@ impl QueryGrams { /// Flushes all buffered characters and emits remaining grams. pub fn flush(mut self, mut consumer: F) where - F: FnMut(NGram, u32, Option), + F: FnMut(NGram, u32, Option, &[u8]), { if self.content_end_idx == 2 { self.extract_gram(2, 2, &mut consumer); @@ -292,7 +317,7 @@ impl QueryGrams { /// 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), + F: FnMut(NGram, u32, Option, &[u8]), { if self.queue.len > 1 { // Emit the gram spanning the first boundary to the next one, mirroring `flush`. @@ -343,12 +368,12 @@ mod tests { let mut q = QueryGrams::default(); let mut out = Vec::new(); for c in input.chars() { - q.append_char(c, |gram, end, _follow| { + q.append_char(c, |gram, end, _follow, _bytes| { let begin = end + 1 - gram.len() as u32; out.push(begin..end - 1); }); } - q.flush(|gram, end, _follow| { + q.flush(|gram, end, _follow, _bytes| { let begin = end + 1 - gram.len() as u32; out.push(begin..end - 1); }); @@ -428,10 +453,10 @@ mod tests { fn append_and_flush_emit_grams() { let mut q = QueryGrams::default(); for c in "hello world".chars() { - q.append_char(c, |_gram, _idx, _follow| {}); + q.append_char(c, |_gram, _idx, _follow, _bytes| {}); } let mut out = Vec::new(); - q.flush(|gram, idx, _follow| out.push((gram, idx))); + q.flush(|gram, idx, _follow, _bytes| out.push((gram, idx))); assert!(!out.is_empty()); assert!( out.iter() @@ -446,17 +471,18 @@ mod tests { 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| {}); + q.append_byte(b, |_gram, _end, _follow, _bytes| {}); } let mut emitted = Vec::new(); - q.flush(|gram, end, follow| emitted.push((gram.len(), end, follow))); + q.flush(|gram, end, follow, _bytes| emitted.push((gram.len(), end, follow))); assert_eq!(emitted, vec![(2usize, 2u32, None)]); } #[test] - fn emitted_follow_is_the_actual_next_char() { + fn emitted_follow_and_bytes_match_the_input() { // 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`). + // The emitted byte slice must likewise be the input slice the gram was built from. for input in [ "ab", "abc", @@ -469,7 +495,17 @@ mod tests { 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| { + let mut check = |gram: NGram, end: u32, follow: Option, gram_bytes: &[u8]| { + assert_eq!( + gram_bytes, + &bytes[end as usize - gram.len()..end as usize], + "wrong gram bytes for {input:?} at end={end}" + ); + assert_eq!( + gram, + NGram::from_bytes(gram_bytes), + "gram bytes disagree with the emitted key for {input:?} at end={end}" + ); if let Some(b) = follow { assert!( end < n, @@ -516,7 +552,7 @@ mod tests { "state moved without an append before byte {i}" ); - q.append_byte(byte, |_, _, _| {}); + q.append_byte(byte, |_, _, _, _| {}); let len = expected_len[i]; let window = &input[i + 1 - len as usize..=i]; @@ -537,12 +573,12 @@ mod tests { let mut b = QueryGrams::default(); for c in "abc".chars() { - a.append_char(c, |_gram, _idx, _follow| {}); + a.append_char(c, |_gram, _idx, _follow, _bytes| {}); } for c in "zabc".chars() { - b.append_char(c, |_gram, _idx, _follow| {}); + b.append_char(c, |_gram, _idx, _follow, _bytes| {}); } - b.consume_first(|_gram, _idx, _follow| {}); + b.consume_first(|_gram, _idx, _follow, _bytes| {}); // Only assert hash consistency when Eq says they are equivalent. if a == b { @@ -724,14 +760,14 @@ mod tests { // plus grams drained via `consume_first`. let mut first_half = Vec::new(); for c in prefix.chars() { - q.append_char(c, |gram, end, _follow| { + q.append_char(c, |gram, end, _follow, _bytes| { 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| { + q.consume_first(|gram, end, _follow, _bytes| { let begin = end + 1 - gram.len() as u32; first_half.push(begin..end - 1); }); @@ -770,12 +806,12 @@ mod tests { let mut remaining = Vec::new(); for c in suffix.chars() { - q.append_char(c, |gram, end, _follow| { + q.append_char(c, |gram, end, _follow, _bytes| { let begin = end + 1 - gram.len() as u32; remaining.push(begin..end - 1); }); } - q.flush(|gram, end, _follow| { + q.flush(|gram, end, _follow, _bytes| { let begin = end + 1 - gram.len() as u32; remaining.push(begin..end - 1); }); @@ -825,12 +861,12 @@ mod tests { for input in ["abc", "abcdef", "hello world", "ababababab", "mississippi"] { let mut q = QueryGrams::default(); for c in input.chars() { - q.append_char(c, |_gram, _idx, _follow| {}); + q.append_char(c, |_gram, _idx, _follow, _bytes| {}); } // Far more calls than there are bytes: the state must be at the default fixpoint well // before this loop ends, and further calls must keep it there. for _ in 0..(input.len() + 8) { - q.consume_first(|_gram, _idx, _follow| {}); + q.consume_first(|_gram, _idx, _follow, _bytes| {}); } assert_eq!( q.state(),