Skip to main content

lz4_flex/block/
compress.rs

1//! The compression algorithm.
2//!
3//! We make use of hash tables to find duplicates. This gives a reasonable compression ratio with a
4//! high performance. It has fixed memory usage, which contrary to other approaches, makes it less
5//! memory hungry.
6
7use crate::block::hashtable::HashTable;
8use crate::block::END_OFFSET;
9use crate::block::LZ4_MIN_LENGTH;
10use crate::block::MAX_DISTANCE;
11use crate::block::MFLIMIT;
12use crate::block::MINMATCH;
13#[cfg(not(feature = "safe-encode"))]
14use crate::sink::PtrSink;
15use crate::sink::Sink;
16use crate::sink::SliceSink;
17#[allow(unused_imports)]
18use alloc::vec;
19
20#[allow(unused_imports)]
21use alloc::vec::Vec;
22
23pub(crate) use super::hashtable::HashTable4K;
24pub(crate) use super::hashtable::HashTable4KU16;
25use super::{CompressError, WINDOW_SIZE};
26
27/// Increase step size after 1<<INCREASE_STEPSIZE_BITSHIFT non matches
28const INCREASE_STEPSIZE_BITSHIFT: usize = 5;
29
30/// Read a 4-byte "batch" from some position.
31///
32/// This will read a native-endian 4-byte integer from some position.
33#[inline]
34#[cfg(not(feature = "safe-encode"))]
35pub(super) fn get_batch(input: &[u8], n: usize) -> u32 {
36    unsafe { read_u32_ptr(input.as_ptr().add(n)) }
37}
38
39#[inline]
40#[cfg(feature = "safe-encode")]
41pub(super) fn get_batch(input: &[u8], n: usize) -> u32 {
42    u32::from_ne_bytes(input[n..n + 4].try_into().unwrap())
43}
44
45/// Read an usize sized "batch" from some position.
46///
47/// This will read a native-endian usize from some position.
48#[inline]
49#[allow(dead_code)]
50#[cfg(not(feature = "safe-encode"))]
51pub(super) fn get_batch_arch(input: &[u8], n: usize) -> usize {
52    unsafe { read_usize_ptr(input.as_ptr().add(n)) }
53}
54
55#[inline]
56#[allow(dead_code)]
57#[cfg(feature = "safe-encode")]
58pub(super) fn get_batch_arch(input: &[u8], n: usize) -> usize {
59    const USIZE_SIZE: usize = core::mem::size_of::<usize>();
60    let arr: &[u8; USIZE_SIZE] = input[n..n + USIZE_SIZE].try_into().unwrap();
61    usize::from_ne_bytes(*arr)
62}
63
64#[inline]
65fn token_from_literal(lit_len: usize) -> u8 {
66    if lit_len < 0xF {
67        // Since we can fit the literals length into it, there is no need for saturation.
68        (lit_len as u8) << 4
69    } else {
70        // We were unable to fit the literals into it, so we saturate to 0xF. We will later
71        // write the extensional value.
72        0xF0
73    }
74}
75
76#[inline]
77fn token_from_literal_and_match_length(lit_len: usize, duplicate_length: usize) -> u8 {
78    let mut token = if lit_len < 0xF {
79        // Since we can fit the literals length into it, there is no need for saturation.
80        (lit_len as u8) << 4
81    } else {
82        // We were unable to fit the literals into it, so we saturate to 0xF. We will later
83        // write the extensional value.
84        0xF0
85    };
86
87    token |= if duplicate_length < 0xF {
88        // We could fit it in.
89        duplicate_length as u8
90    } else {
91        // We were unable to fit it in, so we default to 0xF, which will later be extended.
92        0xF
93    };
94
95    token
96}
97
98/// Counts the number of same bytes in two byte streams.
99/// `input` is the complete input
100/// `cur` is the current position in the input. it will be incremented by the number of matched
101/// bytes `source` either the same as input or an external slice
102/// `candidate` is the candidate position in `source`
103///
104/// The function ignores the last END_OFFSET bytes in input as those should be literals.
105#[inline]
106#[cfg(feature = "safe-encode")]
107fn count_same_bytes(input: &[u8], cur: &mut usize, source: &[u8], candidate: usize) -> usize {
108    const USIZE_SIZE: usize = core::mem::size_of::<usize>();
109    let cur_slice = &input[*cur..input.len() - END_OFFSET];
110    let cand_slice = &source[candidate..];
111
112    let mut num = 0;
113    for (block1, block2) in cur_slice
114        .chunks_exact(USIZE_SIZE)
115        .zip(cand_slice.chunks_exact(USIZE_SIZE))
116    {
117        let input_block = usize::from_ne_bytes(block1.try_into().unwrap());
118        let match_block = usize::from_ne_bytes(block2.try_into().unwrap());
119
120        if input_block == match_block {
121            num += USIZE_SIZE;
122        } else {
123            let diff = input_block ^ match_block;
124            num += (diff.to_le().trailing_zeros() / 8) as usize;
125            *cur += num;
126            return num;
127        }
128    }
129
130    // If we're here we may have 1 to 7 bytes left to check close to the end of input
131    // or source slices. Since this is rare occurrence we mark it cold to get better
132    // ~5% better performance.
133    #[cold]
134    fn count_same_bytes_tail(a: &[u8], b: &[u8], offset: usize) -> usize {
135        a.iter()
136            .zip(b)
137            .skip(offset)
138            .take_while(|(a, b)| a == b)
139            .count()
140    }
141    num += count_same_bytes_tail(cur_slice, cand_slice, num);
142
143    *cur += num;
144    num
145}
146
147/// Counts the number of same bytes in two byte streams.
148/// `input` is the complete input
149/// `cur` is the current position in the input. it will be incremented by the number of matched
150/// bytes `source` either the same as input OR an external slice
151/// `candidate` is the candidate position in `source`
152///
153/// The function ignores the last END_OFFSET bytes in input as those should be literals.
154#[inline]
155#[cfg(not(feature = "safe-encode"))]
156fn count_same_bytes(input: &[u8], cur: &mut usize, source: &[u8], candidate: usize) -> usize {
157    let max_input_match = input.len().saturating_sub(*cur + END_OFFSET);
158    let max_candidate_match = source.len() - candidate;
159    // Considering both limits calc how far we may match in input.
160    let input_end = *cur + max_input_match.min(max_candidate_match);
161
162    let start = *cur;
163    let mut source_ptr = unsafe { source.as_ptr().add(candidate) };
164
165    // compare 4/8 bytes blocks depending on the arch
166    const STEP_SIZE: usize = core::mem::size_of::<usize>();
167    while *cur + STEP_SIZE <= input_end {
168        let diff = read_usize_ptr(unsafe { input.as_ptr().add(*cur) }) ^ read_usize_ptr(source_ptr);
169
170        if diff == 0 {
171            *cur += STEP_SIZE;
172            unsafe {
173                source_ptr = source_ptr.add(STEP_SIZE);
174            }
175        } else {
176            *cur += (diff.to_le().trailing_zeros() / 8) as usize;
177            return *cur - start;
178        }
179    }
180
181    // compare 4 bytes block
182    #[cfg(target_pointer_width = "64")]
183    {
184        if input_end - *cur >= 4 {
185            let diff = read_u32_ptr(unsafe { input.as_ptr().add(*cur) }) ^ read_u32_ptr(source_ptr);
186
187            if diff == 0 {
188                *cur += 4;
189                unsafe {
190                    source_ptr = source_ptr.add(4);
191                }
192            } else {
193                *cur += (diff.to_le().trailing_zeros() / 8) as usize;
194                return *cur - start;
195            }
196        }
197    }
198
199    // compare 2 bytes block
200    if input_end - *cur >= 2
201        && unsafe { read_u16_ptr(input.as_ptr().add(*cur)) == read_u16_ptr(source_ptr) }
202    {
203        *cur += 2;
204        unsafe {
205            source_ptr = source_ptr.add(2);
206        }
207    }
208
209    if *cur < input_end
210        && unsafe { input.as_ptr().add(*cur).read() } == unsafe { source_ptr.read() }
211    {
212        *cur += 1;
213    }
214
215    *cur - start
216}
217
218/// Write an integer to the output.
219///
220/// Each additional byte then represent a value from 0 to 255, which is added to the previous value
221/// to produce a total length. When the byte value is 255, another byte must read and added, and so
222/// on. There can be any number of bytes of value "255" following token
223#[inline]
224pub(super) fn write_integer(output: &mut impl Sink, mut n: usize) {
225    // Note: Since `n` is usually < 0xFF and writing multiple bytes to the output
226    // requires 2 branches of bound check (due to the possibility of add overflows)
227    // the simple byte at a time implementation below is faster in most cases.
228    while n >= 0xFF {
229        n -= 0xFF;
230        push_byte(output, 0xFF);
231    }
232    push_byte(output, n as u8);
233}
234
235/// Handle the last bytes from the input as literals
236#[cold]
237fn handle_last_literals(output: &mut impl Sink, input: &[u8], start: usize) {
238    let lit_len = input.len() - start;
239
240    let token = token_from_literal(lit_len);
241    push_byte(output, token);
242    if lit_len >= 0xF {
243        write_integer(output, lit_len - 0xF);
244    }
245    // Now, write the actual literals.
246    output.extend_from_slice(&input[start..]);
247}
248
249/// Moves the cursors back as long as the bytes match, to find additional bytes in a duplicate
250#[inline]
251#[cfg(feature = "safe-encode")]
252fn backtrack_match(
253    input: &[u8],
254    cur: &mut usize,
255    literal_start: usize,
256    source: &[u8],
257    candidate: &mut usize,
258) {
259    // Note: Even if iterator version of this loop has less branches inside the loop it has more
260    // branches before the loop. That in practice seems to make it slower than the while version
261    // bellow. TODO: It should be possible remove all bounds checks, since we are walking
262    // backwards
263    while *candidate > 0 && *cur > literal_start && input[*cur - 1] == source[*candidate - 1] {
264        *cur -= 1;
265        *candidate -= 1;
266    }
267}
268
269/// Moves the cursors back as long as the bytes match, to find additional bytes in a duplicate
270#[inline]
271#[cfg(not(feature = "safe-encode"))]
272fn backtrack_match(
273    input: &[u8],
274    cur: &mut usize,
275    literal_start: usize,
276    source: &[u8],
277    candidate: &mut usize,
278) {
279    while unsafe {
280        *candidate > 0
281            && *cur > literal_start
282            && input.get_unchecked(*cur - 1) == source.get_unchecked(*candidate - 1)
283    } {
284        *cur -= 1;
285        *candidate -= 1;
286    }
287}
288
289/// Compress all bytes of `input[input_pos..]` into `output`.
290///
291/// Bytes in `input[..input_pos]` are treated as a preamble and can be used for lookback.
292/// This part is known as the compressor "prefix".
293/// Bytes in `ext_dict` logically precede the bytes in `input` and can also be used for lookback.
294///
295/// `input_stream_offset` is the logical position of the first byte of `input`. This allows same
296/// `dict` to be used for many calls to `compress_internal` as we can "readdress" the first byte of
297/// `input` to be something other than 0.
298///
299/// `dict` is the dictionary of previously encoded sequences.
300///
301/// This is used to find duplicates in the stream so they are not written multiple times.
302///
303/// Every four bytes are hashed, and in the resulting slot their position in the input buffer
304/// is placed in the dict. This way we can easily look up a candidate to back references.
305///
306/// Returns the number of bytes written (compressed) into `output`.
307///
308/// # Const parameters
309/// `USE_DICT`: Disables usage of ext_dict (it'll panic if a non-empty slice is used).
310/// In other words, this generates more optimized code when an external dictionary isn't used.
311///
312/// A similar const argument could be used to disable the Prefix mode (eg. USE_PREFIX),
313/// which would impose `input_pos == 0 && input_stream_offset == 0`. Experiments didn't
314/// show significant improvement though.
315// Intentionally avoid inlining.
316// Empirical tests revealed it to be rarely better but often significantly detrimental.
317#[inline(never)]
318pub(crate) fn compress_internal<T: HashTable, const USE_DICT: bool, S: Sink>(
319    input: &[u8],
320    input_pos: usize,
321    output: &mut S,
322    dict: &mut T,
323    ext_dict: &[u8],
324    input_stream_offset: usize,
325) -> Result<usize, CompressError> {
326    assert!(input_pos <= input.len());
327    if USE_DICT {
328        assert!(ext_dict.len() <= super::WINDOW_SIZE);
329        assert!(ext_dict.len() <= input_stream_offset);
330        // Check for overflow hazard when using ext_dict
331        assert!(input_stream_offset
332            .checked_add(input.len())
333            .and_then(|i| i.checked_add(ext_dict.len()))
334            .is_some_and(|i| i <= isize::MAX as usize));
335    } else {
336        assert!(ext_dict.is_empty());
337    }
338    if output.capacity() - output.pos() < get_maximum_output_size(input.len() - input_pos) {
339        return Err(CompressError::OutputTooSmall);
340    }
341
342    let output_start_pos = output.pos();
343    if input.len() - input_pos < LZ4_MIN_LENGTH {
344        handle_last_literals(output, input, input_pos);
345        return Ok(output.pos() - output_start_pos);
346    }
347
348    let ext_dict_stream_offset = input_stream_offset - ext_dict.len();
349    let end_pos_check = input.len() - MFLIMIT;
350    let mut literal_start = input_pos;
351    let mut cur = input_pos;
352
353    if cur == 0 && input_stream_offset == 0 {
354        // According to the spec we can't start with a match,
355        // except when referencing another block.
356        let hash = T::get_hash_at(input, 0);
357        dict.put_at(hash, 0);
358        cur = 1;
359    }
360
361    loop {
362        // Read the next block into two sections, the literals and the duplicates.
363        let mut step_size;
364        let mut candidate;
365        let mut candidate_source;
366        let mut offset;
367        let mut non_match_count = 1 << INCREASE_STEPSIZE_BITSHIFT;
368        // The number of bytes before our cursor, where the duplicate starts.
369        let mut next_cur = cur;
370
371        // In this loop we search for duplicates via the hashtable. 4bytes or 8bytes are hashed and
372        // compared.
373        loop {
374            step_size = non_match_count >> INCREASE_STEPSIZE_BITSHIFT;
375            non_match_count += 1;
376
377            cur = next_cur;
378            next_cur += step_size;
379
380            // Same as cur + MFLIMIT > input.len()
381            if cur > end_pos_check {
382                handle_last_literals(output, input, literal_start);
383                return Ok(output.pos() - output_start_pos);
384            }
385            // Find a candidate in the dictionary with the hash of the current four bytes.
386            // Unchecked is safe as long as the values from the hash function don't exceed the size
387            // of the table. This is ensured by right shifting the hash values
388            // (`dict_bitshift`) to fit them in the table
389
390            // [Bounds Check]: Can be elided due to `end_pos_check` above
391            let hash = T::get_hash_at(input, cur);
392            candidate = dict.get_at(hash);
393            dict.put_at(hash, cur + input_stream_offset);
394
395            // Sanity check: Matches can't be ahead of `cur`.
396            debug_assert!(candidate <= input_stream_offset + cur);
397
398            // Two requirements to the candidate exists:
399            // - We should not return a position which is merely a hash collision, so that the
400            //   candidate actually matches what we search for.
401            // - We can address up to 16-bit offset, hence we are only able to address the candidate
402            //   if its offset is less than or equals to 0xFFFF.
403            if input_stream_offset + cur - candidate > MAX_DISTANCE {
404                continue;
405            }
406
407            if candidate >= input_stream_offset {
408                // match within input
409                offset = (input_stream_offset + cur - candidate) as u16;
410                candidate -= input_stream_offset;
411                candidate_source = input;
412            } else if USE_DICT {
413                // Sanity check, which may fail if we lost history beyond MAX_DISTANCE
414                debug_assert!(
415                    candidate >= ext_dict_stream_offset,
416                    "Lost history in ext dict mode"
417                );
418                // match within ext dict
419                offset = (input_stream_offset + cur - candidate) as u16;
420                candidate -= ext_dict_stream_offset;
421                candidate_source = ext_dict;
422            } else {
423                // Match is not reachable anymore
424                // eg. compressing an independent block frame w/o clearing
425                // the matches tables, only increasing input_stream_offset.
426                // Sanity check
427                debug_assert!(input_pos == 0, "Lost history in prefix mode");
428                continue;
429            }
430            // [Bounds Check]: Candidate is coming from the Hashmap. It can't be out of bounds, but
431            // impossible to prove for the compiler and remove the bounds checks.
432            let cand_bytes: u32 = get_batch(candidate_source, candidate);
433            // [Bounds Check]: Should be able to be elided due to `end_pos_check`.
434            let curr_bytes: u32 = get_batch(input, cur);
435
436            if cand_bytes == curr_bytes {
437                break;
438            }
439        }
440
441        // Extend the match backwards if we can
442        backtrack_match(
443            input,
444            &mut cur,
445            literal_start,
446            candidate_source,
447            &mut candidate,
448        );
449
450        // The length (in bytes) of the literals section.
451        let lit_len = cur - literal_start;
452
453        // Generate the higher half of the token.
454        cur += MINMATCH;
455        candidate += MINMATCH;
456        let duplicate_length = count_same_bytes(input, &mut cur, candidate_source, candidate);
457
458        // Note: The `- 2` offset was copied from the reference implementation, it could be
459        // arbitrary.
460        let hash = T::get_hash_at(input, cur - 2);
461        dict.put_at(hash, cur - 2 + input_stream_offset);
462
463        let token = token_from_literal_and_match_length(lit_len, duplicate_length);
464
465        // Push the token to the output stream.
466        push_byte(output, token);
467        // If we were unable to fit the literals length into the token, write the extensional
468        // part.
469        if lit_len >= 0xF {
470            write_integer(output, lit_len - 0xF);
471        }
472
473        // Now, write the actual literals.
474        //
475        // The unsafe version copies blocks of 8bytes, and therefore may copy up to 7bytes more than
476        // needed. This is safe, because the last 12 bytes (MF_LIMIT) are handled in
477        // handle_last_literals.
478        copy_literals_wild(output, input, literal_start, lit_len);
479        // write the offset in little endian.
480        push_u16(output, offset);
481
482        // If we were unable to fit the duplicates length into the token, write the
483        // extensional part.
484        if duplicate_length >= 0xF {
485            write_integer(output, duplicate_length - 0xF);
486        }
487        literal_start = cur;
488    }
489}
490
491#[inline]
492#[cfg(feature = "safe-encode")]
493fn push_byte(output: &mut impl Sink, el: u8) {
494    output.push(el);
495}
496
497#[inline]
498#[cfg(not(feature = "safe-encode"))]
499fn push_byte(output: &mut impl Sink, el: u8) {
500    unsafe {
501        core::ptr::write(output.pos_mut_ptr(), el);
502        output.set_pos(output.pos() + 1);
503    }
504}
505
506#[inline]
507#[cfg(feature = "safe-encode")]
508fn push_u16(output: &mut impl Sink, el: u16) {
509    output.extend_from_slice(&el.to_le_bytes());
510}
511
512#[inline]
513#[cfg(not(feature = "safe-encode"))]
514fn push_u16(output: &mut impl Sink, el: u16) {
515    unsafe {
516        core::ptr::copy_nonoverlapping(el.to_le_bytes().as_ptr(), output.pos_mut_ptr(), 2);
517        output.set_pos(output.pos() + 2);
518    }
519}
520
521#[inline(always)] // (always) necessary otherwise compiler fails to inline it
522#[cfg(feature = "safe-encode")]
523fn copy_literals_wild(output: &mut impl Sink, input: &[u8], input_start: usize, len: usize) {
524    output.extend_from_slice_wild(&input[input_start..input_start + len], len)
525}
526
527#[inline]
528#[cfg(not(feature = "safe-encode"))]
529fn copy_literals_wild(output: &mut impl Sink, input: &[u8], input_start: usize, len: usize) {
530    debug_assert!(input_start + len / 8 * 8 + ((len % 8) != 0) as usize * 8 <= input.len());
531    debug_assert!(output.pos() + len / 8 * 8 + ((len % 8) != 0) as usize * 8 <= output.capacity());
532    unsafe {
533        // Note: This used to be a wild copy loop of 8 bytes, but the compiler consistently
534        // transformed it into a call to memcopy, which hurts performance significantly for
535        // small copies, which are common.
536        let start_ptr = input.as_ptr().add(input_start);
537        match len {
538            0..=8 => core::ptr::copy_nonoverlapping(start_ptr, output.pos_mut_ptr(), 8),
539            9..=16 => core::ptr::copy_nonoverlapping(start_ptr, output.pos_mut_ptr(), 16),
540            17..=24 => core::ptr::copy_nonoverlapping(start_ptr, output.pos_mut_ptr(), 24),
541            _ => core::ptr::copy_nonoverlapping(start_ptr, output.pos_mut_ptr(), len),
542        }
543        output.set_pos(output.pos() + len);
544    }
545}
546
547/// Compress all bytes of `input` into `output`.
548/// The method chooses an appropriate hashtable to lookup duplicates.
549/// output should be preallocated with a size of
550/// `get_maximum_output_size`.
551///
552/// Returns the number of bytes written (compressed) into `output`.
553#[inline]
554pub(crate) fn compress_into_sink_with_dict<const USE_DICT: bool>(
555    input: &[u8],
556    output: &mut impl Sink,
557    mut dict_data: &[u8],
558) -> Result<usize, CompressError> {
559    if USE_DICT && dict_data.len() < MINMATCH {
560        return compress_into_sink_without_dict(input, output);
561    }
562
563    if dict_data.len() + input.len() < u16::MAX as usize {
564        let mut dict = HashTable4KU16::new();
565        init_dict(&mut dict, &mut dict_data);
566        compress_internal::<_, USE_DICT, _>(input, 0, output, &mut dict, dict_data, dict_data.len())
567    } else {
568        let mut dict = HashTable4K::new();
569        init_dict(&mut dict, &mut dict_data);
570        compress_internal::<_, USE_DICT, _>(input, 0, output, &mut dict, dict_data, dict_data.len())
571    }
572}
573
574/// Slow fallback for when the dictionary is too small to be useful. This avoids the overhead of
575/// inlining.
576#[cold]
577#[inline(never)]
578fn compress_into_sink_without_dict(
579    input: &[u8],
580    output: &mut impl Sink,
581) -> Result<usize, CompressError> {
582    compress_into_sink_with_dict::<false>(input, output, b"")
583}
584
585#[inline]
586fn init_dict<T: HashTable>(dict: &mut T, dict_data: &mut &[u8]) {
587    if dict_data.len() > WINDOW_SIZE {
588        *dict_data = &dict_data[dict_data.len() - WINDOW_SIZE..];
589    }
590    let mut i = 0usize;
591    while i + core::mem::size_of::<usize>() <= dict_data.len() {
592        let hash = T::get_hash_at(dict_data, i);
593        dict.put_at(hash, i);
594        // Note: The 3 byte step was copied from the reference implementation, it could be
595        // arbitrary.
596        i += 3;
597    }
598}
599
600/// Returns the maximum output size of the compressed data.
601/// Can be used to preallocate capacity on the output vector
602#[inline]
603pub const fn get_maximum_output_size(input_len: usize) -> usize {
604    16 + 4 + (input_len as u64 * 110 / 100) as usize
605}
606
607/// Compress all bytes of `input` into `output`.
608/// The method chooses an appropriate hashtable to lookup duplicates.
609/// output should be preallocated with a size of
610/// `get_maximum_output_size`.
611///
612/// Returns the number of bytes written (compressed) into `output`.
613#[inline]
614pub fn compress_into(input: &[u8], output: &mut [u8]) -> Result<usize, CompressError> {
615    compress_into_sink_with_dict::<false>(input, &mut SliceSink::new(output, 0), b"")
616}
617
618/// Compress all bytes of `input` into `output`.
619/// The method chooses an appropriate hashtable to lookup duplicates.
620/// output should be preallocated with a size of
621/// `get_maximum_output_size`.
622///
623/// Returns the number of bytes written (compressed) into `output`.
624#[inline]
625pub fn compress_into_with_dict(
626    input: &[u8],
627    output: &mut [u8],
628    dict_data: &[u8],
629) -> Result<usize, CompressError> {
630    compress_into_sink_with_dict::<true>(input, &mut SliceSink::new(output, 0), dict_data)
631}
632
633#[inline]
634fn compress_into_vec_with_dict<const USE_DICT: bool>(
635    input: &[u8],
636    prepend_size: bool,
637    dict_data: &[u8],
638) -> Vec<u8> {
639    let prepend_size_num_bytes = if prepend_size { 4 } else { 0 };
640    let max_compressed_size = get_maximum_output_size(input.len()) + prepend_size_num_bytes;
641    if USE_DICT && dict_data.len() < MINMATCH {
642        return compress_into_vec_without_dict(input, prepend_size);
643    }
644    #[cfg(feature = "safe-encode")]
645    let mut compressed = {
646        let mut compressed: Vec<u8> = vec![0u8; max_compressed_size];
647        let out = if prepend_size {
648            compressed[..4].copy_from_slice(&(input.len() as u32).to_le_bytes());
649            &mut compressed[4..]
650        } else {
651            &mut compressed
652        };
653        let compressed_len =
654            compress_into_sink_with_dict::<USE_DICT>(input, &mut SliceSink::new(out, 0), dict_data)
655                .unwrap();
656
657        compressed.truncate(prepend_size_num_bytes + compressed_len);
658        compressed
659    };
660    #[cfg(not(feature = "safe-encode"))]
661    let mut compressed = {
662        let mut vec = Vec::with_capacity(max_compressed_size);
663        let start_pos = if prepend_size {
664            vec.extend_from_slice(&(input.len() as u32).to_le_bytes());
665            4
666        } else {
667            0
668        };
669        let compressed_len = compress_into_sink_with_dict::<USE_DICT>(
670            input,
671            &mut PtrSink::from_vec(&mut vec, start_pos),
672            dict_data,
673        )
674        .unwrap();
675        unsafe {
676            vec.set_len(prepend_size_num_bytes + compressed_len);
677        }
678        vec
679    };
680
681    compressed.shrink_to_fit();
682    compressed
683}
684
685#[cold]
686#[inline(never)]
687fn compress_into_vec_without_dict(input: &[u8], prepend_size: bool) -> Vec<u8> {
688    compress_into_vec_with_dict::<false>(input, prepend_size, b"")
689}
690
691/// Compress all bytes of `input` into `output`. The uncompressed size will be prepended as a little
692/// endian u32. Can be used in conjunction with `decompress_size_prepended`
693#[inline]
694pub fn compress_prepend_size(input: &[u8]) -> Vec<u8> {
695    compress_into_vec_with_dict::<false>(input, true, b"")
696}
697
698/// Compress all bytes of `input`.
699#[inline]
700pub fn compress(input: &[u8]) -> Vec<u8> {
701    compress_into_vec_with_dict::<false>(input, false, b"")
702}
703
704/// Compress all bytes of `input` with an external dictionary.
705#[inline]
706pub fn compress_with_dict(input: &[u8], ext_dict: &[u8]) -> Vec<u8> {
707    compress_into_vec_with_dict::<true>(input, false, ext_dict)
708}
709
710/// Compress all bytes of `input` into `output`. The uncompressed size will be prepended as a little
711/// endian u32. Can be used in conjunction with `decompress_size_prepended_with_dict`
712#[inline]
713pub fn compress_prepend_size_with_dict(input: &[u8], ext_dict: &[u8]) -> Vec<u8> {
714    compress_into_vec_with_dict::<true>(input, true, ext_dict)
715}
716
717/// A reusable compression table that avoids re-allocating the internal hash table on every call.
718///
719/// This is useful when compressing many small inputs in a loop. Create one table and pass it
720/// to [`compress_into_with_table`] repeatedly.
721///
722/// # Example
723/// ```
724/// use lz4_flex::block::{compress_into_with_table, get_maximum_output_size, CompressTable};
725///
726/// let mut table = CompressTable::default();
727/// let input = b"hello world, hello world, hello!";
728/// let mut output = vec![0u8; get_maximum_output_size(input.len())];
729/// let compressed_len = compress_into_with_table(input, &mut output, &mut table).unwrap();
730/// ```
731pub enum CompressTable {
732    /// Table using 16-bit entries, suitable for inputs where `input.len() < u16::MAX`.
733    Small(HashTable4KU16),
734    /// Table using 32-bit entries, suitable for any input size.
735    Large(HashTable4K),
736}
737
738impl Default for CompressTable {
739    fn default() -> Self {
740        CompressTable::Small(HashTable4KU16::new())
741    }
742}
743
744impl CompressTable {
745    /// Create a small table (16-bit entries). More memory efficient, but only usable when the
746    /// total input size is less than 65535 bytes.
747    pub fn small() -> Self {
748        CompressTable::Small(HashTable4KU16::new())
749    }
750
751    /// Create a large table (32-bit entries). Works for any input size.
752    pub fn large() -> Self {
753        CompressTable::Large(HashTable4K::new())
754    }
755}
756
757/// Compress all bytes of `input` into `output`, reusing a [`CompressTable`] to avoid
758/// re-allocating the internal hash table.
759///
760/// `output` should be preallocated with a size of [`get_maximum_output_size`].
761///
762/// Returns the number of bytes written (compressed) into `output`.
763///
764/// **Note:** If the table variant doesn't match the input size (e.g. a `Small` table is used
765/// with input >= 64KB), the table will be transparently upgraded. However, it won't be
766/// downgraded automatically.
767#[inline]
768pub fn compress_into_with_table(
769    input: &[u8],
770    output: &mut [u8],
771    table: &mut CompressTable,
772) -> Result<usize, CompressError> {
773    if input.len() >= u16::MAX as usize && matches!(table, CompressTable::Small(_)) {
774        *table = CompressTable::Large(HashTable4K::new());
775    }
776
777    match table {
778        CompressTable::Small(dict) => {
779            dict.clear();
780            compress_internal::<_, false, _>(input, 0, &mut SliceSink::new(output, 0), dict, b"", 0)
781        }
782        CompressTable::Large(dict) => {
783            dict.clear();
784            compress_internal::<_, false, _>(input, 0, &mut SliceSink::new(output, 0), dict, b"", 0)
785        }
786    }
787}
788
789#[inline]
790#[cfg(not(feature = "safe-encode"))]
791fn read_u16_ptr(input: *const u8) -> u16 {
792    let mut num: u16 = 0;
793    unsafe {
794        core::ptr::copy_nonoverlapping(input, &mut num as *mut u16 as *mut u8, 2);
795    }
796    num
797}
798
799#[inline]
800#[cfg(not(feature = "safe-encode"))]
801fn read_u32_ptr(input: *const u8) -> u32 {
802    let mut num: u32 = 0;
803    unsafe {
804        core::ptr::copy_nonoverlapping(input, &mut num as *mut u32 as *mut u8, 4);
805    }
806    num
807}
808
809#[inline]
810#[cfg(not(feature = "safe-encode"))]
811fn read_usize_ptr(input: *const u8) -> usize {
812    let mut num: usize = 0;
813    unsafe {
814        core::ptr::copy_nonoverlapping(
815            input,
816            &mut num as *mut usize as *mut u8,
817            core::mem::size_of::<usize>(),
818        );
819    }
820    num
821}
822
823#[cfg(test)]
824mod tests {
825    use super::*;
826
827    #[test]
828    fn test_count_same_bytes() {
829        // 8byte aligned block, zeros and ones are added because the end/offset
830        let first: &[u8] = &[
831            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,
832        ];
833        let second: &[u8] = &[
834            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,
835        ];
836        assert_eq!(count_same_bytes(first, &mut 0, second, 0), 16);
837
838        // 4byte aligned block
839        let first: &[u8] = &[
840            1, 2, 3, 4, 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,
841            0, 0, 0,
842        ];
843        let second: &[u8] = &[
844            1, 2, 3, 4, 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,
845            1, 1, 1,
846        ];
847        assert_eq!(count_same_bytes(first, &mut 0, second, 0), 20);
848
849        // 2byte aligned block
850        let first: &[u8] = &[
851            1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 3, 4, 0, 0, 0, 0, 0, 0, 0,
852            0, 0, 0, 0, 0,
853        ];
854        let second: &[u8] = &[
855            1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 3, 4, 1, 1, 1, 1, 1, 1, 1,
856            1, 1, 1, 1, 1,
857        ];
858        assert_eq!(count_same_bytes(first, &mut 0, second, 0), 22);
859
860        // 1byte aligned block
861        let first: &[u8] = &[
862            1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 3, 4, 5, 0, 0, 0, 0, 0, 0,
863            0, 0, 0, 0, 0, 0,
864        ];
865        let second: &[u8] = &[
866            1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 3, 4, 5, 1, 1, 1, 1, 1, 1,
867            1, 1, 1, 1, 1, 1,
868        ];
869        assert_eq!(count_same_bytes(first, &mut 0, second, 0), 23);
870
871        // 1byte aligned block - last byte different
872        let first: &[u8] = &[
873            1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 3, 4, 5, 0, 0, 0, 0, 0, 0,
874            0, 0, 0, 0, 0, 0,
875        ];
876        let second: &[u8] = &[
877            1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 3, 4, 6, 1, 1, 1, 1, 1, 1,
878            1, 1, 1, 1, 1, 1,
879        ];
880        assert_eq!(count_same_bytes(first, &mut 0, second, 0), 22);
881
882        // 1byte aligned block
883        let first: &[u8] = &[
884            1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 3, 9, 5, 0, 0, 0, 0, 0, 0,
885            0, 0, 0, 0, 0, 0,
886        ];
887        let second: &[u8] = &[
888            1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 3, 4, 6, 1, 1, 1, 1, 1, 1,
889            1, 1, 1, 1, 1, 1,
890        ];
891        assert_eq!(count_same_bytes(first, &mut 0, second, 0), 21);
892
893        for diff_idx in 8..100 {
894            let first: Vec<u8> = (0u8..255).cycle().take(100 + 12).collect();
895            let mut second = first.clone();
896            second[diff_idx] = 255;
897            for start in 0..=diff_idx {
898                let same_bytes = count_same_bytes(&first, &mut start.clone(), &second, start);
899                assert_eq!(same_bytes, diff_idx - start);
900            }
901        }
902    }
903
904    #[test]
905    fn test_bug() {
906        let input: &[u8] = &[
907            10, 12, 14, 16, 18, 10, 12, 14, 16, 18, 10, 12, 14, 16, 18, 10, 12, 14, 16, 18,
908        ];
909        let _out = compress(input);
910    }
911
912    #[test]
913    fn test_dict() {
914        let input: &[u8] = &[
915            10, 12, 14, 16, 18, 10, 12, 14, 16, 18, 10, 12, 14, 16, 18, 10, 12, 14, 16, 18,
916        ];
917        let dict = input;
918        let compressed = compress_with_dict(input, dict);
919        assert_lt!(compressed.len(), compress(input).len());
920
921        assert!(compressed.len() < compress(input).len());
922        let mut uncompressed = vec![0u8; input.len()];
923        let uncomp_size = crate::block::decompress::decompress_into_with_dict(
924            &compressed,
925            &mut uncompressed,
926            dict,
927        )
928        .unwrap();
929        uncompressed.truncate(uncomp_size);
930        assert_eq!(input, uncompressed);
931    }
932
933    #[test]
934    fn test_dict_no_panic() {
935        let input: &[u8] = &[
936            10, 12, 14, 16, 18, 10, 12, 14, 16, 18, 10, 12, 14, 16, 18, 10, 12, 14, 16, 18,
937        ];
938        let dict = &[10, 12, 14];
939        let _compressed = compress_with_dict(input, dict);
940    }
941
942    #[test]
943    fn compress_into_with_short_dict_does_not_panic() {
944        let input = [0u8; 13];
945
946        for dict_len in 0..MINMATCH {
947            let dict = vec![0u8; dict_len];
948            let mut output = vec![0u8; get_maximum_output_size(input.len())];
949            let compressed_len = compress_into_with_dict(&input, &mut output, &dict).unwrap();
950
951            let mut uncompressed = vec![0u8; input.len()];
952            let uncompressed_len = crate::block::decompress::decompress_into_with_dict(
953                &output[..compressed_len],
954                &mut uncompressed,
955                &dict,
956            )
957            .unwrap();
958            uncompressed.truncate(uncompressed_len);
959            assert_eq!(uncompressed, input);
960        }
961    }
962
963    #[test]
964    #[cfg(all(miri, not(feature = "safe-encode")))]
965    fn miri_compress_into_with_short_dict_reads_past_dict() {
966        let input = [0u8; 13];
967        let dict = [0u8; 1];
968        let mut output = vec![0u8; get_maximum_output_size(input.len())];
969
970        let _ = compress_into_with_dict(&input, &mut output, &dict);
971    }
972
973    #[test]
974    fn test_dict_match_crossing() {
975        let input: &[u8] = &[
976            10, 12, 14, 16, 18, 10, 12, 14, 16, 18, 10, 12, 14, 16, 18, 10, 12, 14, 16, 18,
977        ];
978        let dict = input;
979        let compressed = compress_with_dict(input, dict);
980        assert_lt!(compressed.len(), compress(input).len());
981
982        let mut uncompressed = vec![0u8; input.len() * 2];
983        // copy first half of the input into output
984        let dict_cutoff = dict.len() / 2;
985        let output_start = dict.len() - dict_cutoff;
986        uncompressed[..output_start].copy_from_slice(&dict[dict_cutoff..]);
987        let uncomp_len = {
988            let mut sink = SliceSink::new(&mut uncompressed[..], output_start);
989            crate::block::decompress::decompress_internal::<true, _>(
990                &compressed,
991                &mut sink,
992                &dict[..dict_cutoff],
993            )
994            .unwrap()
995        };
996        assert_eq!(input.len(), uncomp_len);
997        assert_eq!(
998            input,
999            &uncompressed[output_start..output_start + uncomp_len]
1000        );
1001    }
1002
1003    #[test]
1004    fn test_conformant_last_block() {
1005        // From the spec:
1006        // The last match must start at least 12 bytes before the end of block.
1007        // The last match is part of the penultimate sequence. It is followed by the last sequence,
1008        // which contains only literals. Note that, as a consequence, an independent block <
1009        // 13 bytes cannot be compressed, because the match must copy "something",
1010        // so it needs at least one prior byte.
1011        // When a block can reference data from another block, it can start immediately with a match
1012        // and no literal, so a block of 12 bytes can be compressed.
1013        let aaas: &[u8] = b"aaaaaaaaaaaaaaa";
1014
1015        // incompressible
1016        let out = compress(&aaas[..12]);
1017        assert_gt!(out.len(), 12);
1018        // compressible
1019        let out = compress(&aaas[..13]);
1020        assert_le!(out.len(), 13);
1021        let out = compress(&aaas[..14]);
1022        assert_le!(out.len(), 14);
1023        let out = compress(&aaas[..15]);
1024        assert_le!(out.len(), 15);
1025
1026        // dict incompressible
1027        let out = compress_with_dict(&aaas[..11], aaas);
1028        assert_gt!(out.len(), 11);
1029        // compressible
1030        let out = compress_with_dict(&aaas[..12], aaas);
1031        // According to the spec this _could_ compress, but it doesn't in this lib
1032        // as it aborts compression for any input len < LZ4_MIN_LENGTH
1033        assert_gt!(out.len(), 12);
1034        let out = compress_with_dict(&aaas[..13], aaas);
1035        assert_le!(out.len(), 13);
1036        let out = compress_with_dict(&aaas[..14], aaas);
1037        assert_le!(out.len(), 14);
1038        let out = compress_with_dict(&aaas[..15], aaas);
1039        assert_le!(out.len(), 15);
1040    }
1041
1042    #[test]
1043    fn test_dict_size() {
1044        let dict = vec![b'a'; 1024 * 1024];
1045        let input = &b"aaaaaaaaaaaaaaaaaaaaaaaaaaaaa"[..];
1046        let compressed = compress_prepend_size_with_dict(input, &dict);
1047        let decompressed =
1048            crate::block::decompress_size_prepended_with_dict(&compressed, &dict).unwrap();
1049        assert_eq!(decompressed, input);
1050    }
1051}