Merge pull request #142 from github/aneubeck/query-gram-bytes

sparse-ngrams: report each query gram's index-folded bytes

Author: aneubeck

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"

crates/sparse-ngrams/README.md

@@ -61,8 +61,11 @@ 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. A consumer that ignores
+the argument optimizes back to code identical to not reporting the bytes at all.
 
 ```rust
 use sparse_ngrams::{QueryGrams, NGram};
@@ -71,12 +74,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<u8>| {
+    q.append_char(c, |gram: NGram, _end: u32, _follow: Option<u8>, 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);
 });
 

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.

crates/sparse-ngrams/src/query.rs

@@ -94,6 +94,21 @@ 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>, &[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, 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).
@@ -202,26 +217,41 @@ impl QueryGrams {
 
     fn extract_gram<F>(&mut self, begin_index: u32, end_index: u32, consumer: &mut F)
     where
-        F: FnMut(NGram, u32, Option<u8>),
+        F: FnMut(NGram, u32, Option<u8>, &[u8]),
     {
         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);
+        // 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);
+        // `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.min(MAX_SPARSE_GRAM_SIZE)],
+        );
     }
 
     /// 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<F>(&mut self, c: char, consumer: F)
     where
-        F: FnMut(NGram, u32, Option<u8>),
+        F: FnMut(NGram, u32, Option<u8>, &[u8]),
     {
         self.append_byte(index_fold_char(c), consumer);
     }
@@ -233,7 +263,7 @@ impl QueryGrams {
     /// applied. May trigger one or more grams.
     pub fn append_byte<F>(&mut self, right: u8, mut consumer: F)
     where
-        F: FnMut(NGram, u32, Option<u8>),
+        F: FnMut(NGram, u32, Option<u8>, &[u8]),
     {
         let left = (self.content & 0xFF) as u8;
         self.content_end_idx += 1;
@@ -276,7 +306,7 @@ impl QueryGrams {
     /// Flushes all buffered characters and emits remaining grams.
     pub fn flush<F>(mut self, mut consumer: F)
     where
-        F: FnMut(NGram, u32, Option<u8>),
+        F: FnMut(NGram, u32, Option<u8>, &[u8]),
     {
         if self.content_end_idx == 2 {
             self.extract_gram(2, 2, &mut consumer);
@@ -292,7 +322,7 @@ impl QueryGrams {
     /// Consumes and emits at most one queued gram (if available), shrinking state.
     pub fn consume_first<F>(&mut self, mut consumer: F)
     where
-        F: FnMut(NGram, u32, Option<u8>),
+        F: FnMut(NGram, u32, Option<u8>, &[u8]),
     {
         if self.queue.len > 1 {
             // Emit the gram spanning the first boundary to the next one, mirroring `flush`.
@@ -343,12 +373,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 +458,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 +476,18 @@ mod tests {
         let bytes: Vec<u8> = "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 +500,17 @@ mod tests {
             let bytes: Vec<u8> = 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<u8>| {
+            let mut check = |gram: NGram, end: u32, follow: Option<u8>, 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 +557,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 +578,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 +765,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 +811,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 +866,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(),