// Frozen benchmark corpus for `sparse-ngrams` (see benchmarks/performance.rs).
// A snapshot of realistic source code, committed so the "large" benchmark input stays
// fixed and does not drift when the crate sources are edited. Do not regenerate.
//! Bigram priority model.
//!
//! Assigns a frequency-based priority to each byte pair, used by the sparse n-gram
//! extraction algorithm to decide where n-gram boundaries fall.
//!
//! Priorities used to be a full 256×256 `u16` table baked from a `bigrams.bin` frequency
//! ranking (~64kB in memory). They are now reconstructed from a compact *factored* model
//! (~8.5kB) tuned offline against that ranking. The ascii bigram `(a, b)` is scored as
//! `BIGRAM_H[a] + BIGRAM_H[b] + (code << BIGRAM_CODE_SHIFT) + 1`, where [`BIGRAM_H`] is a single
//! shared per-byte weight and `code` is a 4-bit (`0..=15`) per-bigram correction. Bigrams absent
//! from the training data carry `code == 0` and are not special-cased: they fall back to the bare
//! factored score, i.e. the model extrapolates a priority for them. The byte index `idx` is folded
//! into the low 16 bits so every bigram gets a *unique* priority, while a higher score still means
//! a more frequent bigram (~1.9% inversions vs. the exact ranking).
/// A casefolded indexable byte uses 7 bits for ascii characters; any non-ascii (unicode)
/// character is expected to have its high bit set. Only ascii bigrams are ever present, so
/// non-ascii characters always resolve to priority `0`. The model therefore only covers the 128
/// ascii values per character.
const BIGRAM_ALPHABET: usize = 128;
/// The per-bigram 4-bit correction code is scaled by `1 << BIGRAM_CODE_SHIFT` and added to the
/// shared per-byte weights. A plain shift replaces what used to be a lookup into a learned
/// 16-entry offset table, at a negligible accuracy cost (~1.9% vs ~1.87% inversions).
const BIGRAM_CODE_SHIFT: u32 = 8;
/// 4-bit correction code per ascii bigram, packed two codes per byte (the even index in the low
/// nibble). Scaled by `1 << BIGRAM_CODE_SHIFT` and added to the shared per-byte weights.
static BIGRAM_CODE: &[u8; BIGRAM_ALPHABET * BIGRAM_ALPHABET / 2] = include_bytes!("bigram_code.bin");
/// Shared per-byte weight. `BIGRAM_H[b]` contributes to the priority of every bigram containing
/// byte `b`; `3134` is the filler weight for bytes absent from the training data.
static BIGRAM_H: [u16; BIGRAM_ALPHABET] = [
3134, 0, 3134, 3134, 3134, 3134, 3134, 478, 3134, 3259, 3982, 3134, 541, 2332, 3134, 3134,
3134, 3134, 3134, 243, 3134, 3134, 3134, 3134, 3134, 3134, 3134, 671, 3134, 3134, 3134, 3134,
4433, 2411, 3784, 3065, 2238, 2305, 2643, 2950, 3427, 3280, 2982, 2601, 3306, 3261, 3491, 3565,
3296, 3431, 3319, 3142, 3121, 3142, 3122, 3090, 3103, 3063, 3153, 3072, 3065, 3208, 3087, 1901,
2506, 3134, 3134, 3134, 3134, 3134, 3134, 3134, 3134, 3134, 3134, 3134, 3134, 3134, 3134, 3134,
3134, 3134, 3134, 3134, 3134, 3134, 3134, 3134, 3134, 3134, 3134, 2626, 2519, 2739, 2382, 3238,
2578, 3803, 3467, 3773, 3752, 3989, 3500, 3524, 3152, 3571, 2639, 3249, 3542, 3437, 3647, 3554,
3503, 2503, 3741, 4069, 3831, 3564, 3162, 3137, 3082, 3080, 2860, 2794, 2396, 2852, 2450, 3134,
];
/// Reconstructs the priority of the ascii bigram `(a, b)`; see [`bigram_priority`]. This rolling
/// variant avoids re-loading `BIGRAM_H[a]`: consecutive bigrams overlap by one byte, so the caller
/// passes the `h_b` returned for the previous position as `h_a` and gets `BIGRAM_H[b]` back for the
/// next one. The `H` value is only used for ascii bytes; for a non-ascii byte the bigram is absent,
/// so the masked lookup (`b & 0x7f`) merely returns a value that the next step discards.
#[inline]
pub(crate) fn bigram_priority_rolling(a: u8, b: u8, h_a: u32) -> (u32, u32) {
let h_b = BIGRAM_H[(b & (BIGRAM_ALPHABET as u8 - 1)) as usize] as u32;
if (a | b) >= BIGRAM_ALPHABET as u8 {
return (0, h_b);
}
let idx = a as usize * BIGRAM_ALPHABET + b as usize;
let code = (BIGRAM_CODE[idx >> 1] >> ((idx & 1) * 4)) & 0xF;
// The 4-bit `code` is scaled by `1 << BIGRAM_CODE_SHIFT` (a plain shift in place of an offset
// table) and added to the shared per-byte weights. Absent bigrams carry `code == 0`, so they
// fall back to the bare factored score `h_a + h_b + 1`. A higher score still means a more
// frequent bigram; `base` fits in 16 bits, so the unique per-bigram `idx` in the low 16 bits
// keeps every priority unique.
let base = h_a + h_b + ((code as u32) << BIGRAM_CODE_SHIFT) + 1;
((base << 16) | idx as u32, h_b)
}
/// The `BIGRAM_H` weight of a single byte, used to seed [`bigram_priority_rolling`].
#[inline]
pub(crate) fn bigram_h(a: u8) -> u32 {
BIGRAM_H[(a & (BIGRAM_ALPHABET as u8 - 1)) as usize] as u32
}
/// Reconstructs the frequency-ranking priority of the ascii bigram `(a, b)`. Absent or non-ascii
/// bigrams resolve to `0`; present bigrams get a strictly positive, unique priority where a higher
/// value means a more frequent bigram. This priority is used to split strings into smaller
/// n-grams.
pub fn bigram_priority(a: u8, b: u8) -> u32 {
bigram_priority_rolling(a, b, bigram_h(a)).0
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn non_ascii_is_zero() {
assert_eq!(bigram_priority(0x80, b'a'), 0);
assert_eq!(bigram_priority(b'a', 0x80), 0);
assert_eq!(bigram_priority(0xff, 0xff), 0);
}
#[test]
fn ascii_bigrams_are_positive_and_unique() {
// Distinct ascii bigrams get distinct, strictly positive priorities (the `idx` in the low
// 16 bits guarantees uniqueness).
assert!(bigram_priority(b'a', b'b') > 0);
assert_ne!(bigram_priority(b'a', b'b'), bigram_priority(b'b', b'a'));
assert_ne!(bigram_priority(b'a', b'b'), bigram_priority(b'a', b'c'));
}
#[test]
fn rolling_matches_direct() {
// The rolling variant must reproduce the standalone `bigram_priority` for every ascii pair,
// and hand back `BIGRAM_H[b]` for the next step.
for a in 0u8..128 {
for b in 0u8..128 {
let (p, h_b) = bigram_priority_rolling(a, b, bigram_h(a));
assert_eq!(p, bigram_priority(a, b), "mismatch at ({a}, {b})");
assert_eq!(h_b, bigram_h(b), "h_b mismatch at ({a}, {b})");
}
}
}
}
//! Compact n-gram representation.
//!
//! An [`NGram`] packs a substring's byte length and a payload into the low **27 bits** of a
//! `u32` (the top 5 bits are always zero):
//!
//! ```text
//! bit 31 27 24 0
//! +-----------------+-----------+-----------+
//! | 00000 (unused) | len - 2 | payload |
//! +-----------------+-----------+-----------+
//! 5 bits 3 bits 24 bits
//! ```
//!
//! * **Length** (`len - 2`, bits 24..27): substring byte-lengths range from 2 (bigrams) to
//! [`MAX_SPARSE_GRAM_SIZE`] (8), so biasing by 2 fits the 7 possible values into 3 bits.
//! * **Payload** (bits 0..24): for substrings of at most 3 bytes the bytes are packed
//! losslessly (left-aligned, so distinct short grams never collide); longer substrings are
//! hashed down to 24 bits with a multiplicative hash.
//!
//! Because the length lives in its own field, an `NGram` of one size never collides with an
//! `NGram` of another size. The packed value is finally run through a bijective [`mix27`]
//! permutation so the most-significant bits (which callers may use for bucketing or sorting) are
//! well distributed even though the packed value is highly structured.
use std::fmt::{self, Write as _};
use crate::MAX_SPARSE_GRAM_SIZE;
/// Odd multiplicative constant (the golden-ratio / Fibonacci hashing constant) used to hash grams
/// longer than 3 bytes down to the 24-bit payload.
const MULTIPLICATIVE_HASH: u64 = 0x9E37_79B9_7F4A_7C15;
/// A compact n-gram identifier. See the module-level documentation for the bit layout.
///
/// Note: we could store n-grams up to length 8 verbatim in a `u64`. However, that would explode
/// the number of distinct keys in a search dictionary. For that reason we compress n-grams into a
/// `u32`, which puts a more reasonable upper bound on the number of dictionary keys.
///
/// Note: by storing the length explicitly, we ensure that only n-grams of the same length can
/// collide. This is important because there are exponentially more long n-grams than short ones.
/// At the same time, longer n-grams occur less frequently, so colliding long n-grams won't
/// increase the false-positive rate too much.
///
/// # Construction
///
/// Use [`NGram::from_bytes`] for one-off hashing, or the rolling 8-byte window helper inside the
/// extraction loop for amortised O(1) computation per n-gram.
#[derive(Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Default)]
#[repr(transparent)]
pub struct NGram(pub(crate) u32);
impl NGram {
/// Smallest indexed gram length (bigrams); subtracted from the stored length so the biased
/// value fits in [`Self::LEN_BITS`] bits.
const LEN_BIAS: u32 = 2;
/// Number of bits used to encode the biased length.
const LEN_BITS: u32 = 3;
/// Mask selecting the biased-length field (once shifted down).
const LEN_MASK: u32 = (1 << Self::LEN_BITS) - 1;
/// Number of low bits used by the payload; the length sits just above it.
const PAYLOAD_BITS: u32 = 24;
/// Mask selecting the payload field.
const PAYLOAD_MASK: u32 = (1 << Self::PAYLOAD_BITS) - 1;
/// Total number of significant bits in the packed representation (the top 5 bits of the `u32`
/// are always zero).
pub(crate) const BITS: u32 = Self::PAYLOAD_BITS + Self::LEN_BITS;
/// Mask selecting the significant [`Self::BITS`] bits.
pub(crate) const MASK: u32 = (1 << Self::BITS) - 1;
/// Build an `NGram` by hashing the given byte slice from scratch.
///
/// # Panics
///
/// In debug builds, panics if `src.len()` is not in `2..=`[`MAX_SPARSE_GRAM_SIZE`].
pub fn from_bytes(src: &[u8]) -> Self {
debug_assert!(
(Self::LEN_BIAS as usize..=MAX_SPARSE_GRAM_SIZE).contains(&src.len()),
"ngram length {} out of range [{}, {}]",
src.len(),
Self::LEN_BIAS,
MAX_SPARSE_GRAM_SIZE,
);
// 24-bit payload: short grams are packed losslessly, longer ones hashed.
let payload = if src.len() <= 3 {
// Pack the 2-3 bytes into the low 24 bits, most-significant byte first.
let mut p = 0u32;
for &byte in src {
p = (p << 8) | byte as u32;
}
p
} else {
// Grams here are 4..=MAX_SPARSE_GRAM_SIZE (8) bytes, so they fit in a single u64. A
// multiplicative hash is much cheaper than a per-byte loop, and the top bits of the
// product mix in every input byte. `from_le_bytes` keeps the result independent of
// host endianness, and the gram length lives in its own field so the payload hash
// needn't encode it.
let mut buf = [0u8; 8];
buf[..src.len()].copy_from_slice(src);
let product = u64::from_le_bytes(buf).wrapping_mul(MULTIPLICATIVE_HASH);
(product >> (u64::BITS - Self::PAYLOAD_BITS)) as u32
};
Self::pack(src.len(), payload)
}
/// Builds an `NGram` from a big-endian packing of its bytes: the `len` gram bytes occupy the
/// most-significant bytes of `value` (the first gram byte in the top byte) and the low
/// `8 - len` bytes are zero. This is the form the extraction loop's rolling 8-byte window
/// produces, so it can construct grams without re-reading them from a slice. It returns exactly
/// the same value as [`from_bytes`](Self::from_bytes) would for the same gram (see the
/// `from_window_matches_from_bytes` test), so the two paths stay interchangeable.
#[inline]
pub(crate) fn from_window(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,
);
// 24-bit payload: short grams are packed losslessly, longer ones hashed.
let payload = if len <= 3 {
// The bytes sit in the top `len` bytes, most-significant byte first; shifting them down
// to the low bits reproduces `from_bytes`'s lossless packing.
(value >> (u64::BITS - len as u32 * 8)) as u32
} else {
// `swap_bytes` turns the big-endian window into the little-endian byte order that
// `from_bytes` feeds to the multiplicative hash, so both paths agree bit-for-bit.
let product = value.swap_bytes().wrapping_mul(MULTIPLICATIVE_HASH);
(product >> (u64::BITS - Self::PAYLOAD_BITS)) as u32
};
Self::pack(len, payload)
}
/// 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.
#[inline]
fn pack(len: usize, payload: u32) -> Self {
let packed = ((len as u32 - Self::LEN_BIAS) << Self::PAYLOAD_BITS) | payload;
Self(mix27(packed))
}
/// The byte length of the n-gram.
#[inline]
pub fn len(&self) -> usize {
// The length lives in the *packed* value; unmix the stored value first.
let packed = unmix27(self.0);
(((packed >> Self::PAYLOAD_BITS) & Self::LEN_MASK) + Self::LEN_BIAS) as usize
}
/// Whether this represents an empty gram. Valid n-grams are always at least 2 bytes long, so
/// this only holds for a default-constructed placeholder.
#[inline]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
/// The raw packed `u32`. This is an opaque, well-distributed identifier suitable as a hash-map
/// or hash-set key.
#[inline]
pub fn as_u32(&self) -> u32 {
self.0
}
}
/// A bijective 27-bit mixing permutation. The packed gram value is highly structured — the length
/// sits in the top 3 bits and short grams carry raw ASCII bytes — which would make the
/// most-significant bits badly skewed. This xorshift-multiply finalizer, restricted to 27 bits,
/// spreads entropy across all bits while remaining a bijection on `[0, 2^27)` (each step is
/// invertible: xorshifts are triangular GF(2) maps, and multiplication by an odd constant is a
/// unit modulo `2^27`), so distinct grams stay distinct.
fn mix27(mut x: u32) -> u32 {
debug_assert!(x <= NGram::MASK, "mix27 input must be a 27-bit value");
x ^= x >> 15;
x = x.wrapping_mul(0x2c1b_3c6d) & NGram::MASK;
x ^= x >> 12;
x = x.wrapping_mul(0x297a_2d39) & NGram::MASK;
x ^= x >> 15;
x
}
/// Inverse of [`mix27`]: recovers the packed (length + payload) value from the stored value. Each
/// step undoes the corresponding `mix27` step in reverse: the `>> 15` xorshifts are self-inverse
/// (since `2 * 15 >= 27`), the `>> 12` xorshift is undone by the doubling `>> 12` then `>> 24`, and