lz4/

encoder.rs

//! `Write`-based streaming encoder wrapping the LZ4 frame API
//! (`LZ4F_compressBegin` / `LZ4F_compressUpdate` / `LZ4F_compressEnd`).
//!
//! [`Encoder`] consumes an inner [`Write`] sink and emits a complete `.lz4`
//! frame: a frame header (written eagerly on `build`), a sequence of
//! compressed blocks (flushed as the input fills the block-size buffer or on
//! explicit [`flush`](Encoder::flush)), and a frame footer (written by
//! [`finish`](Encoder::finish)). [`EncoderBuilder`] configures the frame
//! preferences (block size, block mode, checksums, compression level, etc.)
//! before constructing the encoder.
//!
//! The API mirrors the `lz4-rs` crate's [`Encoder`] / [`EncoderBuilder`] so
//! that this crate is a drop-in replacement.

use super::liblz4::*;
use super::size_t;
use std::cmp;
use std::io::Result;
use std::io::Write;
use std::ptr;

/// RAII wrapper around an `LZ4F_cctx*` that frees the context on drop.
#[derive(Debug)]
struct EncoderContext {
    c: LZ4FCompressionContext,
}

/// Builder for [`Encoder`] that collects LZ4 frame preferences before
/// constructing the streaming encoder.
///
/// Preferences correspond to the fields of upstream `LZ4F_preferences_t`
/// (see `lz4frame.h`). Defaults match the `lz4-rs` crate, not the upstream
/// `lz4` CLI; for CLI-compatible defaults see `src/bin/lz4.rs`.
#[derive(Clone, Debug)]
pub struct EncoderBuilder {
    block_size: BlockSize,
    block_mode: BlockMode,
    // 1: each block followed by a checksum of block's compressed data; 0: disabled (default)
    block_checksum: BlockChecksum,
    checksum: ContentChecksum,
    // 0 == default (fast mode); values above 16 count as 16; values below 0 count as 0
    level: u32,
    // 1 == always flush (reduce need for tmp buffer)
    auto_flush: bool,
    favor_dec_speed: bool,
    content_size: u64,
}

/// Streaming LZ4 frame encoder wrapping an inner [`Write`] sink.
///
/// Bytes written to the encoder are buffered up to the configured block
/// size and then handed to `LZ4F_compressUpdate`, which emits compressed
/// blocks into `w`. Always call [`finish`](Encoder::finish) to flush the
/// final block and write the frame footer — dropping the encoder without
/// calling `finish` produces a truncated frame.
#[derive(Debug)]
pub struct Encoder<W> {
    c: EncoderContext,
    w: W,
    limit: usize,
    auto_flush: bool,
    input: Vec<u8>,
    buffer: Vec<u8>,
}

impl EncoderBuilder {
    /// Returns a builder pre-populated with the same defaults used by the
    /// `lz4-rs` crate (linked blocks, content + block checksums enabled,
    /// fast-mode compression).
    pub fn new() -> Self {
        EncoderBuilder {
            block_size: BlockSize::Default,
            block_mode: BlockMode::Linked,
            checksum: ContentChecksum::ChecksumEnabled,
            block_checksum: BlockChecksum::BlockChecksumEnabled,
            level: 0,
            auto_flush: false,
            favor_dec_speed: false,
            content_size: 0,
        }
    }

    /// Selects the maximum block size (64 KiB, 256 KiB, 1 MiB, or 4 MiB).
    /// Larger blocks generally improve compression ratio at the cost of
    /// memory.
    pub fn block_size(&mut self, block_size: BlockSize) -> &mut Self {
        self.block_size = block_size;
        self
    }

    /// Sets whether blocks may reference data from previous blocks
    /// ([`BlockMode::Linked`]) or must each compress independently
    /// ([`BlockMode::Independent`]).
    pub fn block_mode(&mut self, block_mode: BlockMode) -> &mut Self {
        self.block_mode = block_mode;
        self
    }

    /// Enables or disables the per-block xxHash32 trailer that follows each
    /// compressed block.
    pub fn block_checksum(&mut self, block_checksum: BlockChecksum) -> &mut Self {
        self.block_checksum = block_checksum;
        self
    }

    /// Enables or disables the xxHash32 checksum of the full decompressed
    /// content that is appended to the frame footer.
    pub fn checksum(&mut self, checksum: ContentChecksum) -> &mut Self {
        self.checksum = checksum;
        self
    }

    /// Sets the compression level. `0` selects fast mode; levels `3..=12`
    /// route through HC. Levels above 12 are clamped by the underlying
    /// implementation.
    pub fn level(&mut self, level: u32) -> &mut Self {
        self.level = level;
        self
    }

    /// When `true`, every [`Encoder::write`] call is flushed to the inner
    /// writer instead of being buffered, reducing internal memory usage at
    /// some cost in compression ratio for small writes.
    pub fn auto_flush(&mut self, auto_flush: bool) -> &mut Self {
        self.auto_flush = auto_flush;
        self
    }

    /// Favor decompression speed over compression ratio. Requires compression
    /// level >=10.
    pub fn favor_dec_speed(&mut self, favor_dec_speed: bool) -> &mut Self {
        self.favor_dec_speed = favor_dec_speed;
        self
    }

    /// Records the total uncompressed content size in the frame header.
    /// Use `0` (the default) when the size is unknown ahead of time.
    pub fn content_size(&mut self, content_size: u64) -> &mut Self {
        self.content_size = content_size;
        self
    }

    /// Consumes the builder, allocates an LZ4 frame compression context,
    /// writes the frame header into `w`, and returns the resulting
    /// streaming encoder.
    pub fn build<W: Write>(&self, w: W) -> Result<Encoder<W>> {
        let block_size = self.block_size.get_size();
        let preferences = LZ4FPreferences {
            frame_info: LZ4FFrameInfo {
                block_size_id: self.block_size.clone(),
                block_mode: self.block_mode.clone(),
                content_checksum_flag: self.checksum.clone(),
                content_size: self.content_size,
                frame_type: FrameType::Frame,
                dict_id: 0,
                block_checksum_flag: self.block_checksum.clone(),
            },
            compression_level: self.level,
            auto_flush: if self.auto_flush { 1 } else { 0 },
            favor_dec_speed: if self.favor_dec_speed { 1 } else { 0 },
            reserved: [0; 3],
        };
        let mut encoder = Encoder {
            w,
            c: EncoderContext::new()?,
            limit: block_size,
            auto_flush: self.auto_flush,
            input: Vec::with_capacity(block_size),
            buffer: Vec::with_capacity(check_error(unsafe {
                LZ4F_compressBound(block_size as size_t, &preferences)
            })?),
        };
        encoder.write_header(&preferences)?;
        Ok(encoder)
    }
}

impl Default for EncoderBuilder {
    fn default() -> Self {
        Self::new()
    }
}

impl<W: Write> Encoder<W> {
    fn write_header(&mut self, preferences: &LZ4FPreferences) -> Result<()> {
        unsafe {
            let len = check_error(LZ4F_compressBegin(
                self.c.c,
                self.buffer.as_mut_ptr(),
                self.buffer.capacity() as size_t,
                preferences,
            ))?;
            self.buffer.set_len(len);
        }
        self.w.write_all(&self.buffer)
    }

    fn write_end(&mut self) -> Result<()> {
        self.write_pending()?;
        unsafe {
            let len = check_error(LZ4F_compressEnd(
                self.c.c,
                self.buffer.as_mut_ptr(),
                self.buffer.capacity() as size_t,
                ptr::null(),
            ))?;
            self.buffer.set_len(len);
        };
        self.w.write_all(&self.buffer)
    }

    /// Immutable writer reference.
    pub fn writer(&self) -> &W {
        &self.w
    }

    fn write_update(&mut self, input: &[u8]) -> Result<()> {
        unsafe {
            let len = check_error(LZ4F_compressUpdate(
                self.c.c,
                self.buffer.as_mut_ptr(),
                self.buffer.capacity() as size_t,
                input.as_ptr(),
                input.len() as size_t,
                ptr::null(),
            ))?;
            self.buffer.set_len(len);
        }
        self.w.write_all(&self.buffer)
    }

    fn write_pending(&mut self) -> Result<()> {
        if self.input.is_empty() {
            return Ok(());
        }
        unsafe {
            let len = check_error(LZ4F_compressUpdate(
                self.c.c,
                self.buffer.as_mut_ptr(),
                self.buffer.capacity() as size_t,
                self.input.as_ptr(),
                self.input.len() as size_t,
                ptr::null(),
            ))?;
            self.buffer.set_len(len);
        }
        self.input.clear();
        self.w.write_all(&self.buffer)
    }

    /// Finalises the LZ4 frame.
    ///
    /// Flushes any buffered input, invokes `LZ4F_compressEnd` to write the
    /// frame footer (and content checksum if enabled), and returns the
    /// wrapped writer together with the result of writing those final
    /// bytes. Always call this; dropping the encoder without `finish`
    /// produces a truncated, undecodable frame.
    pub fn finish(mut self) -> (W, Result<()>) {
        let result = self.write_end();
        (self.w, result)
    }
}

impl<W: Write> Write for Encoder<W> {
    fn write(&mut self, buffer: &[u8]) -> Result<usize> {
        if self.auto_flush {
            self.write_update(buffer)?;
            return Ok(buffer.len());
        }

        let mut offset = 0;
        while offset < buffer.len() {
            let remaining = buffer.len() - offset;
            if self.input.is_empty() && remaining >= self.limit {
                let end = offset + self.limit;
                self.write_update(&buffer[offset..end])?;
                offset = end;
                continue;
            }

            let available = self.limit - self.input.len();
            let size = cmp::min(remaining, available);
            self.input.extend_from_slice(&buffer[offset..offset + size]);
            offset += size;
            if self.input.len() == self.limit {
                self.write_pending()?;
            }
        }
        Ok(buffer.len())
    }

    fn flush(&mut self) -> Result<()> {
        self.write_pending()?;
        loop {
            unsafe {
                let len = check_error(LZ4F_flush(
                    self.c.c,
                    self.buffer.as_mut_ptr(),
                    self.buffer.capacity() as size_t,
                    ptr::null(),
                ))?;
                if len == 0 {
                    break;
                }
                self.buffer.set_len(len);
            };
            self.w.write_all(&self.buffer)?;
        }
        self.w.flush()
    }
}

impl EncoderContext {
    fn new() -> Result<EncoderContext> {
        let mut context = LZ4FCompressionContext(ptr::null_mut());
        check_error(unsafe { LZ4F_createCompressionContext(&mut context, LZ4F_VERSION) })?;
        Ok(EncoderContext { c: context })
    }
}

impl Drop for EncoderContext {
    fn drop(&mut self) {
        unsafe { LZ4F_freeCompressionContext(self.c) };
    }
}

#[cfg(test)]
mod test {
    use super::EncoderBuilder;
    use crate::liblz4::{BlockChecksum, ContentChecksum};
    use std::io::{Read, Write};

    #[test]
    fn test_encoder_smoke() {
        let mut encoder = EncoderBuilder::new().level(1).build(Vec::new()).unwrap();
        encoder.write_all(b"Some ").unwrap();
        encoder.write_all(b"data").unwrap();
        let (_, result) = encoder.finish();
        result.unwrap();
    }

    #[test]
    fn test_encoder_random() {
        let mut encoder = EncoderBuilder::new().level(1).build(Vec::new()).unwrap();
        let mut input = Vec::new();
        let mut rnd: u32 = 42;
        for _ in 0..1024 * 1024 {
            input.push((rnd & 0xFF) as u8);
            rnd = (1664525_u64 * (rnd as u64) + 1013904223_u64) as u32;
        }
        encoder.write_all(&input).unwrap();
        let (compressed, result) = encoder.finish();
        result.unwrap();

        let mut dec = crate::decoder::Decoder::new(&compressed[..]).unwrap();
        let mut output = Vec::new();
        dec.read_to_end(&mut output).unwrap();
        assert_eq!(input, output);
    }

    #[test]
    fn test_encoder_content_size() {
        let mut encoder = EncoderBuilder::new()
            .level(1)
            .content_size(1024 * 1024)
            .build(Vec::new())
            .unwrap();
        let mut input = Vec::new();
        let mut rnd: u32 = 42;
        for _ in 0..1024 * 1024 {
            input.push((rnd & 0xFF) as u8);
            rnd = (1664525_u64 * (rnd as u64) + 1013904223_u64) as u32;
        }
        encoder.write_all(&input).unwrap();
        let (compressed, result) = encoder.finish();
        result.unwrap();

        let mut dec = crate::decoder::Decoder::new(&compressed[..]).unwrap();
        let mut output = Vec::new();
        dec.read_to_end(&mut output).unwrap();
        assert_eq!(input, output);
    }

    #[test]
    fn test_encoder_auto_flush_emits_partial_write() {
        let mut encoder = EncoderBuilder::new()
            .auto_flush(true)
            .block_checksum(BlockChecksum::NoBlockChecksum)
            .checksum(ContentChecksum::NoChecksum)
            .build(Vec::new())
            .unwrap();
        let header_len = encoder.writer().len();
        encoder.write_all(b"partial").unwrap();
        assert!(encoder.writer().len() > header_len);

        let (compressed, result) = encoder.finish();
        result.unwrap();
        let mut dec = crate::decoder::Decoder::new(&compressed[..]).unwrap();
        let mut output = Vec::new();
        dec.read_to_end(&mut output).unwrap();
        assert_eq!(output, b"partial");
    }

    #[test]
    fn test_encoder_send() {
        fn check_send<S: Send>(_: &S) {}
        let enc = EncoderBuilder::new().build(Vec::new());
        check_send(&enc);
    }

    #[test]
    fn test_favor_dec_speed() {
        let mut encoder = EncoderBuilder::new()
            .level(11)
            .favor_dec_speed(true)
            .build(Vec::new())
            .unwrap();
        let mut input = Vec::new();
        let mut rnd: u32 = 42;
        for _ in 0..1024 * 1024 {
            input.push((rnd & 0xFF) as u8);
            rnd = (1664525_u64 * (rnd as u64) + 1013904223_u64) as u32;
        }
        encoder.write_all(&input).unwrap();
        let (compressed, result) = encoder.finish();
        result.unwrap();

        let mut dec = crate::decoder::Decoder::new(&compressed[..]).unwrap();

        let mut output = Vec::new();
        dec.read_to_end(&mut output).unwrap();
        assert_eq!(input, output);
    }
}