From 7432a174e58a3658381a2d9e3693d949e767497f Mon Sep 17 00:00:00 2001 From: Alexander Neubeck Date: Thu, 30 Jul 2026 17:45:43 +0200 Subject: [PATCH 1/2] 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(), From 5009143b57ab875f1ebed72afe7d1e83a70443e4 Mon Sep 17 00:00:00 2001 From: Alexander Neubeck Date: Fri, 31 Jul 2026 08:05:32 +0200 Subject: [PATCH 2/2] Make the reported gram bytes free for consumers that ignore them MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Slicing the stack buffer with a runtime `len` leaves a bounds-check panic path in the generated code. That path is an observable side effect, so the optimizer cannot prove the buffer is dead even when the inlined consumer never touches it, and the byte-swap plus stack traffic survive. Measured on aarch64 (`--release`, an `append_byte` loop whose consumer only sums `gram.as_u32()`): 243 instructions with the plain slice against 229 without reporting bytes at all — a 14-instruction tax on every consumer, paid whether or not the bytes are wanted. Clamping the slice length to `MAX_SPARSE_GRAM_SIZE` makes the bound statically provable, so the panic path disappears and the whole buffer becomes dead code when unused: back to 229 instructions, byte-for-byte the pre-change codegen. Consumers that do read the bytes get cheaper code too (247 -> 240), since the eliminated check was on the hot path. The clamp never changes the slice: `len` is always in `2..=MAX_SPARSE_GRAM_SIZE`, which `from_window` already debug-asserts. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 62aac6cc-77fd-42e4-9d10-27ab508f4062 --- crates/sparse-ngrams/README.md | 3 ++- crates/sparse-ngrams/src/query.rs | 9 +++++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/crates/sparse-ngrams/README.md b/crates/sparse-ngrams/README.md index 9527383..f6ddbb8 100644 --- a/crates/sparse-ngrams/README.md +++ b/crates/sparse-ngrams/README.md @@ -64,7 +64,8 @@ determined, and can be `flush`ed to drain the tail. 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. +call — copy it if you need to keep it — so nothing is allocated per gram. A consumer that ignores +the argument optimizes back to code identical to not reporting the bytes at all. ```rust use sparse_ngrams::{QueryGrams, NGram}; diff --git a/crates/sparse-ngrams/src/query.rs b/crates/sparse-ngrams/src/query.rs index 4c7d8b8..14184c1 100644 --- a/crates/sparse-ngrams/src/query.rs +++ b/crates/sparse-ngrams/src/query.rs @@ -107,7 +107,8 @@ impl Queue { /// 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. +/// must copy them. Nothing is allocated per gram, and a consumer that ignores the argument +/// optimizes back to code identical to not reporting the bytes at all. #[derive(Clone)] pub struct QueryGrams { /// Queue of candidate boundaries (strictly increasing indices and nondecreasing priorities). @@ -232,11 +233,15 @@ impl QueryGrams { // 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); + // `len` is always in `2..=MAX_SPARSE_GRAM_SIZE` (`from_window` debug-asserts it), so the + // clamp never changes the slice. It is what makes the bound *statically* provable: without + // it the slicing keeps a bounds-check panic path alive, and that observable side effect + // stops the optimizer from eliminating the byte buffer for consumers that ignore it. consumer( NGram::from_window(aligned, len), end_index, follow, - &bytes[..len], + &bytes[..len.min(MAX_SPARSE_GRAM_SIZE)], ); }