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

// Copyright 2017, 2019, 2020 Martin Pool.

// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation; either version 2 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

//! Snappy compression glue.

use bytes::{Bytes, BytesMut};
use snap::raw::{Decoder, Encoder};

use crate::Result;

pub(crate) struct Compressor {
    encoder: Encoder,
}

impl Compressor {
    pub fn new() -> Compressor {
        Compressor {
            encoder: Encoder::new(),
        }
    }

    /// Compress bytes into unframed Snappy data.
    pub fn compress(&mut self, input: &[u8]) -> Result<Bytes> {
        let max_len = snap::raw::max_compress_len(input.len());
        let mut out = BytesMut::zeroed(max_len);
        let actual_len = self.encoder.compress(input, &mut out)?;
        out.truncate(actual_len);
        Ok(out.freeze())
    }
}

#[derive(Default)]
pub(crate) struct Decompressor {
    decoder: Decoder,
}

impl Decompressor {
    pub fn new() -> Decompressor {
        Decompressor::default()
    }

    /// Decompressed unframed Snappy data.
    ///
    /// Returns a slice pointing into a reusable object inside the Decompressor.
    pub fn decompress(&mut self, input: &[u8]) -> Result<Bytes> {
        let max_len = snap::raw::decompress_len(input)?;
        let mut out = BytesMut::zeroed(max_len);
        let actual_len = self.decoder.decompress(input, &mut out)?;
        out.truncate(actual_len);
        Ok(out.freeze())
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn compressor_decompressor() {
        let mut compressor = Compressor::new();
        let mut decompressor = Decompressor::new();

        let comp = compressor.compress(b"hello world").unwrap();
        assert_eq!(comp.as_ref(), b"\x0b(hello world");
        assert_eq!(
            decompressor.decompress(&comp).unwrap().as_ref(),
            b"hello world"
        );

        let long_input = b"hello world, hello world, hello world, hello world";
        let comp = compressor.compress(long_input).unwrap();
        assert_eq!(comp.as_ref(), b"\x32\x30hello world, \x92\x0d\0");
        assert_eq!(decompressor.decompress(&comp).unwrap(), &long_input[..]);
    }
}