[C#] ArrayPool を使用して ToArray のアロケーションを1回にする方法

2024-08-08 (木)

一般的には List<T> を使用して要素を追加して行き、最後に ToArray() することで配列を作成します。
この方法では、最後の配列以外に作業用のヒープメモリ確保が発生してしまいます。
プールと構造体を使用することで、最後の完成した配列に対して1回のみメモリ確保する実装例です。

環境

.NET 8.0.7 (SDK 8.0.303)
C# 12.0
Visual Studio 2022 Version 17.10.5
Windows 11 Pro 23H2 22631.3880

前提

Dispose パターンを使用しているため、using ステートメントを使用してください。
.NET 8 前提なので、ビルドできないターゲットフレームワークの場合は、以下を参考に対応してください。
- Span<T> は NuGet 参照追加 https://www.nuget.org/packages/System.Memory
- MemberNotNull は NuGet 参照追加 https://www.nuget.org/packages/PolySharp/
- var clearArray = RuntimeHelpers.IsReferenceOrContainsReferences<T>() は true にしてください

結果

まず、通常の List<T> を使用した例です。

// サンプルデータを取得
IEnumerable<int> source1 = GetSampleData1();
int source2 = GetSampleData2();

// List に追加する
var list = new List<int>();
list.AddRange(source1);
list.Add(source2);

// 最後に ToArray() で配列に変換
var array = list.ToArray();

今回実装する ValueArrayBuilder<T> を使用した例です。使用方法は似ています。

IEnumerable<int> source1 = GetSampleData1();
int source2 = GetSampleData2();

// Dispose() が必要のため using ステートメントを使用してください
using var builder = new ValueArrayBuilder<int>();
builder.AddRange(source1);
builder.Add(source2);

var array = builder.ToArray();

初期バッファに stackalloc を使用した例です。容量が不足する場合は、自動的に ArrayPool から確保します。

IEnumerable<int> source1 = GetSampleData1();
int source2 = GetSampleData2();

// Dispose() が必要のため using ステートメントを使用してください
using var builder = new ValueArrayBuilder<int>(stackalloc int[128]);
builder.AddRange(source1);
builder.Add(source2);

var array = builder.ToArray();

実装

public ref struct ValueArrayBuilder<T>
{
    private Span<T> _span;
    private T[]? _arrayFromPool;
    private int _pos;

    public ValueArrayBuilder(int minimumLength = 256)
    {
        _span = _arrayFromPool = ArrayPool<T>.Shared.Rent(minimumLength);
        _pos = 0;
    }

    public ValueArrayBuilder(Span<T> initialSpan)
    {
        _span = initialSpan;
        _arrayFromPool = null;
        _pos = 0;
    }

    public ref T this[int index] => ref _span[index];
    public int Length => _pos;

    public void Dispose()
    {
        ReturnArray();
    }

    public void Clear()
    {
        _span.Clear();
        _pos = 0;
    }

    public void Add(T item)
    {
        if ((uint)_pos >= (uint)_span.Length)
        {
            Grow();
        }

        _span[_pos] = item;
        _pos++;
    }

    public void AddRange(scoped ReadOnlySpan<T> source)
    {
        var pos = _pos;
        var span = _span;

        if (source.Length == 1 && (uint)pos < (uint)span.Length)
        {
            span[pos] = source[0];
            _pos = pos + 1;
        }
        else
        {
            AddSpan(source);
        }
    }

    private void AddSpan(scoped ReadOnlySpan<T> source)
    {
        if ((uint)(_pos + source.Length) > (uint)_span.Length)
        {
            Grow(_span.Length - _pos + source.Length);
        }

        source.CopyTo(_span.Slice(_pos));
        _pos += source.Length;
    }

    public void AddRange(ICollection<T> source)
    {
        if (_arrayFromPool is null || (uint)(_pos + source.Count) > (uint)_span.Length)
        {
            Grow(_span.Length - _pos + source.Count);
        }

        source.CopyTo(_arrayFromPool, _pos);
        _pos += source.Count;
    }

    public void AddRange(IEnumerable<T> source)
    {
        if (source is ICollection<T> collection)
        {
            AddRange(collection);
            return;
        }

        foreach (var item in source)
        {
            Add(item);
        }
    }

    public T[] ToArray()
    {
        return AsSpan().ToArray();
    }

    public ReadOnlySpan<T> AsSpan()
    {
        return _span.Slice(0, _pos);
    }

    private void ReturnArray()
    {
        if (_arrayFromPool is { } toReturn)
        {
            _arrayFromPool = null;

            var clearArray = RuntimeHelpers.IsReferenceOrContainsReferences<T>();
            ArrayPool<T>.Shared.Return(toReturn, clearArray);
        }
    }

    [MemberNotNull(nameof(_arrayFromPool))]
    private void Grow(int additionalMinimumLength = 1)
    {
        var nextCapacity = Math.Max(_span.Length != 0 ? _span.Length * 2 : 256, _span.Length + additionalMinimumLength);
        var nextArray = ArrayPool<T>.Shared.Rent(nextCapacity);

        _span.CopyTo(nextArray);
        ReturnArray();
        _span = _arrayFromPool = nextArray;
    }
}