snap/

compress.rs

use std::fmt;
use std::ops::{Deref, DerefMut};
use std::ptr;

use crate::bytes;
use crate::error::{Error, Result};
use crate::{MAX_BLOCK_SIZE, MAX_INPUT_SIZE};

/// The total number of slots we permit for our hash table of 4 byte repeat
/// sequences.
const MAX_TABLE_SIZE: usize = 1 << 14;

/// The size of a small hash table. This is useful for reducing overhead when
/// compressing very small blocks of bytes.
const SMALL_TABLE_SIZE: usize = 1 << 10;

/// The total number of bytes that we always leave uncompressed at the end
/// of the buffer. This in particular affords us some wiggle room during
/// compression such that faster copy operations can be used.
const INPUT_MARGIN: usize = 16 - 1;

/// The minimum block size that we're willing to consider for compression.
/// Anything smaller than this gets emitted as a literal.
const MIN_NON_LITERAL_BLOCK_SIZE: usize = 1 + 1 + INPUT_MARGIN;

/// Nice names for the various Snappy tags.
enum Tag {
    Literal = 0b00,
    Copy1 = 0b01,
    Copy2 = 0b10,
    // Compression never actually emits a Copy4 operation and decompression
    // uses tricks so that we never explicitly do case analysis on the copy
    // operation type, therefore leading to the fact that we never use Copy4.
    #[allow(dead_code)]
    Copy4 = 0b11,
}

/// Returns the maximum compressed size given the uncompressed size.
///
/// If the uncompressed size exceeds the maximum allowable size then this
/// returns 0.
pub fn max_compress_len(input_len: usize) -> usize {
    let input_len = input_len as u64;
    if input_len > MAX_INPUT_SIZE {
        return 0;
    }
    let max = 32 + input_len + (input_len / 6);
    if max > MAX_INPUT_SIZE {
        0
    } else {
        max as usize
    }
}

/// Encoder is a raw encoder for compressing bytes in the Snappy format.
///
/// Thie encoder does not use the Snappy frame format and simply compresses the
/// given bytes in one big Snappy block (that is, it has a single header).
///
/// Unless you explicitly need the low-level control, you should use
/// [`read::FrameEncoder`](../read/struct.FrameEncoder.html)
/// or
/// [`write::FrameEncoder`](../write/struct.FrameEncoder.html)
/// instead, which compresses to the Snappy frame format.
///
/// It is beneficial to reuse an Encoder when possible.
pub struct Encoder {
    small: [u16; SMALL_TABLE_SIZE],
    big: Vec<u16>,
}

impl fmt::Debug for Encoder {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Encoder(...)")
    }
}

impl Encoder {
    /// Return a new encoder that can be used for compressing bytes.
    pub fn new() -> Encoder {
        Encoder { small: [0; SMALL_TABLE_SIZE], big: vec![] }
    }

    /// Compresses all bytes in `input` into `output`.
    ///
    /// `input` can be any arbitrary sequence of bytes.
    ///
    /// `output` must be large enough to hold the maximum possible compressed
    /// size of `input`, which can be computed using `max_compress_len`.
    ///
    /// On success, this returns the number of bytes written to `output`.
    ///
    /// # Errors
    ///
    /// This method returns an error in the following circumstances:
    ///
    /// * The total number of bytes to compress exceeds `2^32 - 1`.
    /// * `output` has length less than `max_compress_len(input.len())`.
    pub fn compress(
        &mut self,
        mut input: &[u8],
        output: &mut [u8],
    ) -> Result<usize> {
        match max_compress_len(input.len()) {
            0 => {
                return Err(Error::TooBig {
                    given: input.len() as u64,
                    max: MAX_INPUT_SIZE,
                });
            }
            min if output.len() < min => {
                return Err(Error::BufferTooSmall {
                    given: output.len() as u64,
                    min: min as u64,
                });
            }
            _ => {}
        }
        // Handle an edge case specially.
        if input.is_empty() {
            // Encodes a varint of 0, denoting the total size of uncompressed
            // bytes.
            output[0] = 0;
            return Ok(1);
        }
        // Write the Snappy header, which is just the total number of
        // uncompressed bytes.
        let mut d = bytes::write_varu64(output, input.len() as u64);
        while !input.is_empty() {
            // Find the next block.
            let mut src = input;
            if src.len() > MAX_BLOCK_SIZE {
                src = &src[..MAX_BLOCK_SIZE as usize];
            }
            input = &input[src.len()..];

            // If the block is smallish, then don't waste time on it and just
            // emit a literal.
            let mut block = Block::new(src, output, d);
            if block.src.len() < MIN_NON_LITERAL_BLOCK_SIZE {
                let lit_end = block.src.len();
                unsafe {
                    // SAFETY: next_emit is zero (in bounds) and the end is
                    // the length of the block (in bounds).
                    block.emit_literal(lit_end);
                }
            } else {
                let table = self.block_table(block.src.len());
                block.compress(table);
            }
            d = block.d;
        }
        Ok(d)
    }

    /// Compresses all bytes in `input` into a freshly allocated `Vec`.
    ///
    /// This is just like the `compress` method, except it allocates a `Vec`
    /// with the right size for you. (This is intended to be a convenience
    /// method.)
    ///
    /// This method returns an error under the same circumstances that
    /// `compress` does.
    pub fn compress_vec(&mut self, input: &[u8]) -> Result<Vec<u8>> {
        let mut buf = vec![0; max_compress_len(input.len())];
        let n = self.compress(input, &mut buf)?;
        buf.truncate(n);
        Ok(buf)
    }
}

struct Block<'s, 'd> {
    src: &'s [u8],
    s: usize,
    s_limit: usize,
    dst: &'d mut [u8],
    d: usize,
    next_emit: usize,
}

impl<'s, 'd> Block<'s, 'd> {
    #[inline(always)]
    fn new(src: &'s [u8], dst: &'d mut [u8], d: usize) -> Block<'s, 'd> {
        Block {
            src: src,
            s: 0,
            s_limit: src.len(),
            dst: dst,
            d: d,
            next_emit: 0,
        }
    }

    #[inline(always)]
    fn compress(&mut self, mut table: BlockTable<'_>) {
        debug_assert!(!table.is_empty());
        debug_assert!(self.src.len() >= MIN_NON_LITERAL_BLOCK_SIZE);

        self.s += 1;
        self.s_limit -= INPUT_MARGIN;
        let mut next_hash =
            table.hash(bytes::read_u32_le(&self.src[self.s..]));
        loop {
            let mut skip = 32;
            let mut candidate;
            let mut s_next = self.s;
            loop {
                self.s = s_next;
                let bytes_between_hash_lookups = skip >> 5;
                s_next = self.s + bytes_between_hash_lookups;
                skip += bytes_between_hash_lookups;
                if s_next > self.s_limit {
                    return self.done();
                }
                unsafe {
                    // SAFETY: next_hash is always computed by table.hash
                    // which is guaranteed to be in bounds.
                    candidate = *table.get_unchecked(next_hash) as usize;
                    *table.get_unchecked_mut(next_hash) = self.s as u16;

                    let srcp = self.src.as_ptr();
                    // SAFETY: s_next is guaranteed to be less than s_limit by
                    // the conditional above, which implies s_next is in
                    // bounds.
                    let x = bytes::loadu_u32_le(srcp.add(s_next));
                    next_hash = table.hash(x);
                    // SAFETY: self.s is always less than s_next, so it is also
                    // in bounds by the argument above.
                    //
                    // candidate is extracted from table, which is only ever
                    // set to valid positions in the block and is therefore
                    // also in bounds.
                    //
                    // We only need to compare y/z for equality, so we don't
                    // need to both with endianness. cur corresponds to the
                    // bytes at the current position and cand corresponds to
                    // a potential match. If they're equal, we declare victory
                    // and move below to try and extend the match.
                    let cur = bytes::loadu_u32_ne(srcp.add(self.s));
                    let cand = bytes::loadu_u32_ne(srcp.add(candidate));
                    if cur == cand {
                        break;
                    }
                }
            }
            // While the above found a candidate for compression, before we
            // emit a copy operation for it, we need to make sure that we emit
            // any bytes between the last copy operation and this one as a
            // literal.
            let lit_end = self.s;
            unsafe {
                // SAFETY: next_emit is set to a previous value of self.s,
                // which is guaranteed to be less than s_limit (in bounds).
                // lit_end is set to the current value of self.s, also
                // guaranteed to be less than s_limit (in bounds).
                self.emit_literal(lit_end);
            }
            loop {
                // Look for more matching bytes starting at the position of
                // the candidate and the current src position. We increment
                // self.s and candidate by 4 since we already know the first 4
                // bytes match.
                let base = self.s;
                self.s += 4;
                unsafe {
                    // SAFETY: candidate is always set to a value from our
                    // hash table, which only contains positions in self.src
                    // that have been seen for this block that occurred before
                    // self.s.
                    self.extend_match(candidate + 4);
                }
                let (offset, len) = (base - candidate, self.s - base);
                self.emit_copy(offset, len);
                self.next_emit = self.s;
                if self.s >= self.s_limit {
                    return self.done();
                }
                // Update the hash table with the byte sequences
                // self.src[self.s - 1..self.s + 3] and
                // self.src[self.s..self.s + 4]. Instead of reading 4 bytes
                // twice, we read 8 bytes once.
                //
                // If we happen to get a hit on self.src[self.s..self.s + 4],
                // then continue this loop and extend the match.
                unsafe {
                    let srcp = self.src.as_ptr();
                    // SAFETY: self.s can never exceed s_limit given by the
                    // conditional above and self.s is guaranteed to be
                    // non-zero and is therefore in bounds.
                    let x = bytes::loadu_u64_le(srcp.add(self.s - 1));
                    // The lower 4 bytes of x correspond to
                    // self.src[self.s - 1..self.s + 3].
                    let prev_hash = table.hash(x as u32);
                    // SAFETY: Hash values are guaranteed to be in bounds.
                    *table.get_unchecked_mut(prev_hash) = (self.s - 1) as u16;
                    // The lower 4 bytes of x>>8 correspond to
                    // self.src[self.s..self.s + 4].
                    let cur_hash = table.hash((x >> 8) as u32);
                    // SAFETY: Hash values are guaranteed to be in bounds.
                    candidate = *table.get_unchecked(cur_hash) as usize;
                    *table.get_unchecked_mut(cur_hash) = self.s as u16;

                    // SAFETY: candidate is set from table, which always
                    // contains valid positions in the current block.
                    let y = bytes::loadu_u32_le(srcp.add(candidate));
                    if (x >> 8) as u32 != y {
                        // If we didn't get a hit, update the next hash
                        // and move on. Our initial 8 byte read continues to
                        // pay off.
                        next_hash = table.hash((x >> 16) as u32);
                        self.s += 1;
                        break;
                    }
                }
            }
        }
    }

    /// Emits one or more copy operations with the given offset and length.
    /// offset must be in the range [1, 65535] and len must be in the range
    /// [4, 65535].
    #[inline(always)]
    fn emit_copy(&mut self, offset: usize, mut len: usize) {
        debug_assert!(1 <= offset && offset <= 65535);
        // Copy operations only allow lengths up to 64, but we'll allow bigger
        // lengths and emit as many operations as we need.
        //
        // N.B. Since our block size is 64KB, we never actually emit a copy 4
        // operation.
        debug_assert!(4 <= len && len <= 65535);

        // Emit copy 2 operations until we don't have to.
        // We check on 68 here and emit a shorter copy than 64 below because
        // it is cheaper to, e.g., encode a length 67 copy as a length 60
        // copy 2 followed by a length 7 copy 1 than to encode it as a length
        // 64 copy 2 followed by a length 3 copy 2. They key here is that a
        // copy 1 operation requires at least length 4 which forces a length 3
        // copy to use a copy 2 operation.
        while len >= 68 {
            self.emit_copy2(offset, 64);
            len -= 64;
        }
        if len > 64 {
            self.emit_copy2(offset, 60);
            len -= 60;
        }
        // If we can squeeze the last copy into a copy 1 operation, do it.
        if len <= 11 && offset <= 2047 {
            self.dst[self.d] = (((offset >> 8) as u8) << 5)
                | (((len - 4) as u8) << 2)
                | (Tag::Copy1 as u8);
            self.dst[self.d + 1] = offset as u8;
            self.d += 2;
        } else {
            self.emit_copy2(offset, len);
        }
    }

    /// Emits a "copy 2" operation with the given offset and length. The
    /// offset and length must be valid for a copy 2 operation. i.e., offset
    /// must be in the range [1, 65535] and len must be in the range [1, 64].
    #[inline(always)]
    fn emit_copy2(&mut self, offset: usize, len: usize) {
        debug_assert!(1 <= offset && offset <= 65535);
        debug_assert!(1 <= len && len <= 64);
        self.dst[self.d] = (((len - 1) as u8) << 2) | (Tag::Copy2 as u8);
        bytes::write_u16_le(offset as u16, &mut self.dst[self.d + 1..]);
        self.d += 3;
    }

    /// Attempts to extend a match from the current position in self.src with
    /// the candidate position given.
    ///
    /// This method uses unaligned loads and elides bounds checks, so the
    /// caller must guarantee that cand points to a valid location in self.src
    /// and is less than the current position in src.
    #[inline(always)]
    unsafe fn extend_match(&mut self, mut cand: usize) {
        debug_assert!(cand < self.s);
        while self.s + 8 <= self.src.len() {
            let srcp = self.src.as_ptr();
            // SAFETY: The loop invariant guarantees that there is at least
            // 8 bytes to read at self.src + self.s. Since cand must be
            // guaranteed by the caller to be valid and less than self.s, it
            // also has enough room to read 8 bytes.
            //
            // TODO(ag): Despite my best efforts, I couldn't get this to
            // autovectorize with 128-bit loads. The logic after the loads
            // appears to be a little too clever...
            let x = bytes::loadu_u64_ne(srcp.add(self.s));
            let y = bytes::loadu_u64_ne(srcp.add(cand));
            if x == y {
                // If all 8 bytes are equal, move on...
                self.s += 8;
                cand += 8;
            } else {
                // Otherwise, find the last byte that was equal. We can do
                // this efficiently by interpreted x/y as little endian
                // numbers, which lets us use the number of trailing zeroes
                // as a proxy for the number of equivalent bits (after an XOR).
                let z = x.to_le() ^ y.to_le();
                self.s += z.trailing_zeros() as usize / 8;
                return;
            }
        }
        // When we have fewer than 8 bytes left in the block, fall back to the
        // slow loop.
        while self.s < self.src.len() && self.src[self.s] == self.src[cand] {
            self.s += 1;
            cand += 1;
        }
    }

    /// Executes any cleanup when the current block has finished compressing.
    /// In particular, it emits any leftover bytes as a literal.
    #[inline(always)]
    fn done(&mut self) {
        if self.next_emit < self.src.len() {
            let lit_end = self.src.len();
            unsafe {
                // SAFETY: Both next_emit and lit_end are trivially in bounds
                // given the conditional and definition above.
                self.emit_literal(lit_end);
            }
        }
    }

    /// Emits a literal from self.src[self.next_emit..lit_end].
    ///
    /// This uses unaligned loads and elides bounds checks, so the caller must
    /// guarantee that self.src[self.next_emit..lit_end] is valid.
    #[inline(always)]
    unsafe fn emit_literal(&mut self, lit_end: usize) {
        let lit_start = self.next_emit;
        let len = lit_end - lit_start;
        let n = len.checked_sub(1).unwrap();
        if n <= 59 {
            self.dst[self.d] = ((n as u8) << 2) | (Tag::Literal as u8);
            self.d += 1;
            if len <= 16 && lit_start + 16 <= self.src.len() {
                // SAFETY: lit_start is equivalent to self.next_emit, which is
                // only set to self.s immediately following a copy emit. The
                // conditional above also ensures that there is at least 16
                // bytes of room in both src and dst.
                //
                // dst is big enough because the buffer is guaranteed to
                // be big enough to hold biggest possible compressed size plus
                // an extra 32 bytes, which exceeds the 16 byte copy here.
                let srcp = self.src.as_ptr().add(lit_start);
                let dstp = self.dst.as_mut_ptr().add(self.d);
                ptr::copy_nonoverlapping(srcp, dstp, 16);
                self.d += len;
                return;
            }
        } else if n < 256 {
            self.dst[self.d] = (60 << 2) | (Tag::Literal as u8);
            self.dst[self.d + 1] = n as u8;
            self.d += 2;
        } else {
            self.dst[self.d] = (61 << 2) | (Tag::Literal as u8);
            bytes::write_u16_le(n as u16, &mut self.dst[self.d + 1..]);
            self.d += 3;
        }
        // SAFETY: lit_start is equivalent to self.next_emit, which is only set
        // to self.s immediately following a copy, which implies that it always
        // points to valid bytes in self.src.
        //
        // We can't guarantee that there are at least len bytes though, which
        // must be guaranteed by the caller and is why this method is unsafe.
        let srcp = self.src.as_ptr().add(lit_start);
        let dstp = self.dst.as_mut_ptr().add(self.d);
        ptr::copy_nonoverlapping(srcp, dstp, len);
        self.d += len;
    }
}

/// `BlockTable` is a map from 4 byte sequences to positions of their most
/// recent occurrence in a block. In particular, this table lets us quickly
/// find candidates for compression.
///
/// We expose the `hash` method so that callers can be fastidious about the
/// number of times a hash is computed.
struct BlockTable<'a> {
    table: &'a mut [u16],
    /// The number of bits required to shift the hash such that the result
    /// is less than table.len().
    shift: u32,
}

impl Encoder {
    fn block_table(&mut self, block_size: usize) -> BlockTable<'_> {
        let mut shift: u32 = 32 - 8;
        let mut table_size = 256;
        while table_size < MAX_TABLE_SIZE && table_size < block_size {
            shift -= 1;
            table_size *= 2;
        }
        // If our block size is small, then use a small stack allocated table
        // instead of putting a bigger one on the heap. This particular
        // optimization is important if the caller is using Snappy to compress
        // many small blocks. (The memset savings alone is considerable.)
        let table: &mut [u16] = if table_size <= SMALL_TABLE_SIZE {
            &mut self.small[0..table_size]
        } else {
            if self.big.is_empty() {
                // Interestingly, using `self.big.resize` here led to some
                // very weird code getting generated that led to a large
                // slow down. Forcing the issue with a new vec seems to
                // fix it. ---AG
                self.big = vec![0; MAX_TABLE_SIZE];
            }
            &mut self.big[0..table_size]
        };
        for x in &mut *table {
            *x = 0;
        }
        BlockTable { table: table, shift: shift }
    }
}

impl<'a> BlockTable<'a> {
    #[inline(always)]
    fn hash(&self, x: u32) -> usize {
        (x.wrapping_mul(0x1E35A7BD) >> self.shift) as usize
    }
}

impl<'a> Deref for BlockTable<'a> {
    type Target = [u16];
    fn deref(&self) -> &[u16] {
        self.table
    }
}

impl<'a> DerefMut for BlockTable<'a> {
    fn deref_mut(&mut self) -> &mut [u16] {
        self.table
    }
}