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

use std::io::{Write, Result as IoResult};

use ll;

struct EncoderContext {
    c: ll::ZBUFFCompressionContext,
}

impl EncoderContext {
    fn new() -> Self {
        EncoderContext { c: unsafe { ll::ZBUFF_createCCtx() } }
    }
}

impl Drop for EncoderContext {
    fn drop(&mut self) {
        let code = unsafe { ll::ZBUFF_freeCCtx(self.c) };
        ll::parse_code(code).unwrap();
    }
}


/// An encoder that compress and forward data to another writer.
///
/// The zstd library has an internal input buffer (~128kb).
pub struct Encoder<W: Write> {
    // output writer (compressed data)
    writer: W,
    // output buffer
    buffer: Vec<u8>,

    // compression context
    context: EncoderContext,
}

impl<W: Write> Encoder<W> {
    /// Creates a new encoder.
    ///
    /// `level`: compression level (1-21)
    pub fn new(writer: W, level: i32) -> IoResult<Self> {

        let buffer_size = unsafe { ll::ZBUFF_recommendedCOutSize() };

        let context = EncoderContext::new();

        // Initialize the stream
        try!(ll::parse_code(unsafe { ll::ZBUFF_compressInit(context.c, level) }));

        Ok(Encoder {
            writer: writer,
            buffer: Vec::with_capacity(buffer_size),
            context: context,
        })
    }

    /// Finishes the stream. You need to call this after writing your stuff.
    ///
    /// This returns the inner writer in case you need it.
    pub fn finish(mut self) -> IoResult<W> {

        // First, closes the stream.
        let mut out_size = self.buffer.capacity();
        let remaining = try!(ll::parse_code(unsafe {
            ll::ZBUFF_compressEnd(self.context.c, self.buffer.as_mut_ptr(), &mut out_size)
        }));
        unsafe {
            self.buffer.set_len(out_size);
        }
        if remaining != 0 {
            // Need to flush?
            panic!("Need to flush, but I'm lazy.");
        }

        // Write the end out
        try!(self.writer.write_all(&self.buffer));

        // Return the writer, because why not
        Ok(self.writer)
    }
}

impl<W: Write> Write for Encoder<W> {
    fn write(&mut self, buf: &[u8]) -> IoResult<usize> {
        // How much we've read from this task
        let mut read = 0;
        while read != buf.len() {
            let mut out_size = self.buffer.capacity();
            let mut in_size = buf.len() - read;

            unsafe {
                // Compress the given buffer into our output buffer
                let code = ll::ZBUFF_compressContinue(self.context.c,
                                                      self.buffer.as_mut_ptr(),
                                                      &mut out_size,
                                                      buf[read..].as_ptr(),
                                                      &mut in_size);
                self.buffer.set_len(out_size);

                // Do we care about the hint?
                let hint = try!(ll::parse_code(code));
            }
            try!(self.writer.write_all(&self.buffer));
            read += in_size;
        }
        Ok(read)
    }

    fn flush(&mut self) -> IoResult<()> {
        let mut out_size = self.buffer.capacity();
        let written = try!(ll::parse_code(unsafe {
            ll::ZBUFF_compressFlush(self.context.c, self.buffer.as_mut_ptr(), &mut out_size)
        }));
        unsafe {
            self.buffer.set_len(written);
        }

        try!(self.writer.write_all(&self.buffer));
        Ok(())
    }
}