1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

use std::io::{Read, Result, Write};

use lz4::{
    BlockMode,
    BlockSize,
    Decoder,
    Encoder,
    EncoderBuilder,
};
use serde::{
    Deserialize,
    Serialize,
};

use super::{
    Compression,
};


// From: https://github.com/bozaro/lz4-rs/issues/9
// Kludge to finish Lz4 encoder on Drop.
struct Wrapper<W: Write> {
    s: Option<Encoder<W>>,
}

impl<W: Write> Write for Wrapper<W> {
    fn write(&mut self, buffer: &[u8]) -> Result<usize> {
        self.s.as_mut().unwrap().write(buffer)
    }

    fn flush(&mut self) -> Result<()> {
        self.s.as_mut().unwrap().flush()
    }
}

impl<W: Write> Drop for Wrapper<W> {
    fn drop(&mut self) {
        if let Some(s) = self.s.take() {
            s.finish().1.unwrap();
        }
    }
}


#[derive(Clone, Serialize, Deserialize, PartialEq, Debug)]
#[serde(rename_all = "camelCase")]
pub struct Lz4Compression {
    #[serde(default = "default_lz4_block_size")]
    block_size: i32,
}

impl Lz4Compression {
    /// `lz4` uses an enum for specifying block size, so choose the smallest
    /// larger size from that enum.
    fn get_effective_block_size(&self) -> BlockSize {
        if self.block_size <= BlockSize::Max64KB.get_size() as i32 {
            BlockSize::Max64KB
        } else if self.block_size as usize <= BlockSize::Max256KB.get_size() {
            BlockSize::Max256KB
        } else if self.block_size as usize <= BlockSize::Max1MB.get_size() {
            BlockSize::Max1MB
        } else {
            BlockSize::Max4MB
        }
    }
}

fn default_lz4_block_size() -> i32 {65_536}

impl Default for Lz4Compression {
    fn default() -> Lz4Compression {
        Lz4Compression {
            block_size: default_lz4_block_size(),
        }
    }
}

impl Compression for Lz4Compression {
    fn decoder<'a, R: Read + 'a>(&self, r: R) -> Box<dyn Read + 'a> {
        Box::new(Decoder::new(r).expect("TODO: LZ4 returns a result here"))
    }

    fn encoder<'a, W: Write + 'a>(&self, w: W) -> Box<dyn Write + 'a> {
        let encoder = EncoderBuilder::new()
            .block_size(self.get_effective_block_size())
            .block_mode(BlockMode::Independent)
            .build(w)
            .expect("TODO");
        Box::new(Wrapper {s: Some(encoder)})
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::compression::CompressionType;

    const TEST_BLOCK_I16_LZ4: [u8; 47] = [
        0x00, 0x00,
        0x00, 0x03,
        0x00, 0x00, 0x00, 0x01,
        0x00, 0x00, 0x00, 0x02,
        0x00, 0x00, 0x00, 0x03,
        0x04, 0x22, 0x4d, 0x18,
        0x64, 0x40, 0xa7, 0x0c,
        0x00, 0x00, 0x80, 0x00,
        0x01, 0x00, 0x02, 0x00,
        0x03, 0x00, 0x04, 0x00,
        0x05, 0x00, 0x06, 0x00,
        0x00, 0x00, 0x00, 0x41,
        0x37, 0x33, 0x08,
    ];

    #[test]
    fn test_read_doc_spec_block() {
        crate::tests::test_read_doc_spec_block(
            TEST_BLOCK_I16_LZ4.as_ref(),
            CompressionType::Lz4(Lz4Compression::default()));
    }

    #[test]
    fn test_write_doc_spec_block() {
        crate::tests::test_write_doc_spec_block(
            TEST_BLOCK_I16_LZ4.as_ref(),
            CompressionType::Lz4(Lz4Compression::default()));
    }

    #[test]
    fn test_rw() {
        crate::tests::test_block_compression_rw(CompressionType::Lz4(Lz4Compression::default()));
    }
}