Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
aneubeck wants to merge 1 commit into
base: main
Choose a base branch
from
+68 −46
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion crates/sparse-ngrams/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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"
Expand Down
11 changes: 7 additions & 4 deletions crates/sparse-ngrams/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand All @@ -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<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);
});

Expand Down
17 changes: 0 additions & 17 deletions crates/sparse-ngrams/src/ngram.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
84 changes: 60 additions & 24 deletions crates/sparse-ngrams/src/query.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>, &[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).
Expand Down Expand Up @@ -202,26 +216,37 @@ 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);
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<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);
}
Expand All @@ -233,7 +258,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;
Expand Down Expand Up @@ -276,7 +301,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);
Expand All @@ -292,7 +317,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`.
Expand Down Expand Up @@ -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);
});
Expand Down Expand Up @@ -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()
Expand All @@ -446,17 +471,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",
Expand All @@ -469,7 +495,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,
Expand Down Expand Up @@ -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];
Expand All @@ -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 {
Expand Down Expand Up @@ -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);
});
Expand Down Expand Up @@ -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);
});
Expand Down Expand Up @@ -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(),
Expand Down