少しでも簡単にバッファ配列の確保をしたい

はじめに

言いたいことはタイトルの通りです。
一時的にバッファ領域を用意してデータの受け取りなどにしたい場面が多くありますが，その度に馬鹿正直に配列を確保していたのではアロケーションのコスト+GCのコストが積み重なりパフォーマンスを悪化させます。
そこでc#ではパフォーマンスを悪化させずにバッファ領域の確保をするための方法がいくつかあります。

配列を使い回そう

最初だけヒープアロケーションを受け入れて，それをGCに回収させずに使いまわせば長期的にはアロケーションのコストが下がっていくと考えられます。
これを容易に行うためにC#にはSystem.Buffers.ArrayPool<T>クラスが用意されており，いい感じの長さの配列を使いまわせるようになっています。
このクラスを使用する際には，使う際にRentし，使い終わった配列をReturnすることで再利用できるようにします。

// int length
int[] array;
try
{
    array = ArrayPool<T>.Shared.Rent(length).AsSpan()[..length];
    // 何かする
}
finally
{
    ArrayPool<T>.Shared.Return(array);
}

このように，try-finallyに入れてあげることで確実に返却できるようになりますが，ネストが深くなってちょっと嫌です。

スタックに乗せよう

ヒープに入れずにスタックに乗せてしまえばアロケーションとかGCとかそんな話は関係なく爆速です。

Span<int> buffer = (stackalloc int[length]);

ところが，スタックというのはそんなに広くはない(1 MBとかそのくらい)ので，1要素が大きめで長い領域を確保しようとするとあっという間にスタックオーバーフローします。
そこで，長さに制限をかけてスタックに乗せるのが危険な場合はヒープに入れるという戦略が一般的です。

// 4 KB以内なら大丈夫なはず
Span<int> buffer = length <= 1024 ? stackalloc int[length] : new int[length];

合わせ技にしたいですよね

上記の方法を合わせて，短い時はスタック上に，長い時は同じ配列の使い回しをするともっと良くなるような気がします。
ところが，これを素直に書こうとすると見た目があまりよくないコード(丁寧な言い回し)ができます。

int[]? pooled = null;
try
{
    Span<int> = length <= 1024
        ? stackalloc int[length]
        : ArrayPool<T>.Shared.Rent(length).AsSpan()[..length];
}
finally
{
    if (pooled is not null)
            ArrayPool<T>.Shared.Return(pooled);
}

やりたいことはバッファ領域の確保であって本質的なコードは別なのにtry-finallyがありさらにfinallyの中で分岐が入るのでとても鬱陶しいことになります。

というわけで，前置きが長くなりましたがこれを簡単にしたいというのが今回のテーマになります。

方針

合わせ技が混沌としてしまったのは，「後片付け」が必要なためです。
ところで，C#には後片付けを勝手にやってくれるusingという便利なものがあります。
そこで，配列の返却をDisposeに押し付けてしまえばもっとシンプルに利用できるのではないかということでRent, Returnを隠す仕組みを作ってみました。

実装例

using System.Buffers;
using System.Runtime.CompilerServices;

namespace Qiita;

public ref struct PooledBuffer<T>
{
    private readonly int _length;
    private T[]? _buffer = null;

    public PooledBuffer(int length)
    {
        ArgumentOutOfRangeException.ThrowIfNegativeOrZero(length, nameof(length));
        this._length = length;
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public Span<T> GetSpan()
    {
        this._buffer ??= ArrayPool<T>.Shared.Rent(this._length);
        return this._buffer.AsSpan()[..this._length];
    }

    public void Dispose()
    {
        var buffer = this._buffer;
        if (buffer is null) return;
        this._buffer = null;
        ArrayPool<T>.Shared.Return(buffer);
    }
}

使い方

using var pooled = new PooledBuffer<int>(length);  // ここではまだRentしない
Span<int> buffer = length <= 1024
    ? stackalloc int[length]  // 十分に短い時はスタック上に確保
    ? pooled.GetSpan();  // GetSpanが呼ばれて初めてRentする

// スコープを抜けるときに，もしRentしていたらReturnする

解説

ref struct

ref構造体と呼ばれるもので，スタック上に存在することが強制されます。
どうやってもヒープに存在できませんし，何かの間違いでボクシングが起きることもあり得ません。
そのため，newして余計なアロケーションが発生することはありません。
そもそもアロケーションを減らしたくてこんなことをしているのに余計なコストが増えては本末転倒です。

制約について

ヒープに逃がせないという縛りがあるため，ref構造体はイテレータや非同期メソッドでは使用が制限されます。
しかし，そもそもそのような場合にはstackallocも使用できないため，素直にArrayPool<T>だけを使うパターンを採用することになります。

コンストラクタ

必要な長さを覚えておきます。
もしかしたらスタック上に乗るかもしれず，その場合はRentしても無駄になってしまうので長さを覚えておくだけにしておきます。

GetSpan

このメソッドが呼ばれたときにはじめてRentします(既に配列を持っている場合にはスキップします)。
ArrayPool<T>.Rentは指定した長さ「以上」の長さの配列を返すことが保証されますが，ぴったりの長さであるとは限りません。
そこで，指定された長さに切りだしてあげる必要があります。

Dispose

もしnullでない配列を持っているなら返却します。
nullなら返却は不要なので何もしません。
ref構造体はIDisposableでなくてもpublic void Dispose()を持っていればパターンマッチングでusingに使えます。

さいごに

使わない場合でもインスタンスが作られるのでちょっと気持ち悪さはありますが，それなりに使いやすくなったと思っています。
基本的な作りだけなので改善点等あるかと思いますが，皆様のご参考になれば幸いです。