Add PooledArrayBufferWriter<T>

PooledArrayBufferWriter<T> is an implementation of IBufferWriter<T> that uses shared array pools.

Author: DustinCampbell

src/Shared/Microsoft.AspNetCore.Razor.Utilities.Shared.Test/PooledArrayBufferWriterTests.cs

@@ -0,0 +1,34 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+using System;
+using System.Buffers;
+using Xunit;
+
+namespace Microsoft.AspNetCore.Razor.Utilities.Shared.Test;
+
+public class PooledArrayBufferWriterTests
+{
+    [Theory]
+    [InlineData(42)]
+    [InlineData(256)]
+    [InlineData(400)]
+    [InlineData(1024)]
+    [InlineData(8 * 1024)]
+    public void FillBufferWriter(int arraySize)
+    {
+        // Create an array filled with byte data.
+        Span<byte> span = new byte[arraySize];
+
+        for (var i = 0; i < arraySize; i++)
+        {
+            span[i] = (byte)(i % 255);
+        }
+
+        using var bufferWriter = new PooledArrayBufferWriter<byte>();
+
+        bufferWriter.Write(span);
+
+        Assert.True(span.SequenceEqual(bufferWriter.WrittenMemory.Span));
+    }
+}

src/Shared/Microsoft.AspNetCore.Razor.Utilities.Shared/PooledArrayBufferWriter`1.cs

@@ -0,0 +1,193 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+// Copied from https://github.com/dotnet/runtime
+
+using System.Diagnostics;
+using System.Diagnostics.CodeAnalysis;
+using System.Runtime.CompilerServices;
+using Microsoft.AspNetCore.Razor;
+
+namespace System.Buffers;
+
+/// <summary>
+///  <see cref="IBufferWriter{T}"/> that uses <see cref="ArrayPool{T}"/>.
+/// </summary>
+internal sealed class PooledArrayBufferWriter<T> : IBufferWriter<T>, IDisposable
+{
+    private T[]? _rentedBuffer;
+    private int _index;
+
+    private const int MinimumBufferSize = 256;
+
+    public PooledArrayBufferWriter()
+    {
+        _rentedBuffer = ArrayPool<T>.Shared.Rent(MinimumBufferSize);
+        _index = 0;
+    }
+
+    public PooledArrayBufferWriter(int initialCapacity)
+    {
+        ArgHelper.ThrowIfNegativeOrZero(initialCapacity);
+
+        _rentedBuffer = ArrayPool<T>.Shared.Rent(initialCapacity);
+        _index = 0;
+    }
+
+    public ReadOnlyMemory<T> WrittenMemory
+    {
+        get
+        {
+            CheckIfDisposed();
+
+            return _rentedBuffer.AsMemory(0, _index);
+        }
+    }
+
+    public int WrittenCount
+    {
+        get
+        {
+            CheckIfDisposed();
+
+            return _index;
+        }
+    }
+
+    public int Capacity
+    {
+        get
+        {
+            CheckIfDisposed();
+
+            return _rentedBuffer.Length;
+        }
+    }
+
+    public int FreeCapacity
+    {
+        get
+        {
+            CheckIfDisposed();
+
+            return _rentedBuffer.Length - _index;
+        }
+    }
+
+    public void Clear()
+    {
+        CheckIfDisposed();
+
+        ClearHelper();
+    }
+
+    private void ClearHelper()
+    {
+        Debug.Assert(_rentedBuffer != null);
+
+        _rentedBuffer.AsSpan(0, _index).Clear();
+        _index = 0;
+    }
+
+    // Returns the rented buffer back to the pool
+    public void Dispose()
+    {
+        if (_rentedBuffer == null)
+        {
+            return;
+        }
+
+        ClearHelper();
+        ArrayPool<T>.Shared.Return(_rentedBuffer);
+        _rentedBuffer = null;
+    }
+
+    [MemberNotNull(nameof(_rentedBuffer))]
+    private void CheckIfDisposed()
+    {
+        if (_rentedBuffer == null)
+        {
+            ThrowObjectDisposed();
+        }
+
+        [DoesNotReturn]
+        [MethodImpl(MethodImplOptions.NoInlining)]
+        static void ThrowObjectDisposed()
+        {
+            throw new ObjectDisposedException(nameof(ArrayBufferWriter<T>));
+        }
+    }
+
+    public void Advance(int count)
+    {
+        CheckIfDisposed();
+
+        ArgHelper.ThrowIfNegative(count);
+
+        if (_index > _rentedBuffer.Length - count)
+        {
+            ThrowCannotAdvance(_rentedBuffer.Length);
+        }
+
+        _index += count;
+
+        [DoesNotReturn]
+        [MethodImpl(MethodImplOptions.NoInlining)]
+        static void ThrowCannotAdvance(int capacity)
+        {
+            throw new InvalidOperationException(SR.FormatCannot_advance_past_end_of_the_buffer_which_has_a_size_of_0(capacity));
+        }
+    }
+
+    public Memory<T> GetMemory(int sizeHint = 0)
+    {
+        CheckIfDisposed();
+
+        CheckAndResizeBuffer(sizeHint);
+        return _rentedBuffer.AsMemory(_index);
+    }
+
+    public Span<T> GetSpan(int sizeHint = 0)
+    {
+        CheckIfDisposed();
+
+        CheckAndResizeBuffer(sizeHint);
+        return _rentedBuffer.AsSpan(_index);
+    }
+
+    private void CheckAndResizeBuffer(int sizeHint)
+    {
+        Debug.Assert(_rentedBuffer != null);
+
+        ArgHelper.ThrowIfNegative(sizeHint);
+
+        if (sizeHint == 0)
+        {
+            sizeHint = MinimumBufferSize;
+        }
+
+        var availableSpace = _rentedBuffer!.Length - _index;
+
+        if (sizeHint > availableSpace)
+        {
+            var growBy = Math.Max(sizeHint, _rentedBuffer.Length);
+
+            var newSize = checked(_rentedBuffer.Length + growBy);
+
+            var oldBuffer = _rentedBuffer;
+
+            _rentedBuffer = ArrayPool<T>.Shared.Rent(newSize);
+
+            Debug.Assert(oldBuffer.Length >= _index);
+            Debug.Assert(_rentedBuffer.Length >= _index);
+
+            var previousBuffer = oldBuffer.AsSpan(0, _index);
+            previousBuffer.CopyTo(_rentedBuffer);
+            previousBuffer.Clear();
+            ArrayPool<T>.Shared.Return(oldBuffer);
+        }
+
+        Debug.Assert(_rentedBuffer.Length - _index > 0);
+        Debug.Assert(_rentedBuffer.Length - _index >= sizeHint);
+    }
+}