深入探索 LINQ

Where的原碼探索

前一章我們講到Where的使用方式，Where使用起來很直覺，就像用if else做判斷一樣，使用一個bool回傳型態的Lambda Expression就可以篩選我們所需要的資料，既然Where使用起來這麼單純，那我們來看看它的原始碼是不是也這麼單純吧。

原始碼分析

  • Source Code: https://github.com/dotnet/corefx/blob/master/src/System.Linq/src/System/Linq/Where.cs
  • Methods: Where有兩個Public Methods，我們先來看看其中一個:
public static IEnumerable<TSource> Where<TSource>(this IEnumerable<TSource> source, Func<TSource, int, bool> predicate)
{
    if (source == null)
    {
        throw Error.ArgumentNull(nameof(source));
    }
    
    if (predicate == null)
    {
        throw Error.ArgumentNull(nameof(predicate));
    }

    return WhereIterator(source, predicate);
}
  • 判斷source或是predicate是否為空，如果是空的就拋ArgumentNull的例外
  • 如果都是合法參數則回傳WhereIterator(source, predicate)

這裡我們依然從委派方法有index傳入參數的方法看起，可以看到跟前面介紹的LINQ方法在架構上幾乎沒有差別，所以關鍵依然在Iterator上，我們來看看WhereIterator的實作:

private static IEnumerable<TSource> WhereIterator<TSource>(IEnumerable<TSource> source, Func<TSource, int, bool> predicate)
{
    int index = -1;
    foreach (TSource element in source)
    {
        checked
        {
            index++;
        }

        if (predicate(element, index))
        {
            yield return element;
        }
    }
}
  • yield區塊中的yield return是傳回每個元素的值，而整個方法的回傳型別是IEnumerable
  • 後一個元素會比前個元素多加1
  • 用if接收predicate執行後的結果來判斷是否要將目前的元素回傳

Where這裡毫無意外的用了if的判斷來決定是否要回傳此元素，整段程式的差別也只有這裡，可見只要學會了Iterator，大部分的LINQ方法都能夠很快的理解。

接著我們要來看第二個Public Method了，這個是predicate沒有index傳入參數的方法:

public static IEnumerable<TSource> Where<TSource>(this IEnumerable<TSource> source, Func<TSource, bool> predicate)
{
    #region 判斷傳入參數合法性
    if (source == null)
    {
        throw Error.ArgumentNull(nameof(source));
    }

    if (predicate == null)
    {
        throw Error.ArgumentNull(nameof(predicate));
    }
    #endregion 判斷傳入參數合法性

    #region 依據source的型別決定Iterator
    if (source is Iterator<TSource> iterator)
    {
        return iterator.Where(predicate);
    }

    if (source is TSource[] array)
    {
        return array.Length == 0 ?
            (IEnumerable<TSource>)EmptyPartition<TSource>.Instance :
            new WhereArrayIterator<TSource>(array, predicate);
    }

    if (source is List<TSource> list)
    {
        return new WhereListIterator<TSource>(list, predicate);
    }

    return new WhereEnumerableIterator<TSource>(source, predicate);
    #endregion 依據source的型別決定Iterator
}
  • 判斷傳入參數source及predicate是否為空，空的話回傳ArgumentNull例外
  • 依據source的型別決定Iterator
    • 已經是Iterator的話就直接叫用Iterator的Where
    • 是Array的話回傳WhereArrayIterator
    • 是List的話回傳WhereListIterator
    • 只是IEnumerable的話就回傳WhereEnumerableIterator

每個Iterator的處理都大同小異，主要的差別在對於各個型別的處理而已，我們就挑WhereArrayIterator來說明內部運作吧:

  • MoveNext(): 這是整個Iterator最重要的Method，來看一下Where是怎麼實作MoveNext()的
public override bool MoveNext()
{
    int index = _state - 1;
    TSource[] source = _source;

    while (unchecked((uint)index < (uint)source.Length))
    {
        TSource item = source[index];
        index = _state++;
        if (_predicate(item))
        {
            _current = item;
            return true;
        }
    }

    Dispose();
    return false;
}
  • _state為陣列位置的基準，在GetEnumerator()執行後會被設為1，所以起始的index是_state - 1
  • 巡覽至陣列最尾端，每次都用predicate取得是否要回傳元素的判斷
  • 通過predicate後將current設為目前的元素，然後回傳true
  • 如果巡覽結束則Dispose()跟return false

這裡我們可以看到_state的用法跟前面介紹的SelectMany是不一樣的，SelectMany是用來決定是在第幾層的Enumerator和巡覽結束後要跳回哪層Enumerator，而Where則是完全當作index來使用。

  • GetCount: 取得集合的元素數量
public int GetCount(bool onlyIfCheap)
{
    if (onlyIfCheap)
    {
        return -1;
    }

    int count = 0;

    foreach (TSource item in _source)
    {
        if (_predicate(item))
        {
            checked
            {
                count++;
            }
        }
    }

    return count;
}
  • 如果通過predicate的驗證則count加1

之前我知道了Where是延遲執行後我就很好奇它的GetCount()是怎麼運作的，原來還是會全部執行後才取得元素數量。

測試案例賞析

  • Source Code: https://github.com/dotnet/corefx/blob/master/src/System.Linq/tests/WhereTests.cs

Where_SourceThrowsOnGetEnumerator

這個案例的source是ThrowsOnGetEnumerator()，會在第一次叫用GetEnumerator()時丟出InvalidOperationException例外。

protected class ThrowsOnGetEnumerator : TestEnumerator
{
    private int getEnumeratorCallCount;

    public override IEnumerator<int> GetEnumerator()
    {
        if (getEnumeratorCallCount++ == 0)
        {
            throw new InvalidOperationException();
        }

        return base.GetEnumerator();
    }
}

[Fact]
public void Where_SourceThrowsOnGetEnumerator()
{
    IEnumerable<int> source = new ThrowsOnGetEnumerator();
    Func<int, bool> truePredicate = (value) => true;

    var enumerator = source.Where(truePredicate).GetEnumerator();

    // Ensure the first MoveNext call throws an exception
    Assert.Throws<InvalidOperationException>(() => enumerator.MoveNext());

    // Ensure Current is set to the default value of type T
    int currentValue = enumerator.Current;
    Assert.Equal(default(int), currentValue);
    
    // Ensure subsequent MoveNext calls succeed
    Assert.True(enumerator.MoveNext());
    Assert.Equal(1, enumerator.Current);
}

要看懂這個測試案例要建立一個觀念: Where叫用的GetEnumerator()並不是source的GetEnumerator()，而是Where自己的GetEnumerator()。

知道這觀念後我們再看測試案例，這樣也就說得通為什麼不是在var enumerator = source.Where(truePredicate).GetEnumerator();拋出例外，而是在enumerator.MoveNext()丟出例外。

由於我們剛剛介紹Where的MoveNext()的source是Array的，所以他並沒有叫用GetEnumerator()，現在我們來看看WhereEnumerableIterator的MoveNext():

public override bool MoveNext()
{
    switch (_state)
    {
        case 1:
            _enumerator = _source.GetEnumerator();
            _state = 2;
            goto case 2;
        case 2:
            while (_enumerator.MoveNext())
            {
                TSource item = _enumerator.Current;
                if (_predicate(item))
                {
                    _current = item;
                    return true;
                }
            }

            Dispose();
            break;
    }

    return false;
}

這裡才是source.GetEnumerator()叫用的地方，所以我們第一次叫用source的GetEnumerator()是在MoveNext()而不是在Where叫用GetEnumerator()時。

結語

Where語法的主要運作原理在於巡覽時使用predicate的結果判斷是否將此元素加入結果集合中，簡單的用if判斷就可以做出如此實用的方法真的是太棒了。

參考

dotnet/corefx

← LINQ排序語法(OrderBy、OrderByDescending、ThenBy、ThenByDescending)的應用 →