LINQの内部で行われている最適化

## 導入

LINQが最適化のため使用しているテクニックをSelectの実装に沿って紹介したいと思います。
ここで紹介する最適化は他の似た系統の言語のイテレータアダプタ (mapやfilter等) にも適用できるはずです。
見たところKotlinのSequenceはここで紹介する最適化を[していない](https://github.com/JetBrains/kotlin/blob/master/libraries/stdlib/src/kotlin/collections/Sequences.kt#L205L221)ので貢献するチャンスかもしれません。

## Selectメソッド

Selectのソースコード: https://github.com/dotnet/runtime/blob/main/src/libraries/System.Linq/src/System/Linq/Select.cs

まずSelectメソッドは次のように実装されています。

```c#
public static IEnumerable<TResult> Select<TSource, TResult>(
    this IEnumerable<TSource> source, Func<TSource, TResult> selector)
{
    if (source == null)
    {
        ThrowHelper.ThrowArgumentNullException(ExceptionArgument.source);
    }

    if (selector == null)
    {
        ThrowHelper.ThrowArgumentNullException(ExceptionArgument.selector);
    }

    if (source is Iterator<TSource> iterator)
    {
        return iterator.Select(selector);
    }

    if (source is IList<TSource> ilist)
    {
        if (source is TSource[] array)
        {
            return array.Length == 0 ?
                Empty<TResult>() :
                new SelectArrayIterator<TSource, TResult>(array, selector);
        }

        if (source is List<TSource> list)
        {
            return new SelectListIterator<TSource, TResult>(list, selector);
        }

        return new SelectIListIterator<TSource, TResult>(ilist, selector);
    }

    if (source is IPartition<TSource> partition)
    {
        IEnumerable<TResult>? result = null;
        CreateSelectIPartitionIterator(selector, partition, ref result);
        if (result != null)
        {
            return result;
        }
    }

    return new SelectEnumerableIterator<TSource, TResult>(source, selector);
}
```

Nullチェックをした後に、sourceがよく使われる特定の型だった場合には専用のIEnumerableを返し、そうでなければ汎用的な実装であるSelectEnumerableIteratorを返すという形になっています。
最初の分岐にあるIteratorという型はSelectEnumerableIteratorなどLINQの返り値が継承している抽象クラスです。
つまり `list.Where(...).Select(...)` のようにメソッドチェインしている場合にあたる分岐になります。
`iterator.Select()` はIteratorのメソッドで (LINQのSelectと違って静的拡張メソッドでない) 、Selectの実装を各Iteratorに任せることによって最適化された実装を返しています。

実際作成されるIteratorクラスの確認の前にIEnumerableとIEnumeratorの復習を軽くしておきます。

## IEnumerableとIEnumeratorの復習

`IEnumerable<T>`はループ可能なことを表す`List<T>`や`HashSet<T>`が実装するインターフェイスです。
`IEnumerator<T> GetEnumerator()`というメソッドだけを持ちループ変数の保持など実際のループ処理は`IEnumerator<T>`が担当します。
詳しくは公式ドキュメントを参照してください
: [IEnumerable](https://docs.microsoft.com/ja-jp/dotnet/api/system.collections.generic.ienumerable-1), [IEnumerator](https://docs.microsoft.com/ja-jp/dotnet/api/system.collections.generic.ienumerator-1)

## Iteratorクラス

Iteratorのソースコード: https://github.com/dotnet/runtime/blob/main/src/libraries/System.Linq/src/System/Linq/Iterator.cs

Iteratorクラスのメンバとコンストラクタは次のようになっています。

```c#
internal abstract class Iterator<TSource> : IEnumerable<TSource>, IEnumerator<TSource>
{
    private readonly int _threadId;
    internal int _state;
    internal TSource _current = default!;

    protected Iterator()
    {
        _threadId = Environment.CurrentManagedThreadId;
    }
}
```

_currentはIEnumeratorとして現在保持している値です。
_threadIdと_stateについては後で説明します。

注目すべきはIEnumerableとIEnumeratorの両方を実装していることです。
これがどういうことかを知るためにIEnumerableのメソッドGetEnumeratorの実装を見てみましょう。

```c#
public IEnumerator<TSource> GetEnumerator()
{
    Iterator<TSource> enumerator = _state == 0 && _threadId == Environment.CurrentManagedThreadId ? this : Clone();
    enumerator._state = 1;
    return enumerator;
}
```

基本的には自分自身をIEnumeratorとして返すという実装になっています。
もしIEnumerableとIEnumeratorが別の型に分かれている場合、`list.Select(...).ToList()` としたときSelectでIEnumerableを作成して `.ToList()` で内部的にGetEnumeratorを呼び出してIEnumeratorを作成するという二度のインスタンス化が行われます。
IEnumerableとIEnumeratorを兼ねておくことでインスタンス化が一度で済むことになります。

ただしそれを素朴に実装すると問題があるので分岐で少し処理をすることになっています。
次のようなコードを考えます。

```c#
Iterator<int> enumerable;
// enumerableの初期化...

var enumerator = enumerable.GetEnumerator();
enumerator.MoveNext();
var enumerator2 = enumerable.GetEnumerator();
enumerator2.MoveNext();
Console.WriteLine(enumerator2.Current); // このCurrentは一番目の要素を指すべき
```

もし単純に`enumerable.GetEnumerator()`が常に自分自身を返すなら最終行の`enumerator2.Current`は二番目の要素になってしまいます。
二度目以降の`GetEnumerator()`は新たなインスタンスを返さなければなりません。
ここで_stateが活用されます。
IteratorのGetEnumeratorの実装ではまず一度目の呼び出しでは_stateを0から1にしてGetEnumeratorが呼び出されたとマークしてから自分自身を返します。
そして二度目のGetEnumeratorの呼び出しでは_stateが0でないのでIteratorの抽象メソッドCloneを呼び出して作成した自分自身のクローンを返します。
CloneではIEnumerableで使用するフィールド (Selectではsourceなど) はそのままで、IEnumeratorで使用するフィールド (ループ変数など) は初期化するよう実装することにしておけば二度目以降のGetEnumeratorの呼び出しでは新しいIEnumeratorが作成されることになりうまくいきます。
_stateは0が作成後GetEnumerator呼び出し前、1以上がGetEnumerator呼び出し後Dispose呼び出し前、-1がDispose呼び出し後を表します。

フィールドの_threadIdについてですが、二つのスレッドが同時にGetEnumeratorを呼び出したとき、両方のスレッドが `_state == 0` だと判定してしまう恐れがあるので、enumerableを作成したスレッド自身以外が呼び出したときは常にクローンを返すことで意図せず同じIEnumeratorを返すことを防いでいます。

では実際にSelectではどのようにIteratorを継承しているか見ていきましょう。

## SelectでのIteratorの実装

Selectのソースコード: https://github.com/dotnet/runtime/blob/main/src/libraries/System.Linq/src/System/Linq/Select.cs

まずは汎用的に使われるSelectEnumerableIteratorを見ていきましょう。
フィールドとコンストラクタは次のようになっています。

```c#
private sealed partial class SelectEnumerableIterator<TSource, TResult> : Iterator<TResult>
{
    private readonly IEnumerable<TSource> _source;
    private readonly Func<TSource, TResult> _selector;
    private IEnumerator<TSource>? _enumerator;

    public SelectEnumerableIterator(IEnumerable<TSource> source, Func<TSource, TResult> selector)
    {
        _source = source;
        _selector = selector;
    }
}
```

_sourceと_selectorに加えて、_sourceのenumeratorを保持する_enumeratorフィールドがあります。
Iterator.Select以外は普通に実装されています。

CloneはIEnumerable部分だけ保持するということだったので次のようになっています。

```c#
public override Iterator<TResult> Clone() =>
    new SelectEnumerableIterator<TSource, TResult>(_source, _selector);
```

MoveNextは次のようになっています。

```c#
public override bool MoveNext()
{
    switch (_state)
    {
        case 1:
            _enumerator = _source.GetEnumerator();
            _state = 2;
            goto case 2;
        case 2:
            Debug.Assert(_enumerator != null);
            if (_enumerator.MoveNext())
            {
                _current = _selector(_enumerator.Current);
                return true;
            }

            Dispose();
            break;
    }

    return false;
}
```

まず初回の呼び出しで_sourceのEnumeratorを取得し、`_state = 2` とすることで既にEnumeratorを取得していることをマークします。

`Iterator.Select` はオーバーライドされています。
`list.Select(hoge).Select(fuga)` は `list.Select(x => fuga(hoge(x)))` と同等であるということから、SelectEnumerableIteratorがネストすることを回避し、MoveNextの呼び出しを浅くしています。

```c#
public override IEnumerable<TResult2> Select<TResult2>(Func<TResult, TResult2> selector) =>
    new SelectEnumerableIterator<TSource, TResult2>(_source, CombineSelectors(_selector, selector));
```

CombineSelectorsは次のように定義されます。
ソース: https://github.com/dotnet/runtime/blob/main/src/libraries/System.Linq/src/System/Linq/Utilities.cs#L70L71

```c#
public static Func<TSource, TResult> CombineSelectors<TSource, TMiddle, TResult>(Func<TSource, TMiddle> selector1, Func<TMiddle, TResult> selector2) =>
    x => selector2(selector1(x));
```

### 配列やList特有の実装

sourceがListの場合のSelectの実装であるSelectListIteratorのフィールドは次のようになっています。

```c#
private sealed partial class SelectListIterator<TSource, TResult> : Iterator<TResult>
{
    private readonly List<TSource> _source;
    private readonly Func<TSource, TResult> _selector;
    private List<TSource>.Enumerator _enumerator;
}
```

ListのEnumeratorは構造体であり、これを汎用的な実装と違ってIEnumeratorでなく具体的な型として保持することでボクシングを回避しています。

sourceが配列の場合のSelectArrayIteratorでは_stateの1以上の部分を利用してインデックスを保持しています。

```c#
public override bool MoveNext()
{
    if (_state < 1 | _state == _source.Length + 1)
    {
        Dispose();
        return false;
    }

    int index = _state++ - 1;
    _current = _selector(_source[index]);
    return true;
}
```

## あとがき

後半力尽きた感がありますが、なんとか書ききることができました。
最適化の話は面白いですね。
実際にIEnumerableを実装するときはyieldで事足りてしまいますが…