Skip to main content

lz4rip_encode/
compress.rs

1//! LZ4 block compression.
2
3use core::fmt;
4
5use crate::hashtable::HashTable;
6use crate::verified_sink::VerifiedSliceSink;
7#[cfg(feature = "alloc")]
8use alloc::vec;
9use lz4rip_core::CompressError;
10use lz4rip_core::END_OFFSET;
11use lz4rip_core::LZ4_MIN_LENGTH;
12use lz4rip_core::MAX_DISTANCE;
13use lz4rip_core::MFLIMIT;
14use lz4rip_core::MINMATCH;
15use lz4rip_core::Sink;
16use lz4rip_core::WINDOW_SIZE;
17
18#[cfg(feature = "alloc")]
19use alloc::vec::Vec;
20
21pub(crate) use crate::hashtable::HashTableU32;
22pub(crate) use crate::hashtable::HashTableU32U16;
23pub use crate::hashtable::{DEFAULT_DICT_ENTRIES, DEFAULT_NODICT_ENTRIES, MIN_ENTRIES};
24
25/// Inputs up to this size reuse the no-dict hash table across calls (epoch-based
26/// table reuse); larger inputs clear it. Independent of table entry count.
27const EPOCH_THRESHOLD: usize = 8 * 1024;
28
29/// Skip acceleration: step grows by 1 every `1 << N` consecutive non-matches.
30/// C lz4 uses 6; see DESIGN.md for tradeoff analysis.
31const INCREASE_STEPSIZE_BITSHIFT: usize = 3;
32
33/// Inputs up to this size use the dict table read-only (no per-call clearing or
34/// table writes). Self-references within small inputs are rare; the dict provides
35/// virtually all matches. Skips the 8 KB table.clear() and all put_at writes.
36const DICT_READONLY_LIMIT: usize = 256;
37
38#[inline]
39fn token_from_literal(lit_len: usize) -> u8 {
40    if lit_len < 0xF {
41        (lit_len as u8) << 4
42    } else {
43        0xF0
44    }
45}
46
47#[inline]
48fn token_from_literal_and_match_length(lit_len: usize, duplicate_length: usize) -> u8 {
49    let mut token = if lit_len < 0xF {
50        (lit_len as u8) << 4
51    } else {
52        0xF0
53    };
54
55    token |= if duplicate_length < 0xF {
56        duplicate_length as u8
57    } else {
58        0xF
59    };
60
61    token
62}
63
64/// Write a variable-length integer in the LZ4 encoding.
65#[inline]
66pub fn write_integer(output: &mut impl Sink, mut n: usize) {
67    while n >= 0xFF {
68        n -= 0xFF;
69        push_byte(output, 0xFF);
70    }
71    push_byte(output, n as u8);
72}
73
74#[cold]
75fn handle_last_literals(output: &mut impl Sink, input: &[u8], start: usize) {
76    let lit_len = input.len() - start;
77
78    let token = token_from_literal(lit_len);
79    push_byte(output, token);
80    if lit_len >= 0xF {
81        write_integer(output, lit_len - 0xF);
82    }
83    output.extend_from_slice(&input[start..]);
84}
85
86#[inline]
87fn backtrack_match(
88    input: &[u8],
89    cur: &mut usize,
90    literal_start: usize,
91    source: &[u8],
92    candidate: &mut usize,
93) {
94    while *candidate > 0 && *cur > literal_start && input[*cur - 1] == source[*candidate - 1] {
95        *cur -= 1;
96        *candidate -= 1;
97    }
98}
99
100/// Core block compression loop, monomorphized over hash table type and dict mode.
101#[inline(never)]
102pub(crate) fn compress_internal<
103    T: HashTable,
104    const USE_DICT: bool,
105    const HAS_OFFSET: bool,
106    const READONLY: bool,
107    S: Sink,
108>(
109    input: &[u8],
110    input_pos: usize,
111    output: &mut S,
112    table: &mut T,
113    ext_dict: &[u8],
114    input_stream_offset: usize,
115) -> Result<usize, CompressError> {
116    assert!(input_pos <= input.len());
117    if USE_DICT {
118        assert!(ext_dict.len() <= WINDOW_SIZE);
119        assert!(ext_dict.len() <= input_stream_offset);
120        assert!(
121            input_stream_offset
122                .checked_add(input.len())
123                .and_then(|i| i.checked_add(ext_dict.len()))
124                .is_some_and(|i| i <= isize::MAX as usize)
125        );
126    } else {
127        assert!(ext_dict.is_empty());
128    }
129    if !HAS_OFFSET {
130        debug_assert_eq!(input_stream_offset, 0);
131    }
132    let input_stream_offset = if HAS_OFFSET { input_stream_offset } else { 0 };
133    if output.capacity() - output.pos() < get_maximum_output_size(input.len() - input_pos) {
134        return Err(CompressError::OutputTooSmall);
135    }
136
137    let output_start_pos = output.pos();
138    if input.len() - input_pos < LZ4_MIN_LENGTH {
139        handle_last_literals(output, input, input_pos);
140        return Ok(output.pos() - output_start_pos);
141    }
142
143    let ext_dict_stream_offset = input_stream_offset - ext_dict.len();
144    let end_pos_check = input.len() - MFLIMIT;
145    let mut literal_start = input_pos;
146    let mut cur = input_pos;
147
148    if cur == 0 && input_stream_offset == 0 {
149        let hash = T::get_hash_at_inbounds(input, 0);
150        if !READONLY {
151            table.put_at(hash, 0);
152        }
153        cur = 1;
154    }
155
156    let mut forward_hash = T::get_hash_at_inbounds(input, cur);
157
158    loop {
159        let mut candidate;
160        let mut candidate_source;
161        let mut offset;
162        let mut non_match_count = 1 << INCREASE_STEPSIZE_BITSHIFT;
163
164        loop {
165            let step = non_match_count >> INCREASE_STEPSIZE_BITSHIFT;
166            non_match_count += 1;
167            let next_cur = cur + step;
168
169            if next_cur > end_pos_check + 1 {
170                handle_last_literals(output, input, literal_start);
171                return Ok(output.pos() - output_start_pos);
172            }
173
174            let hash = forward_hash;
175            candidate = table.get_at(hash);
176            forward_hash = T::get_hash_at_inbounds(input, next_cur);
177            if !READONLY {
178                table.put_at(hash, cur + input_stream_offset);
179            }
180
181            debug_assert!(READONLY || candidate <= input_stream_offset + cur);
182
183            if candidate >= input_stream_offset
184                && input_stream_offset + cur - candidate <= MAX_DISTANCE
185            {
186                offset = (input_stream_offset + cur - candidate) as u16;
187                candidate -= input_stream_offset;
188                candidate_source = input;
189            } else if USE_DICT
190                && candidate >= ext_dict_stream_offset
191                && input_stream_offset + cur - candidate <= MAX_DISTANCE
192            {
193                offset = (input_stream_offset + cur - candidate) as u16;
194                candidate -= ext_dict_stream_offset;
195                candidate_source = ext_dict;
196            } else {
197                cur = next_cur;
198                continue;
199            }
200            let cand_bytes: u32 = crate::hashtable::get_batch_inbounds(candidate_source, candidate);
201            let curr_bytes: u32 = crate::hashtable::get_batch_inbounds(input, cur);
202
203            if cand_bytes == curr_bytes {
204                break;
205            }
206            cur = next_cur;
207        }
208
209        loop {
210            backtrack_match(
211                input,
212                &mut cur,
213                literal_start,
214                candidate_source,
215                &mut candidate,
216            );
217
218            let lit_len = cur - literal_start;
219
220            cur += MINMATCH;
221            candidate += MINMATCH;
222            let duplicate_length = crate::hashtable::count_same_bytes_inbounds(
223                input,
224                &mut cur,
225                candidate_source,
226                candidate,
227                END_OFFSET,
228            );
229
230            let hash = T::get_hash_at_inbounds(input, cur - 2);
231            if !READONLY {
232                table.put_at(hash, cur - 2 + input_stream_offset);
233            }
234
235            let token = token_from_literal_and_match_length(lit_len, duplicate_length);
236            push_byte(output, token);
237            if lit_len >= 0xF {
238                write_integer(output, lit_len - 0xF);
239            }
240            if lit_len > 0 {
241                copy_literals_wild(output, input, literal_start, lit_len);
242            }
243            push_u16(output, offset);
244            if duplicate_length >= 0xF {
245                write_integer(output, duplicate_length - 0xF);
246            }
247            literal_start = cur;
248
249            if !USE_DICT && cur <= end_pos_check {
250                let hash = T::get_hash_at_inbounds(input, cur);
251                let rematch = table.get_at(hash);
252
253                if input_stream_offset + cur - rematch <= MAX_DISTANCE
254                    && rematch >= input_stream_offset
255                {
256                    let rc = rematch - input_stream_offset;
257                    if crate::hashtable::get_batch_inbounds(input, cur)
258                        == crate::hashtable::get_batch_inbounds(input, rc)
259                    {
260                        table.put_at(hash, cur + input_stream_offset);
261                        candidate = rc;
262                        candidate_source = input;
263                        offset = (input_stream_offset + cur - rematch) as u16;
264                        continue;
265                    }
266                }
267                forward_hash = hash;
268            } else if cur <= end_pos_check {
269                forward_hash = T::get_hash_at_inbounds(input, cur);
270            }
271            break;
272        }
273    }
274}
275
276/// Compress with a caller-owned `HashTableU32`.
277///
278/// This is cross-crate plumbing for the frame encoder. It keeps the internal
279/// `HashTable` trait private, so downstream safe code cannot corrupt match
280/// finder invariants that protect unchecked reads.
281pub fn compress_into_sink_with_table<
282    const USE_DICT: bool,
283    const HAS_OFFSET: bool,
284    const READONLY: bool,
285    S: Sink,
286>(
287    input: &[u8],
288    input_pos: usize,
289    output: &mut S,
290    table: &mut HashTableU32,
291    ext_dict: &[u8],
292    input_stream_offset: usize,
293) -> Result<usize, CompressError> {
294    compress_internal::<_, USE_DICT, HAS_OFFSET, READONLY, _>(
295        input,
296        input_pos,
297        output,
298        table,
299        ext_dict,
300        input_stream_offset,
301    )
302}
303
304/// Dual-table compression for `CompressorRef::with_dict`.
305#[inline(never)]
306fn compress_with_dict_table<T: HashTable, S: Sink>(
307    input: &[u8],
308    output: &mut S,
309    table: &mut T,
310    dict_table: &T,
311    ext_dict: &[u8],
312    input_stream_offset: usize,
313) -> Result<usize, CompressError> {
314    debug_assert_eq!(input_stream_offset, ext_dict.len());
315    assert!(ext_dict.len() <= WINDOW_SIZE);
316    assert!(ext_dict.len() <= input_stream_offset);
317    assert!(
318        input_stream_offset
319            .checked_add(input.len())
320            .and_then(|i| i.checked_add(ext_dict.len()))
321            .is_some_and(|i| i <= isize::MAX as usize)
322    );
323    if output.capacity() - output.pos() < get_maximum_output_size(input.len()) {
324        return Err(CompressError::OutputTooSmall);
325    }
326
327    let output_start_pos = output.pos();
328    if input.len() < LZ4_MIN_LENGTH {
329        handle_last_literals(output, input, 0);
330        return Ok(output.pos() - output_start_pos);
331    }
332
333    let end_pos_check = input.len() - MFLIMIT;
334    let mut literal_start = 0;
335
336    let hash = T::get_hash_at_inbounds(input, 0);
337    table.put_at(hash, input_stream_offset);
338    let mut cur = 1;
339
340    let mut forward_hash = T::get_hash_at_inbounds(input, cur);
341
342    loop {
343        let mut candidate;
344        let candidate_source;
345        let offset;
346        let mut non_match_count = 1 << INCREASE_STEPSIZE_BITSHIFT;
347
348        loop {
349            let step = non_match_count >> INCREASE_STEPSIZE_BITSHIFT;
350            non_match_count += 1;
351            let next_cur = cur + step;
352
353            if next_cur > end_pos_check + 1 {
354                handle_last_literals(output, input, literal_start);
355                return Ok(output.pos() - output_start_pos);
356            }
357
358            let hash = forward_hash;
359            forward_hash = T::get_hash_at_inbounds(input, next_cur);
360            let curr_bytes: u32 = crate::hashtable::get_batch_inbounds(input, cur);
361
362            let main_candidate = table.get_at(hash);
363            table.put_at(hash, cur + input_stream_offset);
364
365            // Probe dict table first: for small inputs most matches come from
366            // the dict, and even for large inputs the dict covers the first
367            // window of repeated structure.
368            let dict_candidate = dict_table.get_at(hash);
369            if dict_candidate < input_stream_offset
370                && input_stream_offset + cur - dict_candidate <= MAX_DISTANCE
371            {
372                let cand_bytes: u32 =
373                    crate::hashtable::get_batch_inbounds(ext_dict, dict_candidate);
374                if cand_bytes == curr_bytes {
375                    offset = (input_stream_offset + cur - dict_candidate) as u16;
376                    candidate = dict_candidate;
377                    candidate_source = ext_dict;
378                    break;
379                }
380            }
381
382            if main_candidate >= input_stream_offset
383                && input_stream_offset + cur - main_candidate <= MAX_DISTANCE
384            {
385                let cand_bytes: u32 = crate::hashtable::get_batch_inbounds(
386                    input,
387                    main_candidate - input_stream_offset,
388                );
389                if cand_bytes == curr_bytes {
390                    offset = (input_stream_offset + cur - main_candidate) as u16;
391                    candidate = main_candidate - input_stream_offset;
392                    candidate_source = input;
393                    break;
394                }
395            }
396
397            cur = next_cur;
398        }
399
400        backtrack_match(
401            input,
402            &mut cur,
403            literal_start,
404            candidate_source,
405            &mut candidate,
406        );
407
408        let lit_len = cur - literal_start;
409
410        cur += MINMATCH;
411        candidate += MINMATCH;
412        let duplicate_length = crate::hashtable::count_same_bytes_inbounds(
413            input,
414            &mut cur,
415            candidate_source,
416            candidate,
417            END_OFFSET,
418        );
419
420        let hash = T::get_hash_at_inbounds(input, cur - 2);
421        table.put_at(hash, cur - 2 + input_stream_offset);
422
423        let token = token_from_literal_and_match_length(lit_len, duplicate_length);
424        push_byte(output, token);
425        if lit_len >= 0xF {
426            write_integer(output, lit_len - 0xF);
427        }
428        if lit_len > 0 {
429            copy_literals_wild(output, input, literal_start, lit_len);
430        }
431        push_u16(output, offset);
432        if duplicate_length >= 0xF {
433            write_integer(output, duplicate_length - 0xF);
434        }
435        literal_start = cur;
436
437        if cur <= end_pos_check {
438            forward_hash = T::get_hash_at_inbounds(input, cur);
439        }
440    }
441}
442
443#[inline]
444fn push_byte(output: &mut impl Sink, el: u8) {
445    output.push(el);
446}
447
448#[inline]
449fn push_u16(output: &mut impl Sink, el: u16) {
450    output.extend_from_slice(&el.to_le_bytes());
451}
452
453#[inline(always)]
454fn copy_literals_wild(output: &mut impl Sink, input: &[u8], input_start: usize, len: usize) {
455    output.extend_from_slice_wild(&input[input_start..input_start + len], len)
456}
457
458/// Compress `input` into `output` with optional dictionary data.
459pub fn compress_into_sink_with_dict<const USE_DICT: bool>(
460    input: &[u8],
461    output: &mut impl Sink,
462    mut dict_data: &[u8],
463) -> Result<usize, CompressError> {
464    if USE_DICT && dict_data.len() < MINMATCH {
465        return compress_into_sink_with_dict::<false>(input, output, b"");
466    }
467    if dict_data.len() + input.len() < u16::MAX as usize {
468        let mut dict: HashTableU32U16 = HashTableU32U16::new();
469        init_dict(&mut dict, &mut dict_data);
470        compress_internal::<_, USE_DICT, USE_DICT, false, _>(
471            input,
472            0,
473            output,
474            &mut dict,
475            dict_data,
476            dict_data.len(),
477        )
478    } else {
479        let mut dict: HashTableU32 = HashTableU32::new();
480        init_dict(&mut dict, &mut dict_data);
481        compress_internal::<_, USE_DICT, USE_DICT, false, _>(
482            input,
483            0,
484            output,
485            &mut dict,
486            dict_data,
487            dict_data.len(),
488        )
489    }
490}
491
492#[inline]
493pub(crate) fn init_dict<T: HashTable>(dict: &mut T, dict_data: &mut &[u8]) {
494    if dict_data.len() > WINDOW_SIZE {
495        *dict_data = &dict_data[dict_data.len() - WINDOW_SIZE..];
496    }
497    let mut i = 0usize;
498    while i + core::mem::size_of::<usize>() <= dict_data.len() {
499        let hash = T::get_hash_at(dict_data, i);
500        dict.put_at(hash, i);
501        i += 3;
502    }
503}
504
505/// Returns the maximum output size of the compressed data.
506/// Can be used to preallocate capacity on the output vector.
507///
508/// Returns `usize::MAX` if the result would overflow (e.g. on 32-bit with >3.9 GB input).
509#[must_use]
510#[inline]
511pub const fn get_maximum_output_size(input_len: usize) -> usize {
512    let raw = 16u64 + 4 + (input_len as u64 * 110 / 100);
513    if raw > usize::MAX as u64 {
514        usize::MAX
515    } else {
516        raw as usize
517    }
518}
519
520/// Compress all bytes of `input` into `output`.
521/// output should be preallocated with a size of
522/// `get_maximum_output_size`.
523///
524/// Returns the number of bytes written (compressed) into `output`.
525#[inline]
526pub fn compress_into(input: &[u8], output: &mut [u8]) -> Result<usize, CompressError> {
527    compress_into_sink_with_dict::<false>(input, &mut VerifiedSliceSink::new(output, 0), b"")
528}
529
530/// Compress all bytes of `input` into `output` using an external dictionary.
531///
532/// Returns the number of bytes written (compressed) into `output`.
533#[inline]
534pub fn compress_into_with_dict(
535    input: &[u8],
536    output: &mut [u8],
537    dict: &[u8],
538) -> Result<usize, CompressError> {
539    compress_into_sink_with_dict::<true>(input, &mut VerifiedSliceSink::new(output, 0), dict)
540}
541
542/// Compress all bytes of `input`.
543#[must_use]
544#[cfg(feature = "alloc")]
545#[inline]
546pub fn compress(input: &[u8]) -> Vec<u8> {
547    let max_compressed_size = get_maximum_output_size(input.len());
548    let mut compressed: Vec<u8> = vec![0u8; max_compressed_size];
549    let compressed_len = compress_into_sink_with_dict::<false>(
550        input,
551        &mut VerifiedSliceSink::new(&mut compressed, 0),
552        b"",
553    )
554    .unwrap();
555    compressed.truncate(compressed_len);
556
557    compressed
558}
559
560/// A reusable no-dict block compressor with `N` hash-table entries.
561///
562/// [`CompressorRef`] is the standard-sized alias (8 KB table). Use this generic
563/// form to pick a smaller table for memory-constrained (e.g. embedded) targets,
564/// e.g. `CompressorRefN::<512>::new()` for a 2 KB table. `N` must be a power of
565/// two (checked at compile time).
566///
567/// This is the no-alloc API. With `alloc`, use [`Compressor`](crate::Compressor)
568/// instead. For one-shot compression, use [`compress_into`] instead.
569///
570/// # Example
571/// ```
572/// use lz4rip_encode::{CompressorRef, get_maximum_output_size};
573///
574/// let mut comp = CompressorRef::new();
575/// let input = b"hello world, hello world, hello!";
576/// let mut output = vec![0u8; get_maximum_output_size(input.len())];
577/// let compressed_len = comp.compress_into(input, &mut output).unwrap();
578/// ```
579pub struct CompressorRefN<const N: usize = DEFAULT_NODICT_ENTRIES> {
580    table: HashTableU32<N>,
581    stream_offset: usize,
582}
583
584/// A reusable no-dict block compressor with the standard 8 KB hash table.
585pub type CompressorRef = CompressorRefN<DEFAULT_NODICT_ENTRIES>;
586
587impl<const N: usize> fmt::Debug for CompressorRefN<N> {
588    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
589        f.debug_struct("CompressorRef").finish_non_exhaustive()
590    }
591}
592
593impl<const N: usize> CompressorRefN<N> {
594    /// Create a new compressor without a dictionary.
595    #[must_use]
596    pub fn new() -> Self {
597        CompressorRefN {
598            table: HashTableU32::<N>::new(),
599            stream_offset: 0,
600        }
601    }
602
603    /// Compress `input` into `output`, returning the number of compressed bytes.
604    ///
605    /// `output` must be at least [`get_maximum_output_size`]`(input.len())` bytes.
606    pub fn compress_into(
607        &mut self,
608        input: &[u8],
609        output: &mut [u8],
610    ) -> Result<usize, CompressError> {
611        compress_plain_table(&mut self.table, &mut self.stream_offset, input, output)
612    }
613
614    /// Compress `input` into a new `Vec<u8>`.
615    #[cfg(feature = "alloc")]
616    pub fn compress(&mut self, input: &[u8]) -> Vec<u8> {
617        let max_compressed = get_maximum_output_size(input.len());
618        let mut compressed = vec![0u8; max_compressed];
619        let compressed_len = self.compress_into(input, &mut compressed).unwrap();
620        compressed.truncate(compressed_len);
621
622        compressed
623    }
624}
625
626/// A reusable dict block compressor (borrowing) with `N` entries per table.
627///
628/// [`DictCompressorRef`] is the standard-sized alias (two 8 KB tables). Use this
629/// generic form to pick smaller tables for memory-constrained targets, e.g.
630/// `DictCompressorRefN::<1024>::new(dict)` for two 2 KB tables. `N` must be a
631/// power of two (checked at compile time).
632///
633/// This is the no-alloc dict API. With `alloc`, use
634/// [`DictCompressor`](crate::DictCompressor) instead. Without a dictionary, use
635/// [`CompressorRef`].
636///
637/// # Example
638/// ```
639/// use lz4rip_encode::{DictCompressorRef, get_maximum_output_size};
640///
641/// let dict = b"the quick brown fox";
642/// let mut comp = DictCompressorRef::new(dict);
643/// let input = b"the quick brown fox jumps";
644/// let mut output = vec![0u8; get_maximum_output_size(input.len())];
645/// let compressed_len = comp.compress_into(input, &mut output).unwrap();
646/// ```
647pub struct DictCompressorRefN<'a, const N: usize = DEFAULT_DICT_ENTRIES> {
648    table: HashTableU32U16<N>,
649    pristine: HashTableU32U16<N>,
650    dict: &'a [u8],
651}
652
653/// A reusable dict block compressor (borrowing) with the standard 8 KB tables.
654pub type DictCompressorRef<'a> = DictCompressorRefN<'a, DEFAULT_DICT_ENTRIES>;
655
656impl<const N: usize> fmt::Debug for DictCompressorRefN<'_, N> {
657    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
658        f.debug_struct("DictCompressorRef")
659            .field("dict_len", &self.dict.len())
660            .finish()
661    }
662}
663
664impl<'a, const N: usize> DictCompressorRefN<'a, N> {
665    /// Create a new compressor seeded with an external dictionary.
666    ///
667    /// If `dict` is longer than the LZ4 window it is trimmed to the last
668    /// [`WINDOW_SIZE`](lz4rip_core::WINDOW_SIZE) bytes. A dictionary shorter than
669    /// 4 bytes is ignored (no dict matches); use [`CompressorRef`] for that case.
670    #[must_use]
671    pub fn new(dict: &'a [u8]) -> Self {
672        let trimmed = if dict.len() < MINMATCH {
673            b"".as_slice()
674        } else if dict.len() > WINDOW_SIZE {
675            &dict[dict.len() - WINDOW_SIZE..]
676        } else {
677            dict
678        };
679        let mut pristine = HashTableU32U16::<N>::new();
680        let mut dict_ref = trimmed;
681        init_dict(&mut pristine, &mut dict_ref);
682        DictCompressorRefN {
683            table: HashTableU32U16::<N>::new(),
684            pristine,
685            dict: trimmed,
686        }
687    }
688
689    /// Compress `input` into `output`, returning the number of compressed bytes.
690    ///
691    /// `output` must be at least [`get_maximum_output_size`]`(input.len())` bytes.
692    pub fn compress_into(
693        &mut self,
694        input: &[u8],
695        output: &mut [u8],
696    ) -> Result<usize, CompressError> {
697        compress_dict_tables(
698            &mut self.table,
699            &mut self.pristine,
700            self.dict,
701            input,
702            output,
703        )
704    }
705
706    /// Compress `input` into a new `Vec<u8>`.
707    #[cfg(feature = "alloc")]
708    pub fn compress(&mut self, input: &[u8]) -> Vec<u8> {
709        let max_compressed = get_maximum_output_size(input.len());
710        let mut compressed = vec![0u8; max_compressed];
711        let compressed_len = self.compress_into(input, &mut compressed).unwrap();
712        compressed.truncate(compressed_len);
713
714        compressed
715    }
716}
717
718/// Compress `input` using the dict main + pristine tables. Shared by
719/// [`CompressorRef::compress_into`] and the owning [`Compressor`] so the dict
720/// branch logic lives in one place regardless of how the tables are stored.
721pub(crate) fn compress_dict_tables<const N: usize>(
722    table: &mut HashTableU32U16<N>,
723    pristine: &mut HashTableU32U16<N>,
724    dict: &[u8],
725    input: &[u8],
726    output: &mut [u8],
727) -> Result<usize, CompressError> {
728    if input.len() <= DICT_READONLY_LIMIT && dict.len() + input.len() < u16::MAX as usize {
729        compress_internal::<_, true, true, true, _>(
730            input,
731            0,
732            &mut VerifiedSliceSink::new(output, 0),
733            pristine,
734            dict,
735            dict.len(),
736        )
737    } else if dict.len() + input.len() < u16::MAX as usize {
738        table.clear();
739        compress_with_dict_table(
740            input,
741            &mut VerifiedSliceSink::new(output, 0),
742            table,
743            pristine,
744            dict,
745            dict.len(),
746        )
747    } else {
748        // dict + input >= 64 KB: positions overflow u16, so use a u32 table sized
749        // to this compressor's `N` (honors the const-generic knob instead of
750        // allocating a standard 8 KB table).
751        let mut u32_table = HashTableU32::<N>::new();
752        let mut dict_data = dict;
753        init_dict(&mut u32_table, &mut dict_data);
754        compress_internal::<_, true, true, false, _>(
755            input,
756            0,
757            &mut VerifiedSliceSink::new(output, 0),
758            &mut u32_table,
759            dict_data,
760            dict_data.len(),
761        )
762    }
763}
764
765/// Compress `input` using the plain (no-dict) table with epoch-based reuse.
766/// Shared by [`CompressorRef::compress_into`] and the owning [`Compressor`].
767pub(crate) fn compress_plain_table<const N: usize>(
768    table: &mut HashTableU32<N>,
769    stream_offset: &mut usize,
770    input: &[u8],
771    output: &mut [u8],
772) -> Result<usize, CompressError> {
773    let offset = prepare_plain_table(table, stream_offset, input.len());
774    if offset > 0 {
775        compress_internal::<_, false, true, false, _>(
776            input,
777            0,
778            &mut VerifiedSliceSink::new(output, 0),
779            table,
780            b"",
781            offset,
782        )
783    } else {
784        compress_internal::<_, false, false, false, _>(
785            input,
786            0,
787            &mut VerifiedSliceSink::new(output, 0),
788            table,
789            b"",
790            0,
791        )
792    }
793}
794
795#[inline]
796fn prepare_plain_table<const N: usize>(
797    table: &mut HashTableU32<N>,
798    stream_offset: &mut usize,
799    input_len: usize,
800) -> usize {
801    if input_len > EPOCH_THRESHOLD {
802        table.clear();
803        *stream_offset = input_len + MAX_DISTANCE + 1;
804        return 0;
805    }
806    let offset = *stream_offset;
807    let next = offset
808        .checked_add(input_len)
809        .and_then(|v| v.checked_add(MAX_DISTANCE + 1));
810    if let Some(next) = next.filter(|&n| n <= u32::MAX as usize) {
811        *stream_offset = next;
812    } else {
813        table.clear();
814        *stream_offset = input_len + MAX_DISTANCE + 1;
815    }
816    offset
817}
818
819impl<const N: usize> Default for CompressorRefN<N> {
820    fn default() -> Self {
821        Self::new()
822    }
823}
824
825#[cfg(test)]
826mod tests {
827    use super::*;
828
829    fn count_same_bytes(input: &[u8], cur: &mut usize, source: &[u8], candidate: usize) -> usize {
830        const USIZE_SIZE: usize = core::mem::size_of::<usize>();
831        let cur_slice = &input[*cur..input.len() - END_OFFSET];
832        let cand_slice = &source[candidate..];
833
834        let mut num = 0;
835        for (block1, block2) in cur_slice
836            .chunks_exact(USIZE_SIZE)
837            .zip(cand_slice.chunks_exact(USIZE_SIZE))
838        {
839            let input_block = usize::from_ne_bytes(block1.try_into().unwrap());
840            let match_block = usize::from_ne_bytes(block2.try_into().unwrap());
841
842            if input_block == match_block {
843                num += USIZE_SIZE;
844            } else {
845                let diff = input_block ^ match_block;
846                num += (diff.to_le().trailing_zeros() / 8) as usize;
847                *cur += num;
848                return num;
849            }
850        }
851
852        #[cold]
853        fn count_same_bytes_tail(a: &[u8], b: &[u8], offset: usize) -> usize {
854            a.iter()
855                .zip(b)
856                .skip(offset)
857                .take_while(|(a, b)| a == b)
858                .count()
859        }
860        num += count_same_bytes_tail(cur_slice, cand_slice, num);
861
862        *cur += num;
863        num
864    }
865
866    #[test]
867    fn test_count_same_bytes() {
868        let first: &[u8] = &[
869            1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
870        ];
871        let second: &[u8] = &[
872            1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
873        ];
874        assert_eq!(count_same_bytes(first, &mut 0, second, 0), 16);
875
876        for diff_idx in 8..100 {
877            let first: Vec<u8> = (0u8..255).cycle().take(100 + 12).collect();
878            let mut second = first.clone();
879            second[diff_idx] = 255;
880            for start in 0..=diff_idx {
881                let same_bytes = count_same_bytes(&first, &mut start.clone(), &second, start);
882                assert_eq!(same_bytes, diff_idx - start);
883            }
884        }
885    }
886
887    #[test]
888    fn test_bug() {
889        let input: &[u8] = &[
890            10, 12, 14, 16, 18, 10, 12, 14, 16, 18, 10, 12, 14, 16, 18, 10, 12, 14, 16, 18,
891        ];
892        let mut output = [0u8; get_maximum_output_size(20)];
893        let _ = compress_into(input, &mut output).unwrap();
894    }
895
896    #[test]
897    fn test_conformant_last_block() {
898        let aaas: &[u8] = b"aaaaaaaaaaaaaaa";
899
900        let mut out = [0u8; get_maximum_output_size(15)];
901        let n = compress_into(&aaas[..12], &mut out).unwrap();
902        assert!(n > 12);
903        let n = compress_into(&aaas[..13], &mut out).unwrap();
904        assert!(n <= 13);
905        let n = compress_into(&aaas[..14], &mut out).unwrap();
906        assert!(n <= 14);
907        let n = compress_into(&aaas[..15], &mut out).unwrap();
908        assert!(n <= 15);
909    }
910}