懒人Ethan

C# LINQ源码分析之SelectMany

概要

在开发过程中，LINQ代码库中的SelectMany方法作为嵌套循环的语法糖，经常被使用。为了更好的了解该方法，我们从源码的角度对其进行分析，以了解其内部工作方式。

SelectMany方法介绍

SelectMany方法的基本功能是将序列中的每个元素投影到IEnumerable中，并将生成的序列平铺到一个新的序列中。LINQ代码库提供了4个SelectMany的重载方法如下：

方法名称	基本介绍
SelectMany(IEnumerable, Func, Func)	将序列中的每个元素投影到IEnumerable上，将结果平铺成一个序列，在平铺过程中，调用结果选择器函数。在投影过程中，元素索引作为参数被使用。
SelectMany(IEnumerable, Func, Func)	将序列中的每个元素投影到IEnumerable上，将结果平铺成一个序列，在平铺过程中，调用结果选择器函数。
SelectMany(IEnumerable, Func)	序列中的每个元素投影到IEnumerable上，将结果平铺成一个序列。
SelectMany(IEnumerable, Func)	序列中的每个元素投影到IEnumerable上，将结果平铺成一个序列。在投影过程中，元素索引作为参数被使用。

前两个重载在方法和后两个重载方法的区别是前两个方法在平铺过程中提供了结果选择器函数，使得平铺操作更加灵活。

第一个和第四个方法与第二第三两个方法的区别是他们在投影过程中增加了索引参数。

因为索引参数在开发过程中，使用的并不多，所以本文主要分析第二和第三两个重载方法的源码。第一第四两个方法实现方式类似，不再赘述。

SelectMany方法关键代码结构

与其他LINQ方法类似，SelectMany方法是实现基础也是迭代器和yield。迭代器继承了基类迭代器Iterator，从而获得了多线程支持和嵌套循环的支持。如果要了解更过的迭代器基础内容，请参考我的博文 C# LINQ源码分析之迭代器

编号	类名或方法名	基本功能
1.	SelectManySingleSelectorIterator迭代器	在逐个迭代过程中将序列元素投影到IEnumerable中，并将所有的投影结果平铺到一个序列中
2.	SelectIterator 方法	该方法的基本功能和SelectManySingleSelectorIterator迭代器相同，根据参数不同，有多个重载方法以支持平铺过程中的结果选择器函数。

SelectMany关键源码分析

因此，本文主要分析SelectManySingleSelectorIterator迭代器的源码实现。以及通过yield关键字的实现源码。

SelectMany源码分析（无结果选择器）

Select是IEnumerable接口的一个扩展方法，代码如下：

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

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

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

该方法有两个参数，第一个是集合的扩展方法参数，TSource是集合内元素的类型，第二个参数是一个委托类型，它封装了一个TSource到IEnumerable投影方法。
如果调用SelectMany方法的数据源为空，抛出异常。
如果Select方法的第二个参数，投影方法为空，抛出异常。
返回SelectManySingleSelectorIterator实例，迭代数据序列和投影方法是SelectManySingleSelectorIterator的构造方法参数。

SelectManySingleSelectorIterator源码分析

  private sealed partial class SelectManySingleSelectorIterator<TSource, TResult> : Iterator<TResult>
{
    private readonly IEnumerable<TSource> _source;
    private readonly Func<TSource, IEnumerable<TResult>> _selector;
    private IEnumerator<TSource>? _sourceEnumerator;
    private IEnumerator<TResult>? _subEnumerator;

    internal SelectManySingleSelectorIterator(IEnumerable<TSource> source, Func<TSource, IEnumerable<TResult>> selector)
    {
        Debug.Assert(source != null);
        Debug.Assert(selector != null);

        _source = source;
        _selector = selector;
    }

    public override Iterator<TResult> Clone()
    {
        return new SelectManySingleSelectorIterator<TSource, TResult>(_source, _selector);
    }

    public override void Dispose()
    {
        if (_subEnumerator != null)
        {
            _subEnumerator.Dispose();
            _subEnumerator = null;
        }

        if (_sourceEnumerator != null)
        {
            _sourceEnumerator.Dispose();
            _sourceEnumerator = null;
        }

        base.Dispose();
    }

    public override bool MoveNext()
    {
        switch (_state)
        {
            case 1:
                // Retrieve the source enumerator.
                _sourceEnumerator = _source.GetEnumerator();
                _state = 2;
                goto case 2;
            case 2:
                // Take the next element from the source enumerator.
                Debug.Assert(_sourceEnumerator != null);
                if (!_sourceEnumerator.MoveNext())
                {
                    break;
                }

                TSource element = _sourceEnumerator.Current;

                // Project it into a sub-collection and get its enumerator.
                _subEnumerator = _selector(element).GetEnumerator();
                _state = 3;
                goto case 3;
            case 3:
                // Take the next element from the sub-collection and yield.
                Debug.Assert(_subEnumerator != null);
                if (!_subEnumerator.MoveNext())
                {
                    _subEnumerator.Dispose();
                    _subEnumerator = null;
                    _state = 2;
                    goto case 2;
                }

                _current = _subEnumerator.Current;
                return true;
        }

        Dispose();
        return false;
    }
}

SelectManySingleSelectorIterator类继承自Iterator类，从而获得多线程安全的支持和多重循环嵌套访问的支持。
该类包含私有成员_source，用于存储要迭代的序列，在构造函数中初始化。
该类包含封装投影方法_selector，其在构造函数中被初始化。
该类的克隆方法返回SelectManySingleSelectorIterator对象，在多线程访问同一个SelectManySingleSelectorIterator对象时候被调用。
该类覆写了基类的MoveNext方法，从而通过迭代完成其投影操作。
（1）其实状态_state为0，在通过foreach或其他方式调用基类中定义的GetEnumertator方法后，被修改为1。
（2）case 1:下，source序列调用自己的GetEnumerator方法，获取迭代器_sourceEnumerator，将_state改为2，进入case 2。
（3）case 2:下
(a) 如果source序列的已经迭代到最后一个元素，调用Dispose方法，释放当前SelectMany迭代器。
(b) 获取当前souce序列当前迭代到的元素，执行投影方法，将该元素投影到IEnumerable中。
© 获取投影结果的迭代器_subEnumerator，以准备执行平铺操作。
(d) 将_state改为3，进入case 3
（4）case 3:下
(a) 调用_subEnumerator迭代器的MoveNext方法，检查是否还有下一个元素，如果没有，将_state改为2，回到case 2。该操作相当于我们的嵌套foreach循环中，内部foreach执行完了，回到外部foreach中。
(b)获取当前元数，并返回true。
如果整个source序列迭代完成，调用Dispose方法，返回false。

SelectMany源码分析（含结果选择器）

public static IEnumerable<TResult> SelectMany<TSource, TCollection, TResult>(this IEnumerable<TSource> source, Func<TSource, IEnumerable<TCollection>> collectionSelector, Func<TSource, TCollection, TResult> resultSelector)
{
    if (source == null)
    {
        ThrowHelper.ThrowArgumentNullException(ExceptionArgument.source);
    }

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

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

    return SelectManyIterator(source, collectionSelector, resultSelector);
}

该方法有三个参数，第一个是集合的扩展方法参数，TSource是集合内元素的类型，第二个参数是一个投影选择器函数，它封装了一个TSource到IEnumerable投影方法，第三个参数是平铺过程中的结果选择器函数。
如果调用SelectMany方法的数据源source为空，抛出异常。
如果SelectMany方法的第二个参数，投影方法为空，抛出异常。
如果SelectMany方法的第三个参数，平铺过程的结果选择器为空，抛出异常。
调用SelectManyIterator方法，该方法包含三个参数，第一个是迭代的数据源，第二个是投影选择器，第三个是平铺过程的结果选择器。

SelectManyIterator源码分析（含结果选择器）

 private static IEnumerable<TResult> SelectManyIterator<TSource, TCollection, TResult>(IEnumerable<TSource> source, Func<TSource, IEnumerable<TCollection>> collectionSelector, Func<TSource, TCollection, TResult> resultSelector)
 {
       foreach (TSource element in source)
       {
           foreach (TCollection subElement in collectionSelector(element))
           {
               yield return resultSelector(element, subElement);
           }
       }
   }

相比于SelectManySingleSelectorIterator的MoveNext方法，我们可以更加直观的看到，该方法就是通过一个foreach嵌套循环，完成投影和平铺操作。SelectMany其实就是对嵌套foreach的封装。

该方法有三个参数，第一个是集合的扩展方法参数，TSource是集合内元素的类型，第二个参数是一个投影选择器函数，它封装了一个TSource到IEnumerable投影方法，第三个参数是平铺过程中的结果选择器方法。
遍历souce序列的每个元素
以source序列的元素为参数，执行collectionSelector投影方法。
遍历投影后的序列元素subElement ，以souce序列的元素和subElement 为参数，调用resultSelector方法，执行平铺操作。每次返回一个平铺后的结果，foreach和LINQ 查询可以直接使用该结果。

SelectMany的基本执行流程

本文将源码中的SelectMany.cs文件抽取了出来，增加了一些日志，以方便我们更好的了解SelectMany方法的执行流程。增加日志的源码文件详见附录。为了避免命名冲突，我们将SelectMany方法名称改为SelectMany2。

延迟加载

我们以一个老师，学生一对多的关系来打遍历所有老师列表，从而获取全部学生列表。Student/Teacher类代码见附录

 static void Main(string[] args)
 {
     List<Student> studentList = new List<Student>(){ 
                new Student("x001", "Tom", "CN-1" , 90),
                new Student("x002", "Jack", "CN-1", 88),
                new Student("x003", "Mary", "CN-2", 87),
                new Student("x004", "Frank", "CN-2", 97),
            };

            List<Student> studentList1 = new List<Student>(){ 
                new Student("x005", "Henry", "CN-1" , 90),
                new Student("x006", "Lance", "CN-1", 88),
                new Student("x007", "Steven", "CN-2", 87),
                new Student("x008", "Carl", "CN-2", 97),
            };

            Teacher teacher1 = new Teacher{
                Id = "t001",
                Name  = "Jane",
                Students = studentList
            };
            Teacher teacher2 = new Teacher{
                Id = "t002",
                Name  = "David",
                Students = studentList1
            };

            List<Teacher> teachers = new List<Teacher>{
                teacher1,teacher2
            };
            var students = teachers.SelectMany2(t => t.Students);
 }

执行结果如下：

SelectMany2方法被调用，当前source是一个List
(1) source序列非空
(2) 从teacher到IEnumerable的投影器非空
实例化一个SelectManySingleSelectorIterator对象，先调用基类的构造方法，再调用派生类的构造方法。

SelectMany方法是支持延迟加载的，所以在不真正使用投影操作的返回值的时候，它也只返回一个投影操作迭代器SelectManySingleSelectorIterator实例，不返回具体迭代结果。

数组降维

我们使用SelectMany方法将二维数组将为一维数组，并将返回结果字符串转为大写

 static void Main(string[] args)
 {
List<List<string>> lists = new List<List<string>>{
     new List<string>{"aa","bb","cc"},
      new List<string>{"dd","ee"},
      new List<string>{"ff"},
  };
  
  var d1 = lists.SelectMany2(x => x, (x,s) => s.ToUpper());
  foreach( var s in d1){
      System.Console.WriteLine(s);
  }