Search code examples
c#loopsforeachvectorizationienumerator

Using a custom ref struct enumerator in a foreach loop


How can I use a custom ref struct enumerator to be implicitly used in a foreach loop?

I want to be able to do something like foreach (var el in span) {...} using my own ref struct enumerator. The intent of this is to enumerate a span in chunks as Vector<T> instances. It was refactored from a class implementing the IEnumerator<T> interface using Memory<T> instead of Span<T> so it still follows that contract aside from the GetEnumerator and Dispose methods without actually implementing the interfaces.

I did this to guarantee no instantiation onto the heap since anybody vectorizing is in it for performance.

public ref struct VectorIterator<T> where T : struct
{
    public int Index { get; private set; }
    public int Increment { get; }
    public Span<T> VectorizedSpan { get; }

    public VectorIterator() : this(Span<T>.Empty)
    {
    }

    public VectorIterator(T[] array) : this(array.AsSpan())
    {
    }

    public VectorIterator(Span<T> span)
    {
        VectorizedSpan = span;
        Increment = Vector<T>.Count;
        Index = -Increment;
    }

    public bool MoveNext()
    {
        Index += Increment;
        return Index <= VectorizedSpan.Length - Increment;
    }

    public void Reset()
    {
        Index = -Increment;
    }

    public readonly Vector<T> Current => new Vector<T>(VectorizedSpan[Index..]);
    public readonly Span<T> Leftovers => VectorizedSpan[^(VectorizedSpan.Length % Increment)..];
}

public static class VectorIteratorHelper
{
    public static VectorIterator<T> GetVectorIterator<T>(this Span<T> span) where T : struct
        => new VectorIterator<T>(span);

    public static VectorIterator<T> GetVectorIterator<T>(this T[] array) where T : struct
        => array.AsSpan().GetVectorIterator();
}

I have tried to include a GetEnumerator method which did not work because Span<T> already has its own enumerator, Span<T>. Enumerator, which hides my GetEnumerator method. I'm not sure how else I could make the compiler recognize this as an enumerator like how Span<T>. Enumerator is recognized and used in for-each loops.


Solution

  • If I understand what you're asking correctly, you won't be able to replace the default iterator for the type, and you'll need to do something like this:

    foreach(var item in span.GetVectorIterator())
    

    The one addition is foreach loops work by calling the type's GetEnumerator() method, however that resolves. This works even if the type does not implement IEnumerable or IEnumerator at all, as long as an appropriate GetEnumerator() method exists. With that in mind, you may want to add a GetEnumerator() method to your custom VectorIterator type (it can just return this;).

    It also means you can override the method in a derived type to replace/hide the original from the base type, but this is still only good for derived types, and won't help you with built-in types like Span<T>. Finally, as of C# 9 you can tack on the GetEnumerator() method via an extension method, but again: it won't help you hide the original.

    Finally, you may want your VectorIterator<T> to formally implement the IEnumerator and IEnumerator<T> interfaces, especially as you've already done the work and only need to add the notice at the top of the struct definition.