implement extraction of query ngrams

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]]

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<u8>| {
+        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):

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;
                 });

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<NGram> {
     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<NGram> {
 
 /// 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<NGram> {
 /// 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<NGram> {
+    fn collect_to_vec(run: impl FnOnce(&mut dyn FnMut(NGram, u32))) -> Vec<NGram> {
         let mut out = Vec::new();
-        run(&mut |gram| out.push(gram));
+        run(&mut |gram, _idx| out.push(gram));
         out
     }
 

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;
 

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.