//! Raw in-memory stream compression/decompression.

//!

//! This module defines a `Decoder` and an `Encoder` to decode/encode streams

//! of data using buffers.

//!

//! They are mostly thin wrappers around `zstd_safe::{DCtx, CCtx}`.

use std::io;

pub use zstd_safe::{CParameter, DParameter, InBuffer, OutBuffer, WriteBuf};

use crate::dict::{DecoderDictionary, EncoderDictionary};

use crate::map_error_code;

/// Represents an abstract compression/decompression operation.

///

/// This trait covers both `Encoder` and `Decoder`.

pub trait Operation {

    /// Performs a single step of this operation.

    ///

    /// Should return a hint for the next input size.

    ///

    /// If the result is `Ok(0)`, it may indicate that a frame was just

    /// finished.

    fn run<C: WriteBuf + ?Sized>(

        &mut self,

        input: &mut InBuffer<'_>,

        output: &mut OutBuffer<'_, C>,

    ) -> io::Result<usize>;

    /// Performs a single step of this operation.

    ///

    /// This is a comvenience wrapper around `Operation::run` if you don't

    /// want to deal with `InBuffer`/`OutBuffer`.

    fn run_on_buffers(

        &mut self,

        input: &[u8],

        output: &mut [u8],

    ) -> io::Result<Status> {

        let mut input = InBuffer::around(input);

        let mut output = OutBuffer::around(output);

        let remaining = self.run(&mut input, &mut output)?;

        Ok(Status {

            remaining,

            bytes_read: input.pos(),

            bytes_written: output.pos(),

        })

    }

    /// Flushes any internal buffer, if any.

    ///

    /// Returns the number of bytes still in the buffer.

    /// To flush entirely, keep calling until it returns `Ok(0)`.

    fn flush<C: WriteBuf + ?Sized>(

        &mut self,

        output: &mut OutBuffer<'_, C>,

    ) -> io::Result<usize> {

        let _ = output;

        Ok(0)

    }

    /// Prepares the operation for a new frame.

    ///

    /// This is hopefully cheaper than creating a new operation.

    fn reinit(&mut self) -> io::Result<()> {

        Ok(())

    }

    /// Finishes the operation, writing any footer if necessary.

    ///

    /// Returns the number of bytes still to write.

    ///

    /// Keep calling this method until it returns `Ok(0)`,

    /// and then don't ever call this method.

    fn finish<C: WriteBuf + ?Sized>(

        &mut self,

        output: &mut OutBuffer<'_, C>,

        finished_frame: bool,

    ) -> io::Result<usize> {

        let _ = output;

        let _ = finished_frame;

        Ok(0)

    }

}

/// Dummy operation that just copies its input to the output.

pub struct NoOp;

impl Operation for NoOp {

    fn run<C: WriteBuf + ?Sized>(

        &mut self,

        input: &mut InBuffer<'_>,

        output: &mut OutBuffer<'_, C>,

    ) -> io::Result<usize> {

        // Skip the prelude

        let src = &input.src[input.pos..];

        // Safe because `output.pos() <= output.capacity()`.

        let output_pos = output.pos();

        let dst = unsafe { output.as_mut_ptr().add(output_pos) };

        // Ignore anything past the end

        let len = usize::min(src.len(), output.capacity() - output_pos);

        let src = &src[..len];

        // Safe because:

        // * `len` is less than either of the two lengths

        // * `src` and `dst` do not overlap because we have `&mut` to each.

        unsafe { std::ptr::copy_nonoverlapping(src.as_ptr(), dst, len) };

        input.set_pos(input.pos() + len);

        unsafe { output.set_pos(output_pos + len) };

        Ok(0)

    }

}

/// Describes the result of an operation.

pub struct Status {

    /// Number of bytes expected for next input.

    ///

    /// * If `remaining = 0`, then we are at the end of a frame.

    /// * If `remaining > 0`, then it's just a hint for how much there is still

    ///   to read.

    pub remaining: usize,

    /// Number of bytes read from the input.

    pub bytes_read: usize,

    /// Number of bytes written to the output.

    pub bytes_written: usize,

}

/// An in-memory decoder for streams of data.

pub struct Decoder<'a> {

    context: MaybeOwnedDCtx<'a>,

}

fn create_dctx<'a>() -> io::Result<zstd_safe::DCtx<'a>> {

    zstd_safe::DCtx::try_create().ok_or_else(|| {

        io::Error::new(

            io::ErrorKind::Other,

            "zstd could not allocate a decompression context",

        )

    })

}

impl Decoder<'static> {

    /// Creates a new decoder.

    pub fn new() -> io::Result<Self> {

        Self::with_dictionary(&[])

    }

    /// Creates a new decoder initialized with the given dictionary.

    pub fn with_dictionary(dictionary: &[u8]) -> io::Result<Self> {

        let mut context = create_dctx()?;

        context.init().map_err(map_error_code)?;

        context

            .load_dictionary(dictionary)

            .map_err(map_error_code)?;

        Ok(Decoder {

            context: MaybeOwnedDCtx::Owned(context),

        })

    }

}

impl<'a> Decoder<'a> {

    /// Creates a new decoder which employs the provided context for deserialization.

    ///

    /// The context's session is reset first: anything left over from a frame

    /// that a previous user of the context did not finish is discarded.

    /// Parameters and dictionary are kept.

    pub fn with_context(context: &'a mut zstd_safe::DCtx<'static>) -> Self {

        // Otherwise we would start by handing out what is left of the

        // previous decoder's output. Resetting only the session cannot fail.

        let _ = context.reset(zstd_safe::ResetDirective::SessionOnly);

        Self {

            context: MaybeOwnedDCtx::Borrowed(context),

        }

    }

    /// Creates a new decoder, using an existing `DecoderDictionary`.

    pub fn with_prepared_dictionary<'b>(

        dictionary: &'a DecoderDictionary<'b>,

    ) -> io::Result<Self>

    where

        'b: 'a,

    {

        let mut context = create_dctx()?;

        context

            .ref_ddict(dictionary.as_ddict())

            .map_err(map_error_code)?;

        Ok(Decoder {

            context: MaybeOwnedDCtx::Owned(context),

        })

    }

    /// Creates a new decoder, using a ref prefix

    pub fn with_ref_prefix<'b>(ref_prefix: &'b [u8]) -> io::Result<Self>

    where

        'b: 'a,

    {

        let mut context = create_dctx()?;

        context.ref_prefix(ref_prefix).map_err(map_error_code)?;

        Ok(Decoder {

            context: MaybeOwnedDCtx::Owned(context),

        })

    }

    /// Sets a decompression parameter for this decoder.

    pub fn set_parameter(&mut self, parameter: DParameter) -> io::Result<()> {

        match &mut self.context {

            MaybeOwnedDCtx::Owned(x) => x.set_parameter(parameter),

            MaybeOwnedDCtx::Borrowed(x) => x.set_parameter(parameter),

        }

        .map_err(map_error_code)?;

        Ok(())

    }

}

impl Operation for Decoder<'_> {

    fn run<C: WriteBuf + ?Sized>(

        &mut self,

        input: &mut InBuffer<'_>,

        output: &mut OutBuffer<'_, C>,

    ) -> io::Result<usize> {

        match &mut self.context {

            MaybeOwnedDCtx::Owned(x) => x.decompress_stream(output, input),

            MaybeOwnedDCtx::Borrowed(x) => x.decompress_stream(output, input),

        }

        .map_err(map_error_code)

    }

    fn flush<C: WriteBuf + ?Sized>(

        &mut self,

        output: &mut OutBuffer<'_, C>,

    ) -> io::Result<usize> {

        // To flush, we just offer no additional input.

        self.run(&mut InBuffer::around(&[]), output)?;

        // We don't _know_ how much (decompressed data) there is still in buffer.

        if output.pos() < output.capacity() {

            // We only know when there's none (the output buffer is not full).

            Ok(0)

        } else {

            // Otherwise, pretend there's still "1 byte" remaining.

            Ok(1)

        }

    }

    fn reinit(&mut self) -> io::Result<()> {

        match &mut self.context {

            MaybeOwnedDCtx::Owned(x) => {

                x.reset(zstd_safe::ResetDirective::SessionOnly)

            }

            MaybeOwnedDCtx::Borrowed(x) => {

                x.reset(zstd_safe::ResetDirective::SessionOnly)

            }

        }

        .map_err(map_error_code)?;

        Ok(())

    }

    fn finish<C: WriteBuf + ?Sized>(

        &mut self,

        _output: &mut OutBuffer<'_, C>,

        finished_frame: bool,

    ) -> io::Result<usize> {

        if finished_frame {

            Ok(0)

        } else {

            Err(io::Error::new(

                io::ErrorKind::UnexpectedEof,

                "incomplete frame",

            ))

        }

    }

}

/// An in-memory encoder for streams of data.

pub struct Encoder<'a> {

    context: MaybeOwnedCCtx<'a>,

}

impl Encoder<'static> {

    /// Creates a new encoder.

    pub fn new(level: i32) -> io::Result<Self> {

        Self::with_dictionary(level, &[])

    }

    /// Creates a new encoder initialized with the given dictionary.

    pub fn with_dictionary(level: i32, dictionary: &[u8]) -> io::Result<Self> {

        let mut context = zstd_safe::CCtx::create();

        context

            .set_parameter(CParameter::CompressionLevel(level))

            .map_err(map_error_code)?;

        context

            .load_dictionary(dictionary)

            .map_err(map_error_code)?;

        Ok(Encoder {

            context: MaybeOwnedCCtx::Owned(context),

        })

    }

}

impl<'a> Encoder<'a> {

    /// Creates a new encoder that uses the provided context for serialization.

    ///

    /// The context's session is reset first: anything left over from a frame

    /// that a previous user of the context did not finish is discarded, and

    /// so is the pledged source size (set it on the encoder instead).

    /// Parameters and dictionary are kept.

    pub fn with_context(context: &'a mut zstd_safe::CCtx<'static>) -> Self {

        // Otherwise the input buffered by the previous encoder would become

        // the start of our frame. Resetting only the session cannot fail.

        let _ = context.reset(zstd_safe::ResetDirective::SessionOnly);

        Self {

            context: MaybeOwnedCCtx::Borrowed(context),

        }

    }

    /// Creates a new encoder using an existing `EncoderDictionary`.

    pub fn with_prepared_dictionary<'b>(

        dictionary: &'a EncoderDictionary<'b>,

    ) -> io::Result<Self>

    where

        'b: 'a,

    {

        let mut context = zstd_safe::CCtx::create();

        context

            .ref_cdict(dictionary.as_cdict())

            .map_err(map_error_code)?;

        Ok(Encoder {

            context: MaybeOwnedCCtx::Owned(context),

        })

    }

    /// Creates a new encoder initialized with the given ref prefix.

    pub fn with_ref_prefix<'b>(

        level: i32,

        ref_prefix: &'b [u8],

    ) -> io::Result<Self>

    where

        'b: 'a,

    {

        let mut context = zstd_safe::CCtx::create();

        context

            .set_parameter(CParameter::CompressionLevel(level))

            .map_err(map_error_code)?;

        context.ref_prefix(ref_prefix).map_err(map_error_code)?;

        Ok(Encoder {

            context: MaybeOwnedCCtx::Owned(context),

        })

    }

    /// Sets a compression parameter for this encoder.

    pub fn set_parameter(&mut self, parameter: CParameter) -> io::Result<()> {

        match &mut self.context {

            MaybeOwnedCCtx::Owned(x) => x.set_parameter(parameter),

            MaybeOwnedCCtx::Borrowed(x) => x.set_parameter(parameter),

        }

        .map_err(map_error_code)?;

        Ok(())

    }

    /// Sets the size of the input expected by zstd.

    ///

    /// May affect compression ratio.

    ///

    /// It is an error to give an incorrect size (an error _will_ be returned when closing the

    /// stream).

    ///

    /// If `None` is given, it assume the size is not known (default behaviour).

    pub fn set_pledged_src_size(

        &mut self,

        pledged_src_size: Option<u64>,

    ) -> io::Result<()> {

        match &mut self.context {

            MaybeOwnedCCtx::Owned(x) => {

                x.set_pledged_src_size(pledged_src_size)

            }

            MaybeOwnedCCtx::Borrowed(x) => {

                x.set_pledged_src_size(pledged_src_size)

            }

        }

        .map_err(map_error_code)?;

        Ok(())

    }

}

impl<'a> Operation for Encoder<'a> {

    fn run<C: WriteBuf + ?Sized>(

        &mut self,

        input: &mut InBuffer<'_>,

        output: &mut OutBuffer<'_, C>,

    ) -> io::Result<usize> {

        match &mut self.context {

            MaybeOwnedCCtx::Owned(x) => x.compress_stream(output, input),

            MaybeOwnedCCtx::Borrowed(x) => x.compress_stream(output, input),

        }

        .map_err(map_error_code)

    }

    fn flush<C: WriteBuf + ?Sized>(

        &mut self,

        output: &mut OutBuffer<'_, C>,

    ) -> io::Result<usize> {

        match &mut self.context {

            MaybeOwnedCCtx::Owned(x) => x.flush_stream(output),

            MaybeOwnedCCtx::Borrowed(x) => x.flush_stream(output),

        }

        .map_err(map_error_code)

    }

    fn finish<C: WriteBuf + ?Sized>(

        &mut self,

        output: &mut OutBuffer<'_, C>,

        _finished_frame: bool,

    ) -> io::Result<usize> {

        match &mut self.context {

            MaybeOwnedCCtx::Owned(x) => x.end_stream(output),

            MaybeOwnedCCtx::Borrowed(x) => x.end_stream(output),

        }

        .map_err(map_error_code)

    }

    fn reinit(&mut self) -> io::Result<()> {

        match &mut self.context {

            MaybeOwnedCCtx::Owned(x) => {

                x.reset(zstd_safe::ResetDirective::SessionOnly)

            }

            MaybeOwnedCCtx::Borrowed(x) => {

                x.reset(zstd_safe::ResetDirective::SessionOnly)

            }

        }

        .map_err(map_error_code)?;

        Ok(())

    }

}

enum MaybeOwnedCCtx<'a> {

    Owned(zstd_safe::CCtx<'a>),

    Borrowed(&'a mut zstd_safe::CCtx<'static>),

}

enum MaybeOwnedDCtx<'a> {

    Owned(zstd_safe::DCtx<'a>),

    Borrowed(&'a mut zstd_safe::DCtx<'static>),

}

#[cfg(test)]

mod tests {

    // This requires impl for [u8; N] which is currently behind a feature.

    #[cfg(feature = "arrays")]

    #[test]

    fn test_cycle() {

        use super::{Decoder, Encoder, InBuffer, Operation, OutBuffer};

        let mut encoder = Encoder::new(1).unwrap();

        let mut decoder = Decoder::new().unwrap();

        // Step 1: compress

        let mut input = InBuffer::around(b"AbcdefAbcdefabcdef");

        let mut output = [0u8; 128];

        let mut output = OutBuffer::around(&mut output);

        loop {

            encoder.run(&mut input, &mut output).unwrap();

            if input.pos == input.src.len() {

                break;

            }

        }

        encoder.finish(&mut output, true).unwrap();

        let initial_data = input.src;

        // Step 2: decompress

        let mut input = InBuffer::around(output.as_slice());

        let mut output = [0u8; 128];

        let mut output = OutBuffer::around(&mut output);

        loop {

            decoder.run(&mut input, &mut output).unwrap();

            if input.pos == input.src.len() {

                break;

            }

        }

        assert_eq!(initial_data, output.as_slice());

    }

}