1use alloc::vec;
13use alloc::vec::Vec;
14
15use super::suffix_array::suffix_array;
16use crate::encoding::fastpath::{EncodeTier, select_kernel};
17
18const MIN_RATIO: u32 = 4;
20const MIN_MATCH_LENGTH: usize = 7;
22const LENGTH_LIMIT: usize = 64;
24const DICT_LIST_SIZE_DEFAULT: usize = 10_000;
26const NOISE_LENGTH: usize = 32;
29const MAX_SAMPLES_SIZE: usize = 2000 << 20;
32const CONTENT_SIZE_MIN: usize = 128;
34pub(crate) const MIN_SAMPLES_SIZE: usize = CONTENT_SIZE_MIN * MIN_RATIO as usize;
36pub const DEFAULT_SELECTIVITY: u32 = 9;
38
39pub(crate) const DICT_SIZE_MIN: usize = super::MIN_TRAINED_DICT_SIZE;
41
42#[derive(Debug, Clone, Copy, PartialEq, Eq)]
44pub(crate) enum TooSmall {
45 Dictionary,
47 Corpus,
49 Content,
51}
52
53#[derive(Debug, Clone, Copy, Default)]
57struct DictItem {
58 pos: u32,
59 length: u32,
60 savings: u32,
61}
62
63fn noise_band() -> [u8; NOISE_LENGTH] {
65 let mut band = [0u8; NOISE_LENGTH];
66 let mut acc: u32 = 2_654_435_761;
67 for byte in &mut band {
68 acc = acc.wrapping_mul(2_246_822_519);
69 *byte = (acc >> 21) as u8;
70 }
71 band
72}
73
74struct Corpus<'a, K: EncodeTier> {
84 samples: &'a [u8],
85 noise: [u8; NOISE_LENGTH],
86 tier: core::marker::PhantomData<K>,
87}
88
89impl<'a, K: EncodeTier> Corpus<'a, K> {
90 fn new(samples: &'a [u8]) -> Self {
91 Self {
92 samples,
93 noise: noise_band(),
94 tier: core::marker::PhantomData,
95 }
96 }
97
98 #[inline]
100 fn len(&self) -> usize {
101 self.samples.len() + NOISE_LENGTH
102 }
103
104 #[inline]
105 fn byte(&self, at: usize) -> Option<u8> {
106 match self.samples.get(at) {
107 Some(&byte) => Some(byte),
108 None => self.noise.get(at - self.samples.len()).copied(),
109 }
110 }
111
112 #[inline(always)]
114 fn read<const N: usize>(&self, at: usize) -> Option<[u8; N]> {
115 if let Some(bytes) = self.samples.get(at..at + N) {
116 return Some(bytes.try_into().expect("N bytes"));
117 }
118 self.read_across(at)
119 }
120
121 #[cold]
124 #[inline(never)]
125 fn read_across<const N: usize>(&self, at: usize) -> Option<[u8; N]> {
126 let mut out = [0u8; N];
127 for (i, slot) in out.iter_mut().enumerate() {
128 *slot = self.byte(at + i)?;
129 }
130 Some(out)
131 }
132
133 #[inline]
134 fn read16(&self, at: usize) -> Option<u16> {
135 self.read::<2>(at).map(u16::from_le_bytes)
136 }
137
138 #[inline]
139 fn read64(&self, at: usize) -> Option<u64> {
140 self.read::<8>(at).map(u64::from_le_bytes)
141 }
142
143 #[inline]
148 fn common(&self, a: usize, b: usize) -> usize {
149 let samples = self.samples;
150 debug_assert!(a.max(b) <= samples.len());
154 let in_samples = samples.len() - a.max(b);
155 let base = samples.as_ptr();
156 let n = unsafe { K::common_prefix_len_ptr(base.add(a), base.add(b), in_samples) };
159 if n < in_samples {
160 return n;
161 }
162 self.common_tail(a, b, n)
163 }
164
165 #[cold]
168 #[inline(never)]
169 fn common_tail(&self, a: usize, b: usize, mut n: usize) -> usize {
170 let Some(limit) = self.len().checked_sub(a.max(b)) else {
171 return 0;
172 };
173 while n < limit && self.byte(a + n) == self.byte(b + n) {
174 n += 1;
175 }
176 n
177 }
178
179 fn copy_to(&self, from: usize, out: &mut [u8]) {
181 if let Some(bytes) = self.samples.get(from..from + out.len()) {
182 out.copy_from_slice(bytes);
183 return;
184 }
185 for (i, slot) in out.iter_mut().enumerate() {
186 *slot = self
187 .byte(from + i)
188 .expect("a segment lies inside the corpus");
189 }
190 }
191}
192
193const FAR_AHEAD: usize = 16;
196
197const NEAR_AHEAD: usize = 8;
200
201struct Suffixes {
206 sa: Vec<u32>,
207 noise: usize,
208}
209
210impl Suffixes {
211 fn new(sa: Vec<u32>, len: usize) -> Self {
212 Self { sa, noise: len }
213 }
214
215 #[inline]
217 fn ranks(&self) -> i64 {
218 self.sa.len() as i64
219 }
220
221 #[inline]
225 fn at(&self, at: i64) -> usize {
226 match self.sa.get(at as u64 as usize) {
227 Some(&pos) => pos as usize,
228 None => self.noise,
229 }
230 }
231
232 #[inline(always)]
239 fn prefetch_for(&self, rank: &[u32], samples: &[u8], cursor: usize) {
240 use crate::decoding::prefetch::prefetch_l1_at;
241 let len = self.sa.len();
242 if let Some(&far) = rank.get(cursor + FAR_AHEAD) {
243 prefetch_l1_at(self.sa.as_ptr().wrapping_add(far as usize).cast());
244 }
245 if let Some(&near) = rank.get(cursor + NEAR_AHEAD) {
246 let near = near as usize;
247 let below = near.wrapping_sub(1).min(len - 1);
249 let above = (near + 1).min(len - 1);
250 let (below, above) = unsafe {
252 (
253 *self.sa.get_unchecked(below) as usize,
254 *self.sa.get_unchecked(above) as usize,
255 )
256 };
257 prefetch_l1_at(samples.as_ptr().wrapping_add(below));
258 prefetch_l1_at(samples.as_ptr().wrapping_add(above));
259 }
260 }
261}
262
263pub(crate) fn train_legacy_raw(
270 samples: &[u8],
271 sample_sizes: &[usize],
272 dict_size: usize,
273 selectivity: u32,
274) -> Result<Vec<u8>, TooSmall> {
275 train_legacy_raw_under(
276 select_kernel(),
277 samples,
278 sample_sizes,
279 dict_size,
280 selectivity,
281 )
282}
283
284fn train_legacy_raw_under(
287 tier: crate::encoding::fastpath::FastpathKernel,
288 samples: &[u8],
289 sample_sizes: &[usize],
290 dict_size: usize,
291 selectivity: u32,
292) -> Result<Vec<u8>, TooSmall> {
293 let total: usize = sample_sizes.iter().sum();
294 debug_assert_eq!(total, samples.len(), "the sizes describe the corpus");
295 if total < MIN_SAMPLES_SIZE {
298 return Err(TooSmall::Corpus);
299 }
300 if dict_size < DICT_SIZE_MIN {
301 return Err(TooSmall::Dictionary);
302 }
303 let selectivity = if selectivity == 0 {
304 DEFAULT_SELECTIVITY
305 } else {
306 selectivity
307 };
308 let nb_samples = sample_sizes.len();
309 let min_rep = if selectivity > 30 {
310 MIN_RATIO
311 } else {
312 u32::try_from(nb_samples >> selectivity).map_err(|_| TooSmall::Content)?
315 };
316 let list_size = DICT_LIST_SIZE_DEFAULT.max(nb_samples).max(dict_size / 16);
317 let mut list = vec![DictItem::default(); list_size];
318 list[0] = DictItem {
319 pos: 1,
320 length: 0,
321 savings: u32::MAX,
322 };
323
324 let mut used = total;
326 let mut kept = nb_samples;
327 while used > MAX_SAMPLES_SIZE {
328 kept -= 1;
329 used -= sample_sizes[kept];
330 }
331 crate::encoding::fastpath::with_encode_tier!(
333 tier,
334 K => select_content::<K>(&mut list, &samples[..used], min_rep, dict_size)
335 )
336}
337
338fn select_content<K: EncodeTier>(
341 list: &mut [DictItem],
342 samples: &[u8],
343 min_rep: u32,
344 dict_size: usize,
345) -> Result<Vec<u8>, TooSmall> {
346 let corpus = Corpus::<K>::new(samples);
347 find_segments(list, &corpus, samples.len(), min_rep);
348
349 let segments = list[0].pos as usize;
350 let content_size: usize = list[1..segments].iter().map(|d| d.length as usize).sum();
351 if content_size < CONTENT_SIZE_MIN {
352 return Err(TooSmall::Content);
353 }
354 let mut fit = 1;
356 let mut size = 0usize;
357 while fit < segments {
358 let next = size + list[fit].length as usize;
359 if next > dict_size {
360 break;
361 }
362 size = next;
363 fit += 1;
364 }
365 let mut content = vec![0u8; size];
367 let mut end = size;
368 for item in &list[1..fit] {
369 let length = item.length as usize;
370 let start = end - length;
371 let from = item.pos as usize;
372 corpus.copy_to(from, &mut content[start..end]);
373 end = start;
374 }
375 debug_assert_eq!(end, 0, "the kept segments fill the content exactly");
376 Ok(content)
377}
378
379fn find_segments<K: EncodeTier>(
382 list: &mut [DictItem],
383 corpus: &Corpus<'_, K>,
384 len: usize,
385 min_rep: u32,
386) {
387 let min_ratio = min_rep.max(MIN_RATIO);
388 debug_assert_eq!(corpus.samples.len(), len);
389 let sa = suffix_array(corpus.samples);
390 let mut rank = vec![0u32; len];
391 for (at, &pos) in sa.iter().enumerate() {
392 rank[pos as usize] = at as u32;
393 }
394 let suffixes = Suffixes::new(sa, len);
395 let mut done = vec![false; len + 16];
398 let max_size = list.len() as u32;
399
400 let mut cursor = 0usize;
401 while cursor < len {
402 suffixes.prefetch_for(&rank, corpus.samples, cursor);
403 if done[cursor] {
404 cursor += 1;
405 continue;
406 }
407 let solution = analyze_position(
408 &mut done,
409 &suffixes,
410 i64::from(rank[cursor]),
411 corpus,
412 min_ratio,
413 );
414 if solution.length == 0 {
415 cursor += 1;
416 continue;
417 }
418 insert_item(list, max_size, solution, corpus);
419 cursor += solution.length as usize;
420 }
421}
422
423#[inline]
425fn mark(done: &mut [bool], at: usize) {
426 if let Some(slot) = done.get_mut(at) {
427 *slot = true;
428 }
429}
430
431#[inline]
434fn mark_run(done: &mut [bool], at: usize, len: usize) {
435 let end = (at + len).min(done.len());
436 if at < end {
437 done[at..end].fill(true);
438 }
439}
440
441fn analyze_position<K: EncodeTier>(
444 done: &mut [bool],
445 suffixes: &Suffixes,
446 mut start: i64,
447 corpus: &Corpus<'_, K>,
448 min_ratio: u32,
449) -> DictItem {
450 let mut pos = suffixes.at(start);
451 let mut end = start;
452 let empty = DictItem::default();
453 mark(done, pos);
454
455 let repeats = |a: usize, b: usize| matches!((corpus.read16(a), corpus.read16(b)), (Some(x), Some(y)) if x == y);
458 if repeats(pos, pos + 2) || repeats(pos + 1, pos + 3) || repeats(pos + 2, pos + 4) {
459 let pattern = corpus.read16(pos + 4);
460 let mut pattern_end = 6usize;
461 while pattern.is_some() && corpus.read16(pos + pattern_end) == pattern {
462 pattern_end += 2;
463 }
464 if corpus.byte(pos + pattern_end).is_some()
465 && corpus.byte(pos + pattern_end) == corpus.byte(pos + pattern_end - 1)
466 {
467 pattern_end += 1;
468 }
469 mark_run(done, pos + 1, pattern_end - 1);
470 return empty;
471 }
472
473 loop {
480 end += 1;
481 if end >= suffixes.ranks() || corpus.common(pos, suffixes.at(end)) < MIN_MATCH_LENGTH {
482 break;
483 }
484 }
485 while start > 0 && corpus.common(pos, suffixes.at(start - 1)) >= MIN_MATCH_LENGTH {
486 start -= 1;
487 }
488
489 if ((end - start) as u32) < min_ratio {
490 for id in start..end {
491 mark(done, suffixes.at(id));
492 }
493 return empty;
494 }
495
496 let mut refined_start = start;
499 let mut refined_end = end;
500 let mut mml = MIN_MATCH_LENGTH;
501 loop {
502 let mut current_char = 0u8;
503 let mut current_count = 0u32;
504 let mut current_id = refined_start;
505 let mut selected_count = 0u32;
506 let mut selected_id = current_id;
507 for id in refined_start..refined_end {
508 let c = corpus.byte(suffixes.at(id) + mml).unwrap_or(0);
509 if c != current_char {
510 if current_count > selected_count {
511 selected_count = current_count;
512 selected_id = current_id;
513 }
514 current_id = id;
515 current_char = c;
516 current_count = 0;
517 }
518 current_count += 1;
519 }
520 if current_count > selected_count {
521 selected_count = current_count;
522 selected_id = current_id;
523 }
524 if selected_count < min_ratio {
525 break;
526 }
527 refined_start = selected_id;
528 refined_end = refined_start + i64::from(selected_count);
529 mml += 1;
530 }
531
532 start = refined_start;
534 pos = suffixes.at(refined_start);
535 end = start;
536 let mut lengths = [0u32; LENGTH_LIMIT];
537 loop {
538 end += 1;
539 let length = corpus.common(pos, suffixes.at(end)).min(LENGTH_LIMIT - 1);
540 lengths[length] += 1;
541 if length < MIN_MATCH_LENGTH || end >= suffixes.ranks() {
544 break;
545 }
546 }
547 {
548 let mut length = MIN_MATCH_LENGTH;
549 while length >= MIN_MATCH_LENGTH && start > 0 {
550 length = corpus
551 .common(pos, suffixes.at(start - 1))
552 .min(LENGTH_LIMIT - 1);
553 lengths[length] += 1;
554 if length >= MIN_MATCH_LENGTH {
555 start -= 1;
556 }
557 }
558 }
559
560 let mut cumulative = [0u32; LENGTH_LIMIT];
562 cumulative[LENGTH_LIMIT - 1] = lengths[LENGTH_LIMIT - 1];
563 for i in (0..LENGTH_LIMIT - 1).rev() {
564 cumulative[i] = cumulative[i + 1] + lengths[i];
565 }
566 let mut max_length = LENGTH_LIMIT - 1;
567 while max_length >= MIN_MATCH_LENGTH && cumulative[max_length] < min_ratio {
568 max_length -= 1;
569 }
570 {
572 let last = corpus.byte(pos + max_length - 1);
573 let mut l = max_length;
574 while l >= 2 && corpus.byte(pos + l - 2) == last {
575 l -= 1;
576 }
577 max_length = l;
578 }
579 if max_length < MIN_MATCH_LENGTH {
580 return empty;
581 }
582
583 let mut savings = [0u32; LENGTH_LIMIT];
584 for i in MIN_MATCH_LENGTH..=max_length {
585 savings[i] = savings[i - 1].wrapping_add(lengths[i].wrapping_mul((i - 3) as u32));
586 }
587 let solution = DictItem {
588 pos: pos as u32,
589 length: max_length as u32,
590 savings: savings[max_length],
591 };
592
593 for id in start..end {
594 let tested = suffixes.at(id);
595 let length = if tested == pos {
596 max_length
597 } else {
598 corpus.common(pos, tested).min(max_length)
599 };
600 mark_run(done, tested, length);
601 }
602 solution
603}
604
605fn is_included<K: EncodeTier>(corpus: &Corpus<'_, K>, a: usize, b: usize, length: usize) -> bool {
609 let samples = corpus.samples;
610 if let (Some(x), Some(y)) = (samples.get(a..a + length), samples.get(b..b + length)) {
611 return x == y;
612 }
613 a.max(b) + length <= corpus.len()
614 && (0..length).all(|i| corpus.byte(a + i) == corpus.byte(b + i))
615}
616
617fn promote(list: &mut [DictItem], mut at: usize) -> usize {
619 let item = list[at];
620 while at > 1 && list[at - 1].savings < item.savings {
621 list[at] = list[at - 1];
622 at -= 1;
623 }
624 list[at] = item;
625 at
626}
627
628fn try_merge<K: EncodeTier>(
631 list: &mut [DictItem],
632 elt: DictItem,
633 skip: usize,
634 corpus: &Corpus<'_, K>,
635) -> usize {
636 let size = list[0].pos as usize;
637 let elt_end = elt.pos + elt.length;
638
639 let mut hit = 0;
642 for (u, item) in list[..size].iter().enumerate().skip(1) {
643 if item.pos > elt.pos && item.pos <= elt_end && u != skip {
644 hit = u;
645 break;
646 }
647 }
648 if hit != 0 {
649 let item = &mut list[hit];
650 let added = item.pos - elt.pos;
651 item.length += added;
652 item.pos = elt.pos;
653 item.savings = item
654 .savings
655 .wrapping_add(elt.savings.wrapping_mul(added) / elt.length);
656 item.savings = item.savings.wrapping_add(elt.length / 8);
657 return promote(list, hit);
658 }
659
660 let elt_head = corpus.read64(elt.pos as usize + 1);
664 for u in 1..size {
665 let item = list[u];
666 if u == skip {
667 continue;
668 }
669 if item.pos + item.length >= elt.pos && item.pos < elt.pos {
670 let added = elt_end as i64 - i64::from(item.pos + item.length);
671 let item = &mut list[u];
672 item.savings = item.savings.wrapping_add(elt.length / 8);
673 if added > 0 {
674 item.length += added as u32;
675 item.savings = item
676 .savings
677 .wrapping_add(elt.savings.wrapping_mul(added as u32) / elt.length);
678 }
679 return promote(list, u);
680 }
681 if elt_head.is_some()
682 && corpus.read64(item.pos as usize) == elt_head
683 && is_included(
684 corpus,
685 item.pos as usize,
686 elt.pos as usize + 1,
687 item.length as usize,
688 )
689 {
690 let added = (i64::from(elt.length) - i64::from(item.length)).max(1) as u64;
693 let item = &mut list[u];
694 item.pos = elt.pos;
695 item.savings = item
696 .savings
697 .wrapping_add((u64::from(elt.savings) * added / u64::from(elt.length)) as u32);
698 item.length = elt.length.min(item.length + 1);
699 return u;
700 }
701 }
702 0
703}
704
705fn remove_item(list: &mut [DictItem], id: usize) {
707 if id == 0 {
708 return;
709 }
710 let max = list[0].pos as usize;
711 for u in id..max - 1 {
712 list[u] = list[u + 1];
713 }
714 list[0].pos -= 1;
715}
716
717fn insert_item<K: EncodeTier>(
721 list: &mut [DictItem],
722 max_size: u32,
723 elt: DictItem,
724 corpus: &Corpus<'_, K>,
725) {
726 let mut merge_id = try_merge(list, elt, 0, corpus);
727 if merge_id != 0 {
728 loop {
729 let merged = try_merge(list, list[merge_id], merge_id, corpus);
730 if merged == 0 {
731 return;
732 }
733 remove_item(list, merge_id);
734 merge_id = merged;
735 }
736 }
737 let next = list[0].pos.min(max_size - 1) as usize;
738 let mut current = next - 1;
739 while list[current].savings < elt.savings {
740 list[current + 1] = list[current];
741 current -= 1;
742 }
743 list[current + 1] = elt;
744 list[0].pos = next as u32 + 1;
745}
746
747#[cfg(test)]
748mod tests;