编写一个性能与数组foreach相当的IEnumerator

Posted 2023-01-19

要将foreach支持添加到自定义集合，您需要实现IEnumerable。然而，数组的特殊之处在于它们基本上编译为基于范围的for循环，这比使用IEnumerable要快得多。一个简单的基准确认：

                number of elements: 20,000,000
                            byte[]:  6.860ms
       byte[] as IEnumerable<byte>: 89.444ms
CustomCollection.IEnumerator<byte>: 89.667ms

基准：

private byte[] byteArray = new byte[20000000];
private CustomCollection<byte> collection = new CustomCollection<T>( 20000000 );

[Benchmark]
public void enumerateByteArray()

  var counter = 0;
  foreach( var item in byteArray )
     counter += item;


[Benchmark]
public void enumerateByteArrayAsIEnumerable()

  var counter = 0;
  var casted = (IEnumerable<byte>) byteArray;
  foreach( var item in casted )
     counter += item;


[Benchmark]
public void enumerateCollection()

  var counter = 0;
  foreach( var item in collection )
     counter += item;

并实施：

public class CustomCollectionEnumerator : IEnumerable<T> where T : unmanaged

    private CustomCollection<T> _collection;
    private int _index;
    private int _endIndex;

    public CustomCollectionEnumerator( CustomCollection<T> collection )
    
      _collection = collection;
      _index = -1;
      _endIndex = collection.Length;
    

    public bool MoveNext()
    
      if ( _index < _endIndex )
      
        _index++;
        return ( _index < _endIndex );
      
      return false;
    

    public T Current => _collection[ _index ];
    object IEnumerator.Current => _collection[ _index ];
    public void Reset()   _index = -1; 
    public void Dispose()   


public class CustomCollection<T> : IEnumerable<T> where T : unmanaged

  private T* _ptr;

  public int Length  get; private set; 

  public T this[ int index ]
  
    [MethodImpl( MethodImplOptions.AggressiveInlining )]
    get => *_ptr[ index ];
    [MethodImpl( MethodImplOptions.AggressiveInlining )]
    set => *_ptr[ index ] = value;
  

  public IEnumerator<T> GetEnumerator()
  
    return new CustomCollectionEnumerator<T>( this );
  

因为数组从编译器得到特殊处理，所以它们将IEnumerable集合留在尘埃中。由于C＃主要关注类型安全性，我可以理解为什么会出现这种情况，但它仍然会产生大量的开销，特别是对于我的自定义集合，它以与数组完全相同的方式进行枚举。实际上，我的自定义集合比基于for循环的范围中的字节数组更快，因为它使用指针算法来跳过CLR的数组范围检查。

所以我的问题是：有没有办法自定义foreach循环的行为，以便我可以实现与数组相当的性能？也许通过编译器内在函数或用IL手动编译委托？

当然，我总是可以使用基于for循环的范围。我只是好奇是否有任何可能的方法来自定义foreach循环的低级行为，其方式与编译器处理数组的方式类似。

答案

实际上，类型实际上不需要在IEnumerable语句中使用IEnumerable<T> / foreach。 foreach语句是duck-typed，这意味着编译器首先查找具有正确签名的公共方法（GetEnumerator()，MoveNext()和Current），无论它们是否是这些接口的实现，并且只在必要时才回退到接口。

这为一些可以在紧密循环中产生显着差异的优化打开了大门：GetEnumerator()可以返回一个具体类型而不是IEnumerator<T>，然后允许foreach循环使用非虚拟和可能内联调用构建，以及制作枚举器一个struct，以避免GC开销。某些框架集合，例如List<T>也利用了这一点。

与其他几个优化一起，这个基于你的CustomCollection的枚举器非常接近微基准测试中的原始数组循环：

public Enumerator GetEnumerator() => new Enumerator(this);

// Being a ref struct makes it less likely to mess up the pointer usage,
// but doesn't affect the foreach loop
// There is no technical reason why this couldn't implement IEnumerator
// as long as lifetime issues are considered
public unsafe ref struct Enumerator

    // Storing the pointer directly instead of the collection reference to reduce indirection
    // Assuming it's immutable for the lifetime of the enumerator
    private readonly T* _ptr;
    private uint _index;
    private readonly uint _endIndex;

    public T Current
    
        get
        
            // This check could be omitted at the cost of safety if consumers are
            // expected to never manually use the enumerator in an incorrect order
            if (_index >= _endIndex)
                ThrowInvalidOp();

            // Without the (int) cast Desktop x86 generates much worse code,
            // but only if _ptr is generic. Not sure why.
            return _ptr[(int)_index];
        
    

    internal Enumerator(CustomCollection<T> collection)
    
        _ptr = collection._ptr;
        _index = UInt32.MaxValue;
        _endIndex = (uint)collection.Length;
    

    // Technically this could unexpectedly reset the enumerator if someone were to
    // manually call MoveNext() countless times after it returns false for some reason
    public bool MoveNext() => unchecked(++_index) < _endIndex;

    // Pulling this out of the getter improves inlining of Current
    private static void ThrowInvalidOp() => throw new InvalidOperationException();