xpress_rs 0.2.0

//! lz77 module
//!
//! This module contains the LZ77Compressor and LZ77Decompressor structs.
//! Each struct has a constructor and a method to compress/decompress data.

use std::collections::HashMap;
use std::error::Error;
use std::hash::BuildHasher;

use xxhash_rust::xxh3::Xxh3;

// logger
#[cfg(feature = "logging")]
use log::{debug, error, info, trace, warn};

/// LZ77Compressor struct
pub struct LZ77Compressor {
    flags: u32,
    flag_count: u32,
    flag_output_position: usize,
    output_position: usize,
    input_position: usize,
    last_length_half_byte: usize,
    input_buffer: Vec<u8>,
    output_buffer: Vec<u8>,
    sequence_index_map: HashMap<u32, usize, Xxh3BuildHasher>,
    match_length: usize,
    match_offset: usize,
}

struct Xxh3BuildHasher;

impl BuildHasher for Xxh3BuildHasher {
    type Hasher = Xxh3;

    fn build_hasher(&self) -> Xxh3 {
        Xxh3::new()
    }
}

impl LZ77Compressor {
    /// Constructor for LZ77Compressor
    ///
    /// # Arguments
    ///
    /// * `input_buffer` - The input buffer to compress as a Vec<u8>
    ///
    pub fn new(input_buffer: Vec<u8>) -> LZ77Compressor {
        let size = input_buffer.len() + 20;
        LZ77Compressor {
            flags: 0,
            flag_count: 0,
            flag_output_position: 0,
            output_position: 4,
            input_position: 0,
            last_length_half_byte: 0,
            input_buffer,
            output_buffer: vec![0; size], // TODO : may introduce some buffer overflow " user can chose to use a verification method to check if the output buffer is big enough, but can be slow"
            sequence_index_map: HashMap::with_hasher(Xxh3BuildHasher),
            match_length: 0,
            match_offset: 0,
        }
    }

    fn find_match(&mut self) -> bool {
        // Ensure that at least 3 characters are left
        if self.input_position + 2 >= self.input_buffer.len() {
            return false;
        }

        // the largest match offset that can be encoded is 8192
        let start_window = self.input_position.saturating_sub(8192);

        let sequence_u32 = self
            .input_buffer
            .get(self.input_position..self.input_position + 3)
            .map(|three_bytes| {
                let mut result = [0; 4];
                result[1..4].copy_from_slice(three_bytes);
                u32::from_ne_bytes(result)
            });

        // Check if the sequence is in the hashmap and in the correct window
        if let Some(&index) =
            sequence_u32.and_then(|sequence| self.sequence_index_map.get(&sequence))
        {
            // Check if the match is in the correct window
            if index >= start_window && index < self.input_position {
                // Calculate the length of the match.
                let length = self.input_buffer[index..]
                    .iter()
                    .zip(&self.input_buffer[self.input_position..])
                    .take_while(|(a, b)| a == b)
                    .count();
                self.match_length = length;

                // Calculate the offset of the match
                let match_offset = self.input_position - index;
                self.match_offset = match_offset;

                #[cfg(feature = "logging")]
                trace!(
                    "Match found at index: {}, length: {}, offset: {}",
                    index,
                    length,
                    match_offset
                );

                // Add the current sequence to the hashmap, update the position.
                self.sequence_index_map
                    .insert(sequence_u32.unwrap(), self.input_position);

                // match found, return true.
                return true;
            }
        }

        // Add the current sequence to the hashmap.
        self.sequence_index_map
            .insert(sequence_u32.unwrap(), self.input_position);
        // No match found, return false.
        false
    }

    fn encode_extra_length(&mut self) {
        // We've already used 3 bits + 4 bits to encode the length,
        self.match_length -= 15;

        // Check if we can encode the remaining length in a single byte.
        if self.match_length < 255 {
            self.output_buffer[self.output_position] = self.match_length as u8;
            self.output_position += 1;
        } else {
            // Use two more bytes for the length.
            self.output_buffer[self.output_position] = 255;
            self.output_position += 1;
            self.match_length += 7 + 15;

            // if 2 bytes are enough to encode the length, use 2 bytes.
            if self.match_length < (1 << 16) {
                self.output_buffer[self.output_position..self.output_position + 2]
                    .copy_from_slice(&(self.match_length as u16).to_le_bytes());
                self.output_position += 2;
            } else {
                // Use 4 bytes for the length.
                self.output_buffer[self.output_position..self.output_position + 2]
                    .copy_from_slice(&(0 as u16).to_le_bytes());
                self.output_position += 2;
                self.output_buffer[self.output_position..self.output_position + 4]
                    .copy_from_slice(&(self.match_length as u32).to_le_bytes());
                self.output_position += 4;
            }
        }
        #[cfg(feature = "logging")]
        trace!(
            "Encode extra length called, final length: {}",
            self.match_length
        );
    }

    fn write_flags(&mut self) {
        // Encode the flags as little endian u32
        let flags_bytes = (self.flags as u32).to_le_bytes();

        // Write the  flags to the output buffer
        self.output_buffer[self.flag_output_position..self.flag_output_position + 4]
            .copy_from_slice(&flags_bytes);
        self.flag_count = 0;
        self.flag_output_position = self.output_position;
        #[cfg(feature = "logging")]
        trace!(
            "Writing flags: {:?} at position: {}",
            flags_bytes,
            self.flag_output_position
        );
    }

    /// Compress the input buffer and return a result containing the compressed data.
    ///
    /// # Return value
    ///
    /// * `Ok(Vec<u8>)` - The compressed data.
    /// * `Err(Box<dyn Error>)` - An error occured during compression.
    ///
    pub fn compress(mut self) -> Result<Vec<u8>, Box<dyn Error>> {
        #[cfg(feature = "logging")]
        info!(
            "Starting compression, input size: {}",
            self.input_buffer.len()
        );

        while self.input_position < self.input_buffer.len() {
            // while there is still data to compress

            let result = self.find_match(); // try to find a match

            if result == false || self.input_position + 2 > self.input_buffer.len() {
                // if no match was found or there is less than 3 characters left

                // Add a literal to the output buffer
                self.output_buffer[self.output_position] = self.input_buffer[self.input_position];
                self.input_position += 1;
                self.output_position += 1;
                #[cfg(feature = "logging")]
                trace!(
                    "Writing literal: {}",
                    self.input_buffer[self.input_position - 1]
                );

                // update the current flags with a 0 bit, corresponding to a literal.
                self.flags <<= 1;
                self.flag_count += 1;

                // If 32 bit have been written, write the flags to the output buffer
                if self.flag_count == 32 {
                    self.write_flags();
                    self.output_position += 4;
                }
            } else if result {
                // if a match was found

                let end_of_match = self.match_length; // calculate the end of the match

                self.match_length -= 3; // we've already matched 3 characters
                self.match_offset -= 1; // the offset is encoded with 1 less than the actual offset
                                        // shift the offset 3 bits to the left ( matches are encoded with a 16-bit value, where the high 13 bits represent the offset )
                self.match_offset <<= 3;

                // Encode the length of the match :
                if self.match_length < 7 {
                    // If the length is less than 7, we can encode it in 3 bits in the current 2 bytes ( 13 bits for the offset, 3 bits for the length )
                    self.match_offset += self.match_length;
                    self.output_buffer[self.output_position..self.output_position + 2]
                        .copy_from_slice(&(self.match_offset as u16).to_le_bytes());
                    self.output_position += 2;
                } else {
                    // 3 bit is not enough to encode the length, record a special value, for indicate a longer length.
                    self.match_offset |= 7;
                    // write the special value to the output buffer in the next 2 bytes.
                    self.output_buffer[self.output_position..self.output_position + 2]
                        .copy_from_slice(&(self.match_offset as u16).to_le_bytes());
                    self.output_position += 2;

                    // Try to encode the length in the next 4 bits. If we previously encoded a 4-bit length, we'll use the high 4 bits from that byte.

                    self.match_length -= 7; // don't count the 3 bits we've already encoded

                    // if we encoded a 4-bit length in the last match byte
                    if self.last_length_half_byte == 0 {
                        // Indicate the position of the match, where the 4 high bits can be used to encode the length of the next match.
                        self.last_length_half_byte = self.output_position;

                        // If the length is less than 15, we can encode it in the 4 bits.
                        if self.match_length < 15 {
                            self.output_buffer[self.output_position] = self.match_length as u8;
                            self.output_position += 1;
                        }
                        // If the length is greater than 15, write 15 to the 4 bits, and call encode_extra_length() to encode the rest of the length.
                        else {
                            self.output_buffer[self.output_position] = 15;
                            self.output_position += 1;
                            self.encode_extra_length();
                        }
                    }
                    // If we didn't encode a 4-bit length in the last match byte
                    else {
                        // If the length is less than 15, we can encode it in the 4 bits of the previous match byte.
                        if self.match_length < 15 {
                            let encoded_length = (self.match_length as u8) << 4;
                            self.output_buffer[self.last_length_half_byte] |= encoded_length;
                            self.last_length_half_byte = 0; // Indicate that we've used the 4 bits in the previous match byte.
                        }
                        // If the length is greater than 15, write 15 to the 4 bits, and call encode_extra_length() to encode the rest of the length.
                        else {
                            self.output_buffer[self.last_length_half_byte] |= 15 << 4;
                            self.last_length_half_byte = 0; // Indicate that we've used the 4 bits in the previous match byte.
                            self.encode_extra_length();
                        }
                    }
                }

                #[cfg(feature = "logging")]
                trace!("Writing match");

                // update the current flags with a 1 bit, corresponding to a match.
                self.flags = (self.flags << 1) | 1;
                self.flag_count += 1;
                // If 32 bit have been written, write the flags to the output buffer
                if self.flag_count == 32 {
                    self.write_flags();
                    self.output_position += 4;
                }
                //Advance InputPosition to the first byte that was not in the match
                self.input_position += end_of_match;
            }
        }
        // End of input buffer reached, write the remaining flags to the output buffer. The remaining flags are encoded as 1's.
        self.flags <<= 32 - self.flag_count;
        self.flags |= (1 << (32 - self.flag_count)) - 1;
        self.write_flags();

        // Remove the unused part of the output buffer.
        self.output_buffer.truncate(self.output_position);
        self.output_buffer.shrink_to_fit();

        #[cfg(feature = "logging")]
        info!(
            "Compression finished, output buffer size : {}",
            self.output_buffer.len()
        );

        // Return the compressed data.
        Ok(self.output_buffer)
    }
}

/// LZ77Decompressor struct
pub struct LZ77Decompressor {
    buffered_flags: u32,
    buffered_flag_count: u8,
    output_position: usize,
    input_position: usize,
    last_length_half_byte: usize,
    input_buffer: Vec<u8>,
    output_buffer: Vec<u8>,
}

impl LZ77Decompressor {
    /// constructor for LZ77Decompressor
    ///
    /// # Arguments
    ///
    /// * `input_buffer` - The buffer containing the compressed data.
    ///
    pub fn new(input_buffer: Vec<u8>) -> LZ77Decompressor {
        LZ77Decompressor {
            buffered_flags: 0,
            buffered_flag_count: 0,
            output_position: 0,
            input_position: 0,
            last_length_half_byte: 0,
            input_buffer,
            output_buffer: Vec::new(),
        }
    }

    /// Decompress the input buffer. and return a result containing the decompressed data.
    ///
    /// # Return value
    ///
    /// * `Ok(Vec<u8>)` - The decompressed data.
    /// * `Err(Box<dyn Error>)` - An error occurred during decompression.
    ///
    pub fn decompress(mut self) -> Result<Vec<u8>, Box<dyn Error>> {
        #[cfg(feature = "logging")]
        info!(
            "Decompressing started, input buffer size: {}",
            self.input_buffer.len()
        );

        loop {
            // If we-ve read all the bits in the current flags, read the next 32 bits ( corresponding to the next flags )
            if self.buffered_flag_count == 0 {
                self.buffered_flags = u32::from_le_bytes(
                    self.input_buffer[self.input_position..self.input_position + 4]
                        .try_into()
                        .map_err(|e| {
                            #[cfg(feature = "logging")]
                            error!("Error during converting slice to array: {}", e);
                            format!("Error during converting slice to array: {}", e)
                        })?,
                );

                self.input_position += 4;
                self.buffered_flag_count = 32;
            }
            self.buffered_flag_count -= 1;
            #[cfg(feature = "logging")]
            trace!("Current flag: {}", self.buffered_flags);

            // If the current flag is 0, corresponding to a literal byte, copy the next byte from the input buffer to the output buffer.
            if (self.buffered_flags & (1 << self.buffered_flag_count)) == 0 {
                self.output_buffer
                    .push(self.input_buffer[self.input_position]);
                self.input_position += 1;
                self.output_position += 1;
                #[cfg(feature = "logging")]
                trace!(
                    "Literal byte: {} at position {}",
                    self.input_buffer[self.input_position - 1],
                    self.input_position - 1
                );
            }
            // If the current flag is 1, corresponding to a match or the end of the input buffer
            else {
                // If the end of the input buffer has been reached, end the decompression and return the output buffer.
                if self.input_position == self.input_buffer.len() {
                    #[cfg(feature = "logging")]
                    info!(
                        "Decompression finished, output buffer size: {}",
                        self.output_buffer.len()
                    );
                    return Ok(self.output_buffer);
                }

                // Read the next 16 bits from the input buffer, corresponding to the match offset and length.
                let match_bytes = u16::from_le_bytes(
                    self.input_buffer[self.input_position..self.input_position + 2]
                        .try_into()
                        .map_err(|e| {
                            #[cfg(feature = "logging")]
                            error!("Error during converting slice to array: {}", e);
                            format!("Error during converting slice to array: {}", e)
                        })?,
                );
                self.input_position += 2;

                // Decode the match offset and length.
                let mut match_length: usize = (match_bytes % 8) as usize;
                let match_offset: usize = ((match_bytes / 8) + 1) as usize;

                // If the 3-bit length field from the match bytes was set to 7,
                // it means the length of the match is encoded in the next byte(s) in the input buffer
                if match_length == 7 {
                    // If this is the first part of the match length, read the lower 4 bits of the next byte
                    // Store the current position to LastLengthHalfByte for later use
                    if self.last_length_half_byte == 0 {
                        match_length = self.input_buffer[self.input_position] as usize;
                        match_length = match_length % 16;
                        self.last_length_half_byte = self.input_position;
                        self.input_position += 1;
                    }
                    // If this is the second part of the match length, read the higher 4 bits of the byte at LastLengthHalfByte
                    // Reset LastLengthHalfByte to 0 for future reads
                    else {
                        match_length = self.input_buffer[self.last_length_half_byte] as usize;
                        match_length = match_length / 16;
                        self.last_length_half_byte = 0;
                    }

                    // If the 4-bit length field from the match bytes was set to 15,
                    if match_length == 15 {
                        // Read the next byte
                        match_length =
                            self.input_buffer[self.input_position]
                                .try_into()
                                .map_err(|e| {
                                    #[cfg(feature = "logging")]
                                    error!("Error during converting slice to array: {}", e);
                                    format!("Error during converting slice to array: {}", e)
                                })?;
                        self.input_position += 1;
                        // if the decoded length is 255, the length is stored in the next 2 bytes
                        if match_length == 255 {
                            match_length = u16::from_le_bytes(
                                self.input_buffer[self.input_position..self.input_position + 2]
                                    .try_into()
                                    .map_err(|e| {
                                        #[cfg(feature = "logging")]
                                        error!("Error during converting slice to array: {}", e);
                                        format!("Error during converting slice to array: {}", e)
                                    })?,
                            ) as usize;
                            self.input_position += 2;

                            // if the decoded length is 0, the length is stored in the next 4 bytes
                            if match_length == 0 {
                                match_length = u32::from_le_bytes(
                                    self.input_buffer[self.input_position..self.input_position + 4]
                                        .try_into()
                                        .map_err(|e| {
                                            #[cfg(feature = "logging")]
                                            error!("Error during converting slice to array: {}", e);
                                            format!("Error during converting slice to array: {}", e)
                                        })?,
                                ) as usize;
                                self.input_position += 4;
                            }
                            if match_length < 15 + 7 {
                                #[cfg(feature = "logging")]
                                error!("Invalid match length");
                                return Err(Box::new(std::io::Error::new(
                                    std::io::ErrorKind::InvalidData,
                                    "Invalid match length",
                                )));
                            }
                            match_length -= 15 + 7;
                        }
                        match_length += 15;
                    }
                    match_length += 7;
                }
                // Add 3 to the match length to account for the minimum length of 3
                match_length += 3;

                #[cfg(feature = "logging")]
                trace!("Match : offset: {}, length: {}", match_offset, match_length);

                // Copy bytes from the earlier part of the output buffer to the current position
                // The amount of bytes copied is equal to the match length
                for _i in 0..match_length {
                    self.output_buffer
                        .push(self.output_buffer[self.output_position - match_offset]);
                    self.output_position += 1;
                }
            }
        }
    }
}