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

时间:2019-02-23 00:41:10

标签: c# ienumerable roslyn

要将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循环的低级行为。

1 个答案:

答案 0 :(得分: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();
}