Improve the efficiency of decompressing to IBufferWriter<byte>

Snappier/Internal/IsExternalInit.cs

@@ -0,0 +1,15 @@
+#if NETSTANDARD2_0
+
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+#pragma warning disable IDE0079
+#pragma warning disable S3903
+
+namespace System.Runtime.CompilerServices;
+
+internal static class IsExternalInit
+{
+}
+
+#endif

Snappier/Internal/SnappyDecompressor.cs

@@ -1,6 +1,7 @@
 using System;
 using System.Buffers;
 using System.Diagnostics;
+using System.Diagnostics.CodeAnalysis;
 using System.IO;
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
@@ -49,6 +50,11 @@ private struct ScratchBuffer
         /// </remarks>
         public void Decompress(ReadOnlySpan<byte> input)
         {
+            if (AllDataDecompressed)
+            {
+                ThrowHelper.ThrowInvalidOperationException("All data has been decompressed");
+            }
+
             if (!ExpectedLength.HasValue)
             {
                 OperationStatus status = TryReadUncompressedLength(input, out int bytesConsumed);
@@ -82,6 +88,15 @@ public void Decompress(ReadOnlySpan<byte> input)
                     DecompressAllTags(input);
                 }
             }
+
+            if (BufferWriter is not null && AllDataDecompressed)
+            {
+                // Advance the buffer writer to the end of the data
+                BufferWriter.Advance(_lookbackPosition);
+
+                // Release the lookback buffer
+                _lookbackBuffer = default;
+            }
         }
 
         public void Reset()
@@ -92,6 +107,12 @@ public void Reset()
             _lookbackPosition = 0;
             _readPosition = 0;
             ExpectedLength = null;
+
+            if (BufferWriter is not null)
+            {
+                // Don't reuse the lookback buffer when it came from a BufferWriter
+                _lookbackBuffer = default;
+            }
         }
 
         private OperationStatus TryReadUncompressedLength(ReadOnlySpan<byte> input, out int bytesConsumed)
@@ -483,6 +504,11 @@ private uint RefillTag(ref byte input, ref byte inputEnd)
 
         #region Loopback Writer
 
+        /// <summary>
+        /// Buffer writer for the output data. Incompatible with <see cref="ExtractData"/> and <see cref="Read"/>.
+        /// </summary>
+        public IBufferWriter<byte>? BufferWriter { get; init; }
+
         private byte[]? _lookbackBufferArray;
         private Memory<byte> _lookbackBuffer;
         private int _lookbackPosition = 0;
@@ -503,8 +529,15 @@ private int? ExpectedLength
                         ArrayPool<byte>.Shared.Return(_lookbackBufferArray);
                     }
 
-                    _lookbackBufferArray = ArrayPool<byte>.Shared.Rent(value.GetValueOrDefault());
-                    _lookbackBuffer = _lookbackBufferArray.AsMemory(0, _lookbackBufferArray.Length);
+                    if (BufferWriter is not null)
+                    {
+                        _lookbackBuffer = BufferWriter.GetMemory(value.GetValueOrDefault());
+                    }
+                    else
+                    {
+                        _lookbackBufferArray = ArrayPool<byte>.Shared.Rent(value.GetValueOrDefault());
+                        _lookbackBuffer = _lookbackBufferArray.AsMemory(0, _lookbackBufferArray.Length);
+                    }
                 }
             }
         }
@@ -587,6 +620,11 @@ private static void AppendFromSelf(ref byte op, ref byte buffer, ref byte buffer
 
         public int Read(Span<byte> destination)
         {
+            if (BufferWriter is not null)
+            {
+                ThrowCannotUseWithBufferWriter(nameof(Read));
+            }
+
             var bytesToRead = Math.Min(destination.Length, UnreadBytes);
             if (bytesToRead <= 0)
             {
@@ -609,6 +647,11 @@ public int Read(Span<byte> destination)
         /// </remarks>
         public IMemoryOwner<byte> ExtractData()
         {
+            if (BufferWriter is not null)
+            {
+                ThrowCannotUseWithBufferWriter(nameof(ExtractData));
+            }
+
             byte[]? data = _lookbackBufferArray;
             if (!ExpectedLength.HasValue)
             {
@@ -637,6 +680,15 @@ public IMemoryOwner<byte> ExtractData()
             return returnBuffer;
         }
 
+        [DoesNotReturn]
+        [MethodImpl(MethodImplOptions.NoInlining)]
+        private static void ThrowCannotUseWithBufferWriter(string method)
+        {
+            // This is intentionally not inlined to keep the size of Read and ExtractData smaller,
+            // making it more likely they may be inlined.
+            ThrowHelper.ThrowNotSupportedException($"Cannot use {method} when using a BufferWriter.");
+        }
+
         #endregion
 
         #region Test Helpers

Snappier/Internal/ThrowHelper.cs

@@ -41,8 +41,8 @@ public static void ThrowInvalidOperationException(string? message) =>
             throw new InvalidOperationException(message);
 
         [DoesNotReturn]
-        public static void ThrowNotSupportedException() =>
-            throw new NotSupportedException();
+        public static void ThrowNotSupportedException(string? message = null) =>
+            throw new NotSupportedException(message);
 
         [DoesNotReturn]
         public static void ThrowObjectDisposedException(string? objectName) =>

Snappier/Snappy.cs

@@ -52,6 +52,8 @@ public static int Compress(ReadOnlySpan<byte> input, Span<byte> output)
         /// </remarks>
         public static void Compress(ReadOnlySequence<byte> input, IBufferWriter<byte> output)
         {
+            ThrowHelper.ThrowIfNull(output);
+
             using var compressor = new SnappyCompressor();
 
             compressor.Compress(input, output);
@@ -149,9 +151,22 @@ public static int Decompress(ReadOnlySpan<byte> input, Span<byte> output)
         /// <exception cref="InvalidDataException">Invalid Snappy block.</exception>
         public static void Decompress(ReadOnlySequence<byte> input, IBufferWriter<byte> output)
         {
-            using IMemoryOwner<byte> buffer = DecompressToMemory(input);
+            ThrowHelper.ThrowIfNull(output);
+
+            using var decompressor = new SnappyDecompressor()
+            {
+                BufferWriter = output
+            };
+
+            foreach (ReadOnlyMemory<byte> segment in input)
+            {
+                decompressor.Decompress(segment.Span);
+            }
 
-            output.Write(buffer.Memory.Span);
+            if (!decompressor.AllDataDecompressed)
+            {
+                ThrowHelper.ThrowInvalidDataException("Incomplete Snappy block.");
+            }
         }
 
         /// <summary>