use alloc::vec;
use alloc::vec::Vec;
use super::suffix_array::suffix_array;
use crate::encoding::fastpath::{
FastpathKernel, dispatch_common_prefix_len_ptr_with_kernel, select_kernel,
};
const MIN_RATIO: u32 = 4;
const MIN_MATCH_LENGTH: usize = 7;
const LENGTH_LIMIT: usize = 64;
const DICT_LIST_SIZE_DEFAULT: usize = 10_000;
const NOISE_LENGTH: usize = 32;
const MAX_SAMPLES_SIZE: usize = 2000 << 20;
const CONTENT_SIZE_MIN: usize = 128;
pub(crate) const MIN_SAMPLES_SIZE: usize = CONTENT_SIZE_MIN * MIN_RATIO as usize;
pub const DEFAULT_SELECTIVITY: u32 = 9;
pub(crate) const DICT_SIZE_MIN: usize = super::MIN_TRAINED_DICT_SIZE;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum TooSmall {
Dictionary,
Corpus,
Content,
}
#[derive(Debug, Clone, Copy, Default)]
struct DictItem {
pos: u32,
length: u32,
savings: u32,
}
fn noise_band() -> [u8; NOISE_LENGTH] {
let mut band = [0u8; NOISE_LENGTH];
let mut acc: u32 = 2_654_435_761;
for byte in &mut band {
acc = acc.wrapping_mul(2_246_822_519);
*byte = (acc >> 21) as u8;
}
band
}
struct Corpus<'a> {
samples: &'a [u8],
noise: [u8; NOISE_LENGTH],
kernel: FastpathKernel,
}
impl<'a> Corpus<'a> {
fn new(samples: &'a [u8]) -> Self {
Self {
samples,
noise: noise_band(),
kernel: select_kernel(),
}
}
#[inline]
fn len(&self) -> usize {
self.samples.len() + NOISE_LENGTH
}
#[inline]
fn byte(&self, at: usize) -> Option<u8> {
match self.samples.get(at) {
Some(&byte) => Some(byte),
None => self.noise.get(at - self.samples.len()).copied(),
}
}
#[inline(always)]
fn read<const N: usize>(&self, at: usize) -> Option<[u8; N]> {
if let Some(bytes) = self.samples.get(at..at + N) {
return Some(bytes.try_into().expect("N bytes"));
}
self.read_across(at)
}
#[cold]
#[inline(never)]
fn read_across<const N: usize>(&self, at: usize) -> Option<[u8; N]> {
let mut out = [0u8; N];
for (i, slot) in out.iter_mut().enumerate() {
*slot = self.byte(at + i)?;
}
Some(out)
}
#[inline]
fn read16(&self, at: usize) -> Option<u16> {
self.read::<2>(at).map(u16::from_le_bytes)
}
#[inline]
fn read64(&self, at: usize) -> Option<u64> {
self.read::<8>(at).map(u64::from_le_bytes)
}
#[inline]
fn common(&self, a: usize, b: usize) -> usize {
let samples = self.samples;
debug_assert!(a.max(b) <= samples.len());
let in_samples = samples.len() - a.max(b);
let base = samples.as_ptr();
let n = unsafe {
dispatch_common_prefix_len_ptr_with_kernel(
self.kernel,
base.add(a),
base.add(b),
in_samples,
)
};
if n < in_samples {
return n;
}
self.common_tail(a, b, n)
}
#[cold]
#[inline(never)]
fn common_tail(&self, a: usize, b: usize, mut n: usize) -> usize {
let Some(limit) = self.len().checked_sub(a.max(b)) else {
return 0;
};
while n < limit && self.byte(a + n) == self.byte(b + n) {
n += 1;
}
n
}
fn copy_to(&self, from: usize, out: &mut [u8]) {
if let Some(bytes) = self.samples.get(from..from + out.len()) {
out.copy_from_slice(bytes);
return;
}
for (i, slot) in out.iter_mut().enumerate() {
*slot = self
.byte(from + i)
.expect("a segment lies inside the corpus");
}
}
}
const FAR_AHEAD: usize = 16;
const NEAR_AHEAD: usize = 8;
struct Suffixes {
sa: Vec<u32>,
noise: usize,
}
impl Suffixes {
fn new(sa: Vec<u32>, len: usize) -> Self {
Self { sa, noise: len }
}
#[inline]
fn ranks(&self) -> i64 {
self.sa.len() as i64
}
#[inline]
fn at(&self, at: i64) -> usize {
match self.sa.get(at as u64 as usize) {
Some(&pos) => pos as usize,
None => self.noise,
}
}
#[inline(always)]
fn prefetch_for(&self, rank: &[u32], samples: &[u8], cursor: usize) {
use crate::decoding::prefetch::prefetch_l1_at;
let len = self.sa.len();
if let Some(&far) = rank.get(cursor + FAR_AHEAD) {
prefetch_l1_at(self.sa.as_ptr().wrapping_add(far as usize).cast());
}
if let Some(&near) = rank.get(cursor + NEAR_AHEAD) {
let near = near as usize;
let below = near.wrapping_sub(1).min(len - 1);
let above = (near + 1).min(len - 1);
let (below, above) = unsafe {
(
*self.sa.get_unchecked(below) as usize,
*self.sa.get_unchecked(above) as usize,
)
};
prefetch_l1_at(samples.as_ptr().wrapping_add(below));
prefetch_l1_at(samples.as_ptr().wrapping_add(above));
}
}
}
pub(crate) fn train_legacy_raw(
samples: &[u8],
sample_sizes: &[usize],
dict_size: usize,
selectivity: u32,
) -> Result<Vec<u8>, TooSmall> {
let total: usize = sample_sizes.iter().sum();
debug_assert_eq!(total, samples.len(), "the sizes describe the corpus");
if total < MIN_SAMPLES_SIZE {
return Err(TooSmall::Corpus);
}
if dict_size < DICT_SIZE_MIN {
return Err(TooSmall::Dictionary);
}
let selectivity = if selectivity == 0 {
DEFAULT_SELECTIVITY
} else {
selectivity
};
let nb_samples = sample_sizes.len();
let min_rep = if selectivity > 30 {
MIN_RATIO
} else {
u32::try_from(nb_samples >> selectivity).map_err(|_| TooSmall::Content)?
};
let list_size = DICT_LIST_SIZE_DEFAULT.max(nb_samples).max(dict_size / 16);
let mut list = vec![DictItem::default(); list_size];
list[0] = DictItem {
pos: 1,
length: 0,
savings: u32::MAX,
};
let mut used = total;
let mut kept = nb_samples;
while used > MAX_SAMPLES_SIZE {
kept -= 1;
used -= sample_sizes[kept];
}
let corpus = Corpus::new(&samples[..used]);
find_segments(&mut list, &corpus, used, min_rep);
let segments = list[0].pos as usize;
let content_size: usize = list[1..segments].iter().map(|d| d.length as usize).sum();
if content_size < CONTENT_SIZE_MIN {
return Err(TooSmall::Content);
}
let mut fit = 1;
let mut size = 0usize;
while fit < segments {
let next = size + list[fit].length as usize;
if next > dict_size {
break;
}
size = next;
fit += 1;
}
let mut content = vec![0u8; size];
let mut end = size;
for item in &list[1..fit] {
let length = item.length as usize;
let start = end - length;
let from = item.pos as usize;
corpus.copy_to(from, &mut content[start..end]);
end = start;
}
debug_assert_eq!(end, 0, "the kept segments fill the content exactly");
Ok(content)
}
fn find_segments(list: &mut [DictItem], corpus: &Corpus<'_>, len: usize, min_rep: u32) {
let min_ratio = min_rep.max(MIN_RATIO);
debug_assert_eq!(corpus.samples.len(), len);
let sa = suffix_array(corpus.samples);
let mut rank = vec![0u32; len];
for (at, &pos) in sa.iter().enumerate() {
rank[pos as usize] = at as u32;
}
let suffixes = Suffixes::new(sa, len);
let mut done = vec![false; len + 16];
let max_size = list.len() as u32;
let mut cursor = 0usize;
while cursor < len {
suffixes.prefetch_for(&rank, corpus.samples, cursor);
if done[cursor] {
cursor += 1;
continue;
}
let solution = analyze_position(
&mut done,
&suffixes,
i64::from(rank[cursor]),
corpus,
min_ratio,
);
if solution.length == 0 {
cursor += 1;
continue;
}
insert_item(list, max_size, solution, corpus);
cursor += solution.length as usize;
}
}
#[inline]
fn mark(done: &mut [bool], at: usize) {
if let Some(slot) = done.get_mut(at) {
*slot = true;
}
}
#[inline]
fn mark_run(done: &mut [bool], at: usize, len: usize) {
let end = (at + len).min(done.len());
if at < end {
done[at..end].fill(true);
}
}
fn analyze_position(
done: &mut [bool],
suffixes: &Suffixes,
mut start: i64,
corpus: &Corpus<'_>,
min_ratio: u32,
) -> DictItem {
let mut pos = suffixes.at(start);
let mut end = start;
let empty = DictItem::default();
mark(done, pos);
let repeats = |a: usize, b: usize| matches!((corpus.read16(a), corpus.read16(b)), (Some(x), Some(y)) if x == y);
if repeats(pos, pos + 2) || repeats(pos + 1, pos + 3) || repeats(pos + 2, pos + 4) {
let pattern = corpus.read16(pos + 4);
let mut pattern_end = 6usize;
while pattern.is_some() && corpus.read16(pos + pattern_end) == pattern {
pattern_end += 2;
}
if corpus.byte(pos + pattern_end).is_some()
&& corpus.byte(pos + pattern_end) == corpus.byte(pos + pattern_end - 1)
{
pattern_end += 1;
}
mark_run(done, pos + 1, pattern_end - 1);
return empty;
}
loop {
end += 1;
if end >= suffixes.ranks() || corpus.common(pos, suffixes.at(end)) < MIN_MATCH_LENGTH {
break;
}
}
while start > 0 && corpus.common(pos, suffixes.at(start - 1)) >= MIN_MATCH_LENGTH {
start -= 1;
}
if ((end - start) as u32) < min_ratio {
for id in start..end {
mark(done, suffixes.at(id));
}
return empty;
}
let mut refined_start = start;
let mut refined_end = end;
let mut mml = MIN_MATCH_LENGTH;
loop {
let mut current_char = 0u8;
let mut current_count = 0u32;
let mut current_id = refined_start;
let mut selected_count = 0u32;
let mut selected_id = current_id;
for id in refined_start..refined_end {
let c = corpus.byte(suffixes.at(id) + mml).unwrap_or(0);
if c != current_char {
if current_count > selected_count {
selected_count = current_count;
selected_id = current_id;
}
current_id = id;
current_char = c;
current_count = 0;
}
current_count += 1;
}
if current_count > selected_count {
selected_count = current_count;
selected_id = current_id;
}
if selected_count < min_ratio {
break;
}
refined_start = selected_id;
refined_end = refined_start + i64::from(selected_count);
mml += 1;
}
start = refined_start;
pos = suffixes.at(refined_start);
end = start;
let mut lengths = [0u32; LENGTH_LIMIT];
loop {
end += 1;
let length = corpus.common(pos, suffixes.at(end)).min(LENGTH_LIMIT - 1);
lengths[length] += 1;
if length < MIN_MATCH_LENGTH || end >= suffixes.ranks() {
break;
}
}
{
let mut length = MIN_MATCH_LENGTH;
while length >= MIN_MATCH_LENGTH && start > 0 {
length = corpus
.common(pos, suffixes.at(start - 1))
.min(LENGTH_LIMIT - 1);
lengths[length] += 1;
if length >= MIN_MATCH_LENGTH {
start -= 1;
}
}
}
let mut cumulative = [0u32; LENGTH_LIMIT];
cumulative[LENGTH_LIMIT - 1] = lengths[LENGTH_LIMIT - 1];
for i in (0..LENGTH_LIMIT - 1).rev() {
cumulative[i] = cumulative[i + 1] + lengths[i];
}
let mut max_length = LENGTH_LIMIT - 1;
while max_length >= MIN_MATCH_LENGTH && cumulative[max_length] < min_ratio {
max_length -= 1;
}
{
let last = corpus.byte(pos + max_length - 1);
let mut l = max_length;
while l >= 2 && corpus.byte(pos + l - 2) == last {
l -= 1;
}
max_length = l;
}
if max_length < MIN_MATCH_LENGTH {
return empty;
}
let mut savings = [0u32; LENGTH_LIMIT];
for i in MIN_MATCH_LENGTH..=max_length {
savings[i] = savings[i - 1].wrapping_add(lengths[i].wrapping_mul((i - 3) as u32));
}
let solution = DictItem {
pos: pos as u32,
length: max_length as u32,
savings: savings[max_length],
};
for id in start..end {
let tested = suffixes.at(id);
let length = if tested == pos {
max_length
} else {
corpus.common(pos, tested).min(max_length)
};
mark_run(done, tested, length);
}
solution
}
fn is_included(corpus: &Corpus<'_>, a: usize, b: usize, length: usize) -> bool {
let samples = corpus.samples;
if let (Some(x), Some(y)) = (samples.get(a..a + length), samples.get(b..b + length)) {
return x == y;
}
a.max(b) + length <= corpus.len()
&& (0..length).all(|i| corpus.byte(a + i) == corpus.byte(b + i))
}
fn promote(list: &mut [DictItem], mut at: usize) -> usize {
let item = list[at];
while at > 1 && list[at - 1].savings < item.savings {
list[at] = list[at - 1];
at -= 1;
}
list[at] = item;
at
}
fn try_merge(list: &mut [DictItem], elt: DictItem, skip: usize, corpus: &Corpus<'_>) -> usize {
let size = list[0].pos as usize;
let elt_end = elt.pos + elt.length;
let mut hit = 0;
for (u, item) in list[..size].iter().enumerate().skip(1) {
if item.pos > elt.pos && item.pos <= elt_end && u != skip {
hit = u;
break;
}
}
if hit != 0 {
let item = &mut list[hit];
let added = item.pos - elt.pos;
item.length += added;
item.pos = elt.pos;
item.savings = item
.savings
.wrapping_add(elt.savings.wrapping_mul(added) / elt.length);
item.savings = item.savings.wrapping_add(elt.length / 8);
return promote(list, hit);
}
let elt_head = corpus.read64(elt.pos as usize + 1);
for u in 1..size {
let item = list[u];
if u == skip {
continue;
}
if item.pos + item.length >= elt.pos && item.pos < elt.pos {
let added = elt_end as i64 - i64::from(item.pos + item.length);
let item = &mut list[u];
item.savings = item.savings.wrapping_add(elt.length / 8);
if added > 0 {
item.length += added as u32;
item.savings = item
.savings
.wrapping_add(elt.savings.wrapping_mul(added as u32) / elt.length);
}
return promote(list, u);
}
if elt_head.is_some()
&& corpus.read64(item.pos as usize) == elt_head
&& is_included(
corpus,
item.pos as usize,
elt.pos as usize + 1,
item.length as usize,
)
{
let added = (i64::from(elt.length) - i64::from(item.length)).max(1) as u64;
let item = &mut list[u];
item.pos = elt.pos;
item.savings = item
.savings
.wrapping_add((u64::from(elt.savings) * added / u64::from(elt.length)) as u32);
item.length = elt.length.min(item.length + 1);
return u;
}
}
0
}
fn remove_item(list: &mut [DictItem], id: usize) {
if id == 0 {
return;
}
let max = list[0].pos as usize;
for u in id..max - 1 {
list[u] = list[u + 1];
}
list[0].pos -= 1;
}
fn insert_item(list: &mut [DictItem], max_size: u32, elt: DictItem, corpus: &Corpus<'_>) {
let mut merge_id = try_merge(list, elt, 0, corpus);
if merge_id != 0 {
loop {
let merged = try_merge(list, list[merge_id], merge_id, corpus);
if merged == 0 {
return;
}
remove_item(list, merge_id);
merge_id = merged;
}
}
let next = list[0].pos.min(max_size - 1) as usize;
let mut current = next - 1;
while list[current].savings < elt.savings {
list[current + 1] = list[current];
current -= 1;
}
list[current + 1] = elt;
list[0].pos = next as u32 + 1;
}
#[cfg(test)]
mod tests;