krabka-compression 0.8.0

//! LZ4 frame format (LZ4F), independent blocks.
//!
//! Kafka writes LZ4 in the frame format, with magic `0x04 22 4D 18`, and makes
//! these choices: 64 KiB block size, independent blocks, no block checksum, and
//! no content-size in the header. We match those defaults, so our bytes agree
//! with the output of `KafkaLZ4BlockOutputStream` for differential testing.
//!
//! `KafkaLZ4BlockOutputStream` (`Lz4BlockOutputStream.java`) picks the block
//! compressor by level: the fast compressor at the default level (9), and LZ4
//! HC (`LZ4Factory.highCompressor(level)`) at every other level in `1..=17`.
//! `compress_with_level` matches that split. `lz4rip` has only the fast
//! compressor. Both paths build the same independent-block frame directly;
//! the HC path uses `lzzzz`'s binding to the reference `liblz4` HC compressor
//! for the block bytes. Both paths use [`FRAME_HEADER`] and the same end mark.
//!
//! Kafka's level range is `1..=17`, but reference `liblz4`'s HC compressor
//! caps out at `LZ4HC_CLEVEL_MAX` (12): `lz4hc.h` documents that "values
//! greater than `LZ4HC_CLEVEL_MAX` behave the same as `LZ4HC_CLEVEL_MAX`",
//! and `LZ4_compress_HC` clamps internally to the same effect. A level above
//! 12 therefore compresses exactly as hard as level 12 -- this binding has no
//! access to a distinct, more exhaustive parse for 13 through 17 the way
//! `lz4-java`'s own implementation might. [`write_hc_block`] clamps
//! explicitly rather than relying on `liblz4`'s internal clamp, so that fact
//! is visible at the call site instead of buried in the C library.

use std::{cell::RefCell, io::Read};

use bytes::Bytes;
use lz4rip::{
    block::{CompressorRef, compress_into, get_maximum_output_size},
    frame::FrameDecoder,
};
use lzzzz::{lz4::max_compressed_size, lz4_hc};

use crate::CompressionError;

/// Kafka's default lz4 level (`CompressionType.LZ4.defaultLevel()`). Only this
/// level uses the fast compressor; every other level in `1..=17` uses HC.
const DEFAULT_LEVEL: i32 = 9;

/// Kafka's independent-block max size (`Lz4BlockOutputStream` uses the LZ4F
/// default of 64 KiB), matched here so HC blocks line up with the fast path.
const BLOCK_SIZE: usize = 64 * 1024;

/// lz4rip emits only literals below 13 bytes, which cannot shrink a block.
const MIN_COMPRESSIBLE_BLOCK: usize = 13;

/// The high bit of a block's 4-byte little-endian size that marks it as
/// stored uncompressed, per the LZ4 frame format.
const BLOCK_UNCOMPRESSED_BIT: u32 = 0x8000_0000;

/// Independent 64 KiB blocks, without block/content checksums, content size or
/// dictionary. The final byte is the descriptor's header checksum.
const FRAME_HEADER: &[u8] = &[0x04, 0x22, 0x4D, 0x18, 0x60, 0x40, 0x82];

// Compress borrowed blocks into the result without a streaming source buffer
// and separate compressed-block buffer.
pub fn compress(data: &[u8]) -> Result<Bytes, CompressionError> {
    // Include framing and one block's compression slack so the scratch slice
    // does not double the result allocation for a nearly full block.
    let largest_block = data.len().min(BLOCK_SIZE);
    let slack = if largest_block < MIN_COMPRESSIBLE_BLOCK {
        0
    } else {
        get_maximum_output_size(largest_block) - largest_block
    };
    let mut out = frame_buffer(data, slack);
    for block in data.chunks(BLOCK_SIZE) {
        write_block(
            &mut out,
            block,
            get_maximum_output_size(block.len()),
            |scratch| {
                // At 65,535 bytes the stateless encoder switches to the u32
                // table. Reuse that same table; CompressorRef clears it for
                // these large inputs.
                let compressed = if block.len() >= usize::from(u16::MAX) {
                    thread_local! {
                        static COMPRESSOR: RefCell<CompressorRef> =
                            RefCell::new(CompressorRef::new());
                    }
                    COMPRESSOR
                        .with(|compressor| compressor.borrow_mut().compress_into(block, scratch))
                } else {
                    compress_into(block, scratch)
                };
                compressed.map_err(|e| CompressionError::InvalidData(format!("lz4 compress: {e}")))
            },
        )?;
    }
    out.extend_from_slice(&0u32.to_le_bytes());
    Ok(Bytes::from(out))
}

/// Compress at `level`, which the caller has checked against Kafka's `1..=17`
/// range. Level 9 (Kafka's default) gives the same bytes as [`compress`].
/// Every other level uses LZ4 HC, in the same independent-64-KiB-block frame.
pub fn compress_with_level(data: &[u8], level: i32) -> Result<Bytes, CompressionError> {
    if level == DEFAULT_LEVEL {
        return compress(data);
    }
    let largest_block = data.len().min(BLOCK_SIZE);
    let slack = if largest_block < MIN_COMPRESSIBLE_BLOCK {
        0
    } else {
        max_compressed_size(largest_block) - largest_block
    };
    let mut out = frame_buffer(data, slack);
    // `[T]::chunks` yields no chunks for empty input, matching `compress`,
    // which never writes a block for an empty frame either.
    for block in data.chunks(BLOCK_SIZE) {
        write_hc_block(&mut out, block, level)?;
    }
    out.extend_from_slice(&0u32.to_le_bytes()); // end mark
    Ok(Bytes::from(out))
}

/// A buffer holding [`FRAME_HEADER`], with room for the blocks of `data`,
/// their sizes, the end mark and `slack` bytes of compression scratch space.
fn frame_buffer(data: &[u8], slack: usize) -> Vec<u8> {
    let capacity = data
        .len()
        .saturating_add(data.len().div_ceil(BLOCK_SIZE).saturating_mul(4))
        .saturating_add(FRAME_HEADER.len() + 4 + slack);
    let mut out = Vec::with_capacity(capacity);
    out.extend_from_slice(FRAME_HEADER);
    out
}

/// `liblz4`'s `LZ4HC_CLEVEL_MAX`: the highest level its HC compressor
/// actually distinguishes. See the module doc for why levels above this
/// (Kafka goes up to 17) get no more effort than this level does.
const LIBLZ4_HC_CLEVEL_MAX: i32 = 12;

/// Append one LZ4F block (4-byte size, then the block bytes) compressed with
/// HC at `level`. Falls back to storing the block uncompressed, exactly as
/// the LZ4 frame format allows, when HC does not shrink it.
fn write_hc_block(out: &mut Vec<u8>, block: &[u8], level: i32) -> Result<(), CompressionError> {
    let level = level.min(LIBLZ4_HC_CLEVEL_MAX);
    write_block(out, block, max_compressed_size(block.len()), |scratch| {
        lz4_hc::compress(block, scratch, level)
            .map_err(|e| CompressionError::InvalidData(format!("lz4 hc compress: {e}")))
    })
}

/// Append one LZ4F block: its 4-byte little-endian size, then its bytes.
///
/// `compress` writes the compressed block into a scratch slice of
/// `max_compressed` bytes and returns its length. A block shorter than
/// [`MIN_COMPRESSIBLE_BLOCK`], or one that compression does not shrink, is
/// stored uncompressed, exactly as the LZ4 frame format allows.
fn write_block(
    out: &mut Vec<u8>,
    block: &[u8],
    max_compressed: usize,
    compress: impl FnOnce(&mut [u8]) -> Result<usize, CompressionError>,
) -> Result<(), CompressionError> {
    if block.len() < MIN_COMPRESSIBLE_BLOCK {
        let size = block_size(block.len())? | BLOCK_UNCOMPRESSED_BIT;
        out.extend_from_slice(&size.to_le_bytes());
        out.extend_from_slice(block);
        return Ok(());
    }
    let header = out.len();
    out.resize(header + 4 + max_compressed, 0);
    let (size_slot, scratch) = out
        .split_at_mut_checked(header)
        .and_then(|(_, tail)| tail.split_first_chunk_mut::<4>())
        .ok_or_else(|| scratch_error(block.len()))?;
    let n = compress(scratch)?;
    // `block` is at most `BLOCK_SIZE` (64 KiB), so both lengths always fit in
    // a `u32`.
    let (size, stored) = if n < block.len() {
        (block_size(n)?, n)
    } else {
        scratch
            .get_mut(..block.len())
            .ok_or_else(|| scratch_error(block.len()))?
            .copy_from_slice(block);
        (
            block_size(block.len())? | BLOCK_UNCOMPRESSED_BIT,
            block.len(),
        )
    };
    *size_slot = size.to_le_bytes();
    out.truncate(header + 4 + stored);
    Ok(())
}

/// A block length as the LZ4F 4-byte size field.
fn block_size(len: usize) -> Result<u32, CompressionError> {
    u32::try_from(len)
        .map_err(|_| CompressionError::InvalidData(format!("lz4 block of {len} bytes exceeds u32")))
}

/// The error for a scratch slice too short to hold a block of `len` bytes.
fn scratch_error(len: usize) -> CompressionError {
    CompressionError::InvalidData(format!(
        "lz4 scratch space too short for a {len}-byte block"
    ))
}

pub fn decompress(data: &[u8], max_output: usize) -> Result<Bytes, CompressionError> {
    if data.is_empty() {
        return Err(CompressionError::InvalidData("empty lz4 payload".into()));
    }
    let decoder = FrameDecoder::new(data);
    // Read at most `max_output + 1` bytes so we can detect overflow without
    // materializing the oversized output.
    let mut limited = decoder.take((max_output as u64).saturating_add(1));
    let mut out = Vec::with_capacity(data.len().saturating_mul(2).min(max_output));
    limited
        .read_to_end(&mut out)
        .map_err(|e| CompressionError::InvalidData(format!("lz4 decode: {e}")))?;
    if out.len() > max_output {
        return Err(CompressionError::TooLarge { limit: max_output });
    }
    Ok(Bytes::from(out))
}

#[cfg(test)]
mod tests {

    use super::*;

    const HELLO: &[u8] = b"hello kafka, this is a moderately repetitive payload to compress";
    const BIG_CAP: usize = 256 * 1024 * 1024;

    #[test]
    fn roundtrip() {
        let z = compress(HELLO).unwrap();
        let back = decompress(&z, BIG_CAP).unwrap();
        assert2::assert!(back.as_ref() == HELLO);
    }

    #[test]
    fn decompress_empty_rejected() {
        assert2::assert!(matches!(
            decompress(b"", BIG_CAP),
            Err(CompressionError::InvalidData(_))
        ));
    }

    #[test]
    fn decompress_garbage_rejected() {
        assert2::assert!(matches!(
            decompress(b"this is not lz4", BIG_CAP),
            Err(CompressionError::InvalidData(_))
        ));
    }

    #[test]
    fn larger_payload_roundtrips() {
        let big = vec![0xABu8; 128 * 1024]; // 128 KiB -> multiple 64 KiB blocks
        let z = compress(&big).unwrap();
        let back = decompress(&z, BIG_CAP).unwrap();
        assert2::assert!(back.as_ref() == big.as_slice());
    }

    #[test]
    fn decompression_bomb_rejected() {
        let bomb = vec![0u8; 64 * 1024 * 1024];
        let z = compress(&bomb).unwrap();
        assert2::assert!(matches!(
            decompress(&z, 1024),
            Err(CompressionError::TooLarge { limit: 1024 })
        ));
        let back = decompress(&z, BIG_CAP).unwrap();
        assert2::assert!(back.as_ref() == bomb.as_slice());
    }

    #[test]
    fn decompress_at_exact_cap_succeeds() {
        let z = compress(HELLO).unwrap();
        // Output of exactly `max_output` bytes is allowed (cap check is
        // `len > max_output`, not `>=`).
        let back = decompress(&z, HELLO.len()).unwrap();
        assert2::assert!(back.as_ref() == HELLO);
        // One byte under the exact size is rejected.
        assert2::assert!(matches!(
            decompress(&z, HELLO.len() - 1),
            Err(CompressionError::TooLarge { limit }) if limit == HELLO.len() - 1
        ));
    }

    #[test]
    fn frame_uses_64kib_independent_blocks() {
        // Compress a payload larger than 64 KiB so the block-size choice is
        // observable in the frame header: our explicit `Max64KB` must stay
        // 64 KiB rather than grow to an auto-selected larger block.
        let big = vec![0xCDu8; 128 * 1024];
        let z = compress(&big).unwrap();
        // LZ4 frame layout: [magic:4][FLG][BD]...
        let flg = z[4];
        let bd = z[5];
        // BD bits 4..6 encode the block max size; value 4 == 64 KiB.
        assert2::assert!(&z[..7] == &[0x04, 0x22, 0x4D, 0x18, 0x60, 0x40, 0x82]);
        assert2::assert!((bd >> 4) & 0x7 == 4);
        assert2::assert!((flg >> 5) & 1 == 1);
        assert2::assert!((flg >> 4) & 1 == 0);
        assert2::assert!((flg >> 2) & 1 == 0);
    }

    /// A payload compressible enough that HC's larger search window finds
    /// matches the fast compressor misses, over more than one 64 KiB block.
    fn compressible_payload() -> Vec<u8> {
        let unit = b"the quick brown fox jumps over the lazy dog, again and again; ";
        unit.iter().copied().cycle().take(3 * BLOCK_SIZE).collect()
    }

    /// Table-driven over Kafka's lz4 levels: the default level (9) gives the
    /// same bytes as the fast compressor, every other tested level uses HC and
    /// beats the fast compressor's size on a compressible payload, and every
    /// level decompresses back to the original payload.
    #[test]
    fn compress_with_level_uses_hc_except_at_the_default() {
        let payload = compressible_payload();
        let fast = compress(&payload).unwrap();

        for (level, expect_fast_bytes) in [(9, true), (1, false), (17, false)] {
            let out = compress_with_level(&payload, level).unwrap();

            assert2::assert!(
                (out.as_ref() == fast.as_ref()) == expect_fast_bytes,
                "level {level}"
            );
            if !expect_fast_bytes {
                assert2::assert!(out.len() < fast.len(), "level {level}");
            }

            let back = decompress(&out, BIG_CAP).unwrap();
            assert2::assert!(back.as_ref() == payload.as_slice(), "level {level}");
        }
    }

    /// Every level above `liblz4`'s `LZ4HC_CLEVEL_MAX` (12) compresses
    /// exactly as hard as 12: `write_hc_block` clamps explicitly, matching
    /// what `liblz4` documents doing internally for a caller that skips the
    /// clamp. Kafka's own range goes to 17, so a level in `13..=17` must
    /// still produce Kafka-decodable output -- it just cannot beat what 12
    /// already achieves.
    #[test]
    fn levels_above_liblz4_max_clamp_to_liblz4_max() {
        let payload = compressible_payload();
        let at_max = compress_with_level(&payload, LIBLZ4_HC_CLEVEL_MAX).unwrap();
        for level in [13, 17] {
            let out = compress_with_level(&payload, level).unwrap();
            assert2::assert!(out.as_ref() == at_max.as_ref(), "level {level}");
            let back = decompress(&out, BIG_CAP).unwrap();
            assert2::assert!(back.as_ref() == payload.as_slice(), "level {level}");
        }
    }

    /// `write_hc_block` stores a block uncompressed, with the LZ4F
    /// uncompressed-block high bit set, when HC does not shrink it -- the
    /// fallback the frame format allows for incompressible input.
    #[test]
    fn incompressible_block_falls_back_to_stored() {
        // High-entropy input HC cannot shrink: every byte value once, so
        // there is no repeated byte or short pattern to exploit within a
        // single block.
        let incompressible: Vec<u8> = (0..=255u8).collect();
        let out = compress_with_level(&incompressible, 1).unwrap();
        let block_size = u32::from_le_bytes(out[FRAME_HEADER.len()..][..4].try_into().unwrap());
        assert2::assert!(block_size & BLOCK_UNCOMPRESSED_BIT != 0);
        assert2::assert!((block_size & !BLOCK_UNCOMPRESSED_BIT) as usize == incompressible.len());
        let back = decompress(&out, BIG_CAP).unwrap();
        assert2::assert!(back.as_ref() == incompressible.as_slice());
    }

    #[test]
    fn compress_with_level_hc_frame_uses_the_same_header_as_the_fast_frame() {
        // The HC path builds its own frame by hand; its header must still be
        // byte-identical to `lz4rip`'s, since that is what makes the frame
        // settings (independent 64 KiB blocks, no checksums, no content size)
        // match Kafka's regardless of which compressor filled the blocks.
        let payload = compressible_payload();
        let fast = compress(&payload).unwrap();
        let hc = compress_with_level(&payload, 1).unwrap();
        assert2::assert!(fast[..FRAME_HEADER.len()] == hc[..FRAME_HEADER.len()]);
    }
}