async-compression 0.1.0-alpha.4

use std::{
    io::Result,
    mem,
    pin::Pin,
    task::{Context, Poll},
};

use crate::unshared::Unshared;
use bytes::{Bytes, BytesMut};
use futures::{ready, stream::Stream};
use libzstd::stream::raw::{Decoder, Encoder, Operation};
use pin_project::unsafe_project;

#[derive(Debug)]
enum State {
    Reading,
    Writing(Bytes),
    Flushing,
    Done,
    Invalid,
}

#[derive(Debug)]
enum DeState {
    Reading,
    Writing(Bytes),
    Done,
    Invalid,
}

/// A zstd encoder, or compressor.
///
/// This structure implements a [`Stream`] interface and will read uncompressed data from an
/// underlying stream and emit a stream of compressed data.
#[unsafe_project(Unpin)]
#[cfg_attr(docsrs, doc(cfg(feature = "zstd")))]
#[derive(Debug)]
pub struct ZstdEncoder<S: Stream<Item = Result<Bytes>>> {
    #[pin]
    inner: S,
    state: State,
    output: BytesMut,
    encoder: Unshared<Encoder>,
}

/// A zstd decoder, or decompressor.
///
/// This structure implements a [`Stream`] interface and will read compressed data from an
/// underlying stream and emit a stream of uncompressed data.
#[unsafe_project(Unpin)]
#[cfg_attr(docsrs, doc(cfg(feature = "zstd")))]
#[derive(Debug)]
pub struct ZstdDecoder<S: Stream<Item = Result<Bytes>>> {
    #[pin]
    inner: S,
    state: DeState,
    output: BytesMut,
    decoder: Unshared<Decoder>,
}

impl<S: Stream<Item = Result<Bytes>>> ZstdEncoder<S> {
    /// Creates a new encoder which will read uncompressed data from the given stream and emit a
    /// compressed stream.
    ///
    /// The `level` argument here can range from 1-21. A level of `0` will use zstd's default, which is `3`.
    pub fn new(stream: S, level: i32) -> ZstdEncoder<S> {
        ZstdEncoder {
            inner: stream,
            // zstd needs to have 0 bytes written to it to create a valid compressed 0 byte stream,
            // just flushing it after not writing to it is not enough.
            state: State::Writing(Bytes::new()),
            output: BytesMut::new(),
            encoder: Unshared::new(Encoder::new(level).unwrap()),
        }
    }

    /// Acquires a reference to the underlying stream that this encoder is wrapping.
    pub fn get_ref(&self) -> &S {
        &self.inner
    }

    /// Acquires a mutable reference to the underlying stream that this encoder is wrapping.
    ///
    /// Note that care must be taken to avoid tampering with the state of the stream which may
    /// otherwise confuse this encoder.
    pub fn get_mut(&mut self) -> &mut S {
        &mut self.inner
    }

    /// Acquires a pinned mutable reference to the underlying stream that this encoder is wrapping.
    ///
    /// Note that care must be taken to avoid tampering with the state of the stream which may
    /// otherwise confuse this encoder.
    pub fn get_pin_mut<'a>(self: Pin<&'a mut Self>) -> Pin<&'a mut S> {
        self.project().inner
    }

    /// Consumes this encoder returning the underlying stream.
    ///
    /// Note that this may discard internal state of this encoder, so care should be taken
    /// to avoid losing resources when this is called.
    pub fn into_inner(self) -> S {
        self.inner
    }
}

impl<S: Stream<Item = Result<Bytes>>> ZstdDecoder<S> {
    /// Creates a new decoder which will read compressed data from the given stream and emit an
    /// uncompressed stream.
    pub fn new(stream: S) -> ZstdDecoder<S> {
        ZstdDecoder {
            inner: stream,
            state: DeState::Reading,
            output: BytesMut::new(),
            decoder: Unshared::new(Decoder::new().unwrap()),
        }
    }

    /// Acquires a reference to the underlying stream that this decoder is wrapping.
    pub fn get_ref(&self) -> &S {
        &self.inner
    }

    /// Acquires a mutable reference to the underlying stream that this decoder is wrapping.
    ///
    /// Note that care must be taken to avoid tampering with the state of the stream which may
    /// otherwise confuse this decoder.
    pub fn get_mut(&mut self) -> &mut S {
        &mut self.inner
    }

    /// Acquires a pinned mutable reference to the underlying stream that this decoder is wrapping.
    ///
    /// Note that care must be taken to avoid tampering with the state of the stream which may
    /// otherwise confuse this decoder.
    pub fn get_pin_mut<'a>(self: Pin<&'a mut Self>) -> Pin<&'a mut S> {
        self.project().inner
    }

    /// Consumes this decoder returning the underlying stream.
    ///
    /// Note that this may discard internal state of this decoder, so care should be taken
    /// to avoid losing resources when this is called.
    pub fn into_inner(self) -> S {
        self.inner
    }
}

impl<S: Stream<Item = Result<Bytes>>> Stream for ZstdEncoder<S> {
    type Item = Result<Bytes>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Result<Bytes>>> {
        let mut this = self.project();

        fn compress(
            encoder: &mut Encoder,
            input: &mut Bytes,
            output: &mut BytesMut,
        ) -> Result<Bytes> {
            const OUTPUT_BUFFER_SIZE: usize = 8_000;

            if output.len() < OUTPUT_BUFFER_SIZE {
                output.resize(OUTPUT_BUFFER_SIZE, 0);
            }

            let status = encoder.run_on_buffers(input, output)?;
            input.advance(status.bytes_read);
            Ok(output.split_to(status.bytes_written).freeze())
        }

        #[allow(clippy::never_loop)] // https://github.com/rust-lang/rust-clippy/issues/4058
        loop {
            break match mem::replace(this.state, State::Invalid) {
                State::Reading => {
                    *this.state = State::Reading;
                    *this.state = match ready!(this.inner.as_mut().poll_next(cx)) {
                        Some(chunk) => State::Writing(chunk?),
                        None => State::Flushing,
                    };
                    continue;
                }
                State::Writing(mut input) => {
                    let chunk = compress(this.encoder.get_mut(), &mut input, &mut this.output)?;

                    if input.is_empty() {
                        *this.state = State::Reading;
                    } else {
                        *this.state = State::Writing(input);
                    }

                    if chunk.is_empty() {
                        continue;
                    }

                    Poll::Ready(Some(Ok(chunk)))
                }
                State::Flushing => {
                    let mut output = zstd_safe::OutBuffer::around(this.output);

                    let bytes_left = this.encoder.get_mut().flush(&mut output).unwrap();
                    *this.state = if bytes_left == 0 {
                        let _ = this.encoder.get_mut().finish(&mut output, true);
                        State::Done
                    } else {
                        State::Flushing
                    };
                    Poll::Ready(Some(Ok(output.as_slice().into())))
                }
                State::Done => Poll::Ready(None),
                State::Invalid => panic!("ZstdEncoder reached invalid state"),
            };
        }
    }
}

impl<S: Stream<Item = Result<Bytes>>> Stream for ZstdDecoder<S> {
    type Item = Result<Bytes>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Result<Bytes>>> {
        let mut this = self.project();

        fn decompress(
            decoder: &mut Decoder,
            input: &mut Bytes,
            output: &mut BytesMut,
        ) -> Result<Bytes> {
            const OUTPUT_BUFFER_SIZE: usize = 8_000;

            if output.len() < OUTPUT_BUFFER_SIZE {
                output.resize(OUTPUT_BUFFER_SIZE, 0);
            }

            let status = decoder.run_on_buffers(input, output)?;
            input.advance(status.bytes_read);
            Ok(output.split_to(status.bytes_written).freeze())
        }

        #[allow(clippy::never_loop)] // https://github.com/rust-lang/rust-clippy/issues/4058
        loop {
            break match mem::replace(this.state, DeState::Invalid) {
                DeState::Reading => {
                    *this.state = DeState::Reading;
                    *this.state = match ready!(this.inner.as_mut().poll_next(cx)) {
                        Some(chunk) => DeState::Writing(chunk?),
                        None => DeState::Done,
                    };
                    continue;
                }
                DeState::Writing(mut input) => {
                    if input.is_empty() {
                        *this.state = DeState::Reading;
                        continue;
                    }

                    let chunk = decompress(this.decoder.get_mut(), &mut input, &mut this.output)?;

                    *this.state = DeState::Writing(input);

                    Poll::Ready(Some(Ok(chunk)))
                }
                DeState::Done => Poll::Ready(None),
                DeState::Invalid => panic!("ZstdDecoder reached invalid state"),
            };
        }
    }
}

fn _assert() {
    crate::util::_assert_send::<ZstdEncoder<Pin<Box<dyn Stream<Item = Result<Bytes>> + Send>>>>();
    crate::util::_assert_sync::<ZstdEncoder<Pin<Box<dyn Stream<Item = Result<Bytes>> + Sync>>>>();
    crate::util::_assert_send::<ZstdDecoder<Pin<Box<dyn Stream<Item = Result<Bytes>> + Send>>>>();
    crate::util::_assert_sync::<ZstdDecoder<Pin<Box<dyn Stream<Item = Result<Bytes>> + Sync>>>>();
}