GetEnumeratorの最適化

自作のコレクションクラスにIEnumerable<T>.GetEnumeratorを実装すればforeachで使えるようになります。このGetEnumeratorを実装するときに少し工夫するとメモリアロケーションを減らせます。

yield returnを使った実装

GetEnumeratorを実装するのに一番簡単なのはyield returnを使う方法です。

using System.Collections;
using System.Collections.Generic;

public class YieldReturnList : IEnumerable<int>
{
    private readonly List<int> list = new List<int>() { 1, 2, 3 };

    public IEnumerator<int> GetEnumerator()
    {
        foreach (var value in list)
        {
            yield return value;
        }
    }

    IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
}

このYieldReturnListをforeachで回すと56バイトのアロケーションが発生します。yield returnは内部的にクラスを自動生成して使用するのでその分のメモリが確保されてしまいます。(計測用のコードは記事末尾に載せています)

これからこのメモリアロケーションを減らしていきます。

IEnumeratorを自分で実装する

yield returnを使わない場合は自前でIEnumeratorを実装します。

using System.Collections;
using System.Collections.Generic;

public class EnumeratorList : IEnumerable<int>
{
    private readonly List<int> list = new List<int>() { 1, 2, 3 };

    public IEnumerator<int> GetEnumerator() => new Enumerator(list);
    IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();

    // IEnumeratorを実装する
    public class Enumerator : IEnumerator<int>
    {
        private readonly List<int> list;
        private int index;

        internal Enumerator(List<int> list)
        {
            this.list = list;
            index = -1;
        }

        public bool MoveNext()
        {
            index++;
            return (uint)index < (uint)list.Count;
        }

        public void Reset()
        {
            index = -1;
        }

        public int Current => list[index];
        object IEnumerator.Current => Current;

        public void Dispose() { }
    }
}

このEnumeratorクラスはyield returnで自動生成されるクラスよりは軽量なのでアロケーションは32バイトまで減ります。

structにする

Enumeratorをclassではなくstructにすればさらに減らせます。

ここで大切なのはGetEnumeratorの戻り値の型をIEnumerable<int>ではなくEnumeratorにすることです。戻り値がIEnumerable<int>のままでは、GetEnumeratorから返すときにボックス化が発生してstructにした意味が無くなってしまいます。

IEnumerable<int>で定義されているIEnumerable<int>を返すGetEnumeratorは明示的実装にします。そしてEnumeratorをそのまま返すGetEnumeratorを実装すれば、foreachでもこのGetEnumeratorが使われるようになります。

using System.Collections;
using System.Collections.Generic;

public class StructEnumeratorList : IEnumerable<int>
{
    private readonly List<int> list = new List<int>() { 1, 2, 3 };

    // IEnumerable<int>ではなくEnumeratorを返す
    public Enumerator GetEnumerator() => new Enumerator(list);
    IEnumerator<int> IEnumerable<int>.GetEnumerator() => GetEnumerator();
    IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();

    public struct Enumerator : IEnumerator<int>
    {
        private readonly List<int> list;
        private int index;

        internal Enumerator(List<int> list)
        {
            this.list = list;
            index = -1;
        }

        public bool MoveNext()
        {
            index++;
            return (uint)index < (uint)list.Count;
        }

        public void Reset()
        {
            index = -1;
        }

        public int Current => list[index];
        object IEnumerator.Current => Current;

        public void Dispose() { }
    }
}

これによってGetEnumeratorのメモリアロケーションが0になります。

List<T>.Enumeratorを使う

いままで例として使ってきたクラスでは値はC#標準のListに入れていました。こういう時は自分で Enumeratorを実装せずとも、ListのEnumeratorがそのまま使えます。

実はListのGetEnumeratorもIEnumerable<T>ではなくList<int>.Enumeratorというstruct型を返しています。これによってListのforeaceではメモリアロケーションが発生しないようになっています。このEnumeratorを使うには次のコードのようになります。

using System.Collections;
using System.Collections.Generic;

public class ListEnumeratorList : IEnumerable<int>
{
    private readonly List<int> list = new List<int>() { 1, 2, 3 };

    // list.GetEnumerator()を返す
    public List<int>.Enumerator GetEnumerator() => list.GetEnumerator();
    IEnumerator<int> IEnumerable<int>.GetEnumerator() => list.GetEnumerator();
    IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
}

パフォーマンス計測

最後に今までのクラスの計測コードとその結果を載せます。

using BenchmarkDotNet.Attributes;
using BenchmarkDotNet.Jobs;

[ShortRunJob(RuntimeMoniker.CoreRt60)]
[MemoryDiagnoser]
[MaxColumn, MinColumn, MeanColumn]
public class Foreach
{
    private readonly YieldReturnList yieldReturnList = new YieldReturnList();
    private readonly EnumeratorList enumeratorList = new EnumeratorList();
    private readonly StructEnumeratorList structEnumeratorList = new StructEnumeratorList();
    private readonly ListEnumeratorList listEnumeratorList = new ListEnumeratorList();

    [Benchmark]
    public int YieldReturn()
    {
        int sum = 0;
        foreach (var value in yieldReturnList)
        {
            sum += value;
        }
        return sum;
    }

    [Benchmark]
    public int Enumerator()
    {
        int sum = 0;
        foreach (var value in enumeratorList)
        {
            sum += value;
        }
        return sum;
    }

    [Benchmark]
    public int StructEnumerator()
    {
        int sum = 0;
        foreach (var value in structEnumeratorList)
        {
            sum += value;
        }
        return sum;
    }

    [Benchmark]
    public int ListEnumerator()
    {
        int sum = 0;
        foreach (var value in listEnumeratorList)
        {
            sum += value;
        }
        return sum;
    }
}