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

Hau*_*aus 5 c# ienumerable roslyn

要添加foreach对自定义集合的支持,您需要实现IEnumerable. 阵列,但是,是特殊的,因为它们基本上是编译成一个范围为基础的循环,这是很多比使用一个IEnumerable更快。一个简单的基准测试证实:

                number of elements: 20,000,000
                            byte[]:  6.860ms
       byte[] as IEnumerable<byte>: 89.444ms
CustomCollection.IEnumerator<byte>: 89.667ms
Run Code Online (Sandbox Code Playgroud)

基准:

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;
}
Run Code Online (Sandbox Code Playgroud)

和实施:

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 );
  }
}
Run Code Online (Sandbox Code Playgroud)

因为数组得到了编译器的特殊处理,所以它们将IEnumerable集合留在了尘埃中。由于 C# 非常注重类型安全,我可以理解为什么会这样,但它仍然会产生大量的开销,尤其是对于我的自定义集合,它以数组完全相同的方式进行枚举。事实上,我的自定义集合比基于范围的 for 循环中的字节数组更快,因为它使用指针算法跳过 CLR 的数组范围检查。

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

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

小智 2

类型实际上不需要实现IEnumerable/IEnumerable<T>才能在foreach语句中使用。该foreach语句是鸭子类型的,这意味着编译器首先查找具有正确签名(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();
}
Run Code Online (Sandbox Code Playgroud)