implement extraction of query ngrams

crates/sparse-ngrams/Cargo.toml

@@ -16,6 +16,7 @@ bench = false
 serde = ["dep:serde"]
 
 [dependencies]
+casefold = { version = "0.1", path = "../casefold" }
 serde = { version = "1", features = ["derive"], optional = true }
 
 [[bench]]

crates/sparse-ngrams/README.md

@@ -46,7 +46,7 @@ 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.
 });

crates/sparse-ngrams/benchmarks/performance.rs

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

@@ -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,10 @@ 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 inclusive end index of `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 +72,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 +109,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);
 
         // 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 +124,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,
+            );
             let bval = val_buf[slot];
             if bval < value {
                 break;
@@ -145,12 +146,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 inclusive end index of
+/// `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 +170,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);
 
         // 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 +186,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);
                 running_min = v_p;
             }
         }
@@ -197,9 +199,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,9 +40,11 @@
 
 mod extract;
 mod ngram;
+mod query;
 mod table;
 
 pub use ngram::NGram;
+pub use query::QueryGrams;
 pub use table::bigram_priority;
 
 /// Maximum length (in bytes) of a sparse n-gram.

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.