sharpcompress/src/SharpCompress/Compressors/ZStandard/CompressionStream.cs

using System;
using System.Buffers;
using System.IO;
using System.Threading;
using System.Threading.Tasks;
using SharpCompress.Compressors.ZStandard.Unsafe;

namespace SharpCompress.Compressors.ZStandard;

public partial class CompressionStream : Stream
{
    private readonly Stream innerStream;
    private readonly byte[] outputBuffer;
    private readonly bool preserveCompressor;
    private readonly bool leaveOpen;
    private Compressor? compressor;
    private ZSTD_outBuffer_s output;

    public CompressionStream(
        Stream stream,
        int level = Compressor.DefaultCompressionLevel,
        int bufferSize = 0,
        bool leaveOpen = true
    )
        : this(stream, new Compressor(level), bufferSize, false, leaveOpen) { }

    public CompressionStream(
        Stream stream,
        Compressor compressor,
        int bufferSize = 0,
        bool preserveCompressor = true,
        bool leaveOpen = true
    )
    {
        SharpCompress.ThrowHelper.ThrowIfNull(stream);

        if (!stream.CanWrite)
        {
            throw new ArgumentException("Stream is not writable", nameof(stream));
        }

        SharpCompress.ThrowHelper.ThrowIfNegative(bufferSize);

        innerStream = stream;
        this.compressor = compressor;
        this.preserveCompressor = preserveCompressor;
        this.leaveOpen = leaveOpen;

        var outputBufferSize =
            bufferSize > 0
                ? bufferSize
                : (int)Unsafe.Methods.ZSTD_CStreamOutSize().EnsureZstdSuccess();
        outputBuffer = ArrayPool<byte>.Shared.Rent(outputBufferSize);
        output = new ZSTD_outBuffer_s { pos = 0, size = (nuint)outputBufferSize };
    }

    public void SetParameter(ZSTD_cParameter parameter, int value)
    {
        EnsureNotDisposed();
        compressor.NotNull().SetParameter(parameter, value);
    }

    public int GetParameter(ZSTD_cParameter parameter)
    {
        EnsureNotDisposed();
        return compressor.NotNull().GetParameter(parameter);
    }

    public void LoadDictionary(byte[] dict)
    {
        EnsureNotDisposed();
        compressor.NotNull().LoadDictionary(dict);
    }

    ~CompressionStream() => Dispose(false);

    protected override void Dispose(bool disposing)
    {
        if (compressor == null)
        {
            base.Dispose(disposing);
            return;
        }

        try
        {
            if (disposing)
            {
                FlushInternal(ZSTD_EndDirective.ZSTD_e_end);
            }
        }
        finally
        {
            ReleaseUnmanagedResources();
        }
        base.Dispose(disposing);
    }

    private void ReleaseUnmanagedResources()
    {
        if (!preserveCompressor)
        {
            compressor.NotNull().Dispose();
        }
        compressor = null;

        if (outputBuffer != null)
        {
            ArrayPool<byte>.Shared.Return(outputBuffer);
        }

        if (!leaveOpen)
        {
            innerStream.Dispose();
        }
    }

    public override void Flush() => FlushInternal(ZSTD_EndDirective.ZSTD_e_flush);

    private void FlushInternal(ZSTD_EndDirective directive) =>
        WriteInternal(ReadOnlySpan<byte>.Empty, directive);

    public override void Write(byte[] buffer, int offset, int count) =>
        Write(new ReadOnlySpan<byte>(buffer, offset, count));

#if !LEGACY_DOTNET || NETSTANDARD2_1
    public override void Write(ReadOnlySpan<byte> buffer) =>
        WriteInternal(buffer, ZSTD_EndDirective.ZSTD_e_continue);
#else
    public void Write(ReadOnlySpan<byte> buffer) =>
        WriteInternal(buffer, ZSTD_EndDirective.ZSTD_e_continue);
#endif

    private void WriteInternal(ReadOnlySpan<byte> buffer, ZSTD_EndDirective directive)
    {
        EnsureNotDisposed();

        var input = new ZSTD_inBuffer_s { pos = 0, size = (nuint)buffer.Length };
        nuint remaining;
        do
        {
            output.pos = 0;
            remaining = CompressStream(ref input, buffer, directive);

            var written = (int)output.pos;
            if (written > 0)
            {
                innerStream.Write(outputBuffer, 0, written);
            }
        } while (
            directive == ZSTD_EndDirective.ZSTD_e_continue ? input.pos < input.size : remaining > 0
        );
    }

    internal unsafe nuint CompressStream(
        ref ZSTD_inBuffer_s input,
        ReadOnlySpan<byte> inputBuffer,
        ZSTD_EndDirective directive
    )
    {
        fixed (byte* inputBufferPtr = inputBuffer)
        fixed (byte* outputBufferPtr = outputBuffer)
        {
            input.src = inputBufferPtr;
            output.dst = outputBufferPtr;
            return compressor
                .NotNull()
                .CompressStream(ref input, ref output, directive)
                .EnsureZstdSuccess();
        }
    }

    public override bool CanRead => false;
    public override bool CanSeek => false;
    public override bool CanWrite => true;

    public override long Length => throw new NotSupportedException();

    public override long Position
    {
        get => throw new NotSupportedException();
        set => throw new NotSupportedException();
    }

    public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException();

    public override void SetLength(long value) => throw new NotSupportedException();

    public override int Read(byte[] buffer, int offset, int count) =>
        throw new NotSupportedException();

    private void EnsureNotDisposed()
    {
        if (compressor == null)
        {
            throw new ObjectDisposedException(nameof(CompressionStream));
        }
    }

    public void SetPledgedSrcSize(ulong pledgedSrcSize)
    {
        EnsureNotDisposed();
        compressor.NotNull().SetPledgedSrcSize(pledgedSrcSize);
    }
}