We’ve seen how to create powerful, Burst-compatible generic algorithms already, but today we’ll take another approach to generic algorithms and implement them in the style of C#’s LINQ. Along the way, we’ll tackle a new challenge by implementing a generic algorithm that allocates a new collection.

LINQ

The extension methods in System.Linq.Enumerable can yield some very compact, readable code. Consider this small script that makes a list of 0–9, finds all the even numbers, squares them, puts the results in another list, and prints them:

class TestScript : MonoBehaviour
{
    void Start()
    {
        List<int> list = new List<int>(10);
        for (int i = 0; i < 10; ++i)
        {
            list.Add(i);
        }
 
        List<int> evensSquared = list
            .Where(x => (x & 1) == 0)
            .Select(x => x * x)
            .ToList();
 
        for (int i = 0; i < evensSquared.Count; ++i)
        {
            print(evensSquared[i]);
        }
    }
}

That’s a remarkable amount of functionality for just a few lines of code! Unfortunately, from a performance perspective it’s extremely expensive. An incredible amount of garbage is created here to allocate List<int> objects and their underlying int[] arrays, copy elements of those arrays when their capacities are reached, box struct Enumerator types into IEnumerator<int> interfaces, make indirect function calls via the IEnumerable<int> and IEnumerator<int> interfaces, allocate and call delegates, and so forth. On top of all that, none of this is Burst-compatible.

Burst

We’ve explored Burst-compatible approaches before, but none of them ended up looking like LINQ. We always adopted an iterator strategy in the style of C++’s generic algorithms. That’s very flexible and powerful, but what if we wanted the simplicity of LINQ? Let’s revisit to see if we can achieve that and still keep compatibility with Burst.

To start, we need to declare interfaces for the LINQ functionality we’re replicating:

public static class Algorithms
{
    public interface ISelector<T>
    {
        T Select(T val);
    }
 
    public interface IPredicate<T>
    {
        bool Check(T val);
    }
 
    public interface ICollection<T, TEnumerator, TCollection>
        where TEnumerator : IEnumerator<T>
    {
        TEnumerator GetEnumerator();
        TCollection Allocate();
        void Add(T element);
    }
}

Here we have ISelector and IPredicate to replace the delegates we used with LINQ and aren’t supported by Burst.

We also have an ICollection to define what we need out of a collection. In this case, T is the type of element, TEnumerator is the type of IEnumerator<T> that enumerates it, and TCollection is the type that is allocated when we Allocate a new collection.

This is actually all we need to implement Where and Select!

public static class Algorithms
{
    public static TCollection Where
        <T, TEnumerator, TPredicate, TCollection>
        (this TCollection collection, TPredicate predicate)
        where TPredicate : IPredicate<T>
        where TEnumerator : IEnumerator<T>
        where TCollection : ICollection<T, TEnumerator, TCollection>
    {
        TCollection results = collection.Allocate();
        TEnumerator enumerator = collection.GetEnumerator();
        while (enumerator.MoveNext())
        {
            T value = enumerator.Current;
            if (predicate.Check(value))
            {
                results.Add(value);
            }
        }
        enumerator.Dispose();
        return results;
    }
 
    public static TCollection Select
        <T, TEnumerator, TSelector, TCollection>
        (this TCollection collection, TSelector selector)
        where TEnumerator : IEnumerator<T>
        where TCollection : ICollection<T, TEnumerator, TCollection>
        where TSelector : ISelector<T>
    {
        TCollection results = collection.Allocate();
        TEnumerator enumerator = collection.GetEnumerator();
        while (enumerator.MoveNext())
        {
            T value = enumerator.Current;
            T selected = selector.Select(value);
            results.Add(selected);
        }
        enumerator.Dispose();
        return results;
    }
}

Both of these follow the same pattern:

  1. Allocate a results collection
  2. Get an enumerator for the input collection
  3. Loop until the enumerator becomes invalid
  4. In the loop, get the current value and optionally add to the results
  5. Dispose the enumerator
  6. Return the results collection

There are many type parameters involved, but only two variable parameters. The first is this TCollection and it’s the collection being operated on, optionally via extension method syntax: x.Where() not Algorithms.Where(x). The second is the functionality to perform on each element: a TPredicate or TSelector in these cases. The return value is the collection created as a result of running the function.

Usage

Next up, we need an ICollection. We can’t use List<T> directly because it’s not an ICollection. We could write a wrapper struct for it that implements ICollection, but List<T> will never be compatible with Burst so that won’t work. Instead, we’ll create our own List<T>-style type that is backed by a NativeArray<T>:

struct DynamicNativeArray<T>
    : Algorithms.ICollection<
        T,
        DynamicNativeArrayEnumerator<T>,
        DynamicNativeArray<T>>
    where T : struct
{
    private NativeArray<T> array;
    private int length;
    private Allocator allocator;
 
    public DynamicNativeArray(int capacity, Allocator allocator)
    {
        array = new NativeArray<T>(capacity, allocator);
        length = 0;
        this.allocator = allocator;
    }
 
    public int Length => length;
 
    public T this[int index]
    {
        get
        {
            if (index >= length)
            {
                throw new IndexOutOfRangeException($"{index} >= {length}");
            }
            return array[index];
        }
        set
        {
            if (index >= length)
            {
                throw new IndexOutOfRangeException($"{index} >= {length}");
            }
            array[index] = value;
        }
    }
 
    public DynamicNativeArrayEnumerator<T> GetEnumerator()
    {
        return new DynamicNativeArrayEnumerator<T>(this, -1);
    }
 
    public DynamicNativeArray<T> Allocate()
    {
        return new DynamicNativeArray<T>(4, allocator);
    }
 
    public void Add(T element)
    {
        if (length == array.Length)
        {
            NativeArray<T> newArray = new NativeArray<T>(
                array.Length * 2,
                allocator);
            NativeArray<T>.Copy(array, newArray, array.Length);
            array = newArray;
        }
        array[length] = element;
        length++;
    }
}

This is a pretty standard implementation of such a “dynamic array” type. It implements the ICollection interface with GetEnumerator, Allocate, and Add and additionally provides an indexer and Length property to mimic NativeArray<T>.

Along with this type, we need an IEnumerator<T> that works with it:

struct DynamicNativeArrayEnumerator<T> : IEnumerator<T>
    where T : struct
{
    private DynamicNativeArray<T> array;
    private int index;
 
    public DynamicNativeArrayEnumerator(
        DynamicNativeArray<T> array,
        int index)
    {
        this.array = array;
        this.index = index;
    }
 
    public void Reset()
    {
        index = -1;
    }
 
    object IEnumerator.Current => Current;
 
    public T Current
    {
        get => array[index];
    }
 
    public bool MoveNext()
    {
        index++;
        return index < array.Length;
    }
 
    public void Dispose()
    {
    }
}

This is a super straightforward implementation that mostly consists of boilerplate imposed by IEnumerator<T>.

Now we just need our IPredicate to get even numbers and our ISelector to get squares of numbers:

struct EvenPredicate : Algorithms.IPredicate<int>
{
    public bool Check(int val)
    {
        return (val & 1) == 0;
    }
}
 
struct SquareSelector : Algorithms.ISelector<int>
{
    public int Select(int val)
    {
        return val * val;
    }
}

Finally, we can implement the test script:

class TestScript : MonoBehaviour
{
    void Start()
    {
        DynamicNativeArray<int> array = new DynamicNativeArray<int>(
            10,
            Allocator.Temp);
        for (int i = 0; i < 10; ++i)
        {
            array.Add(i);
        }
 
        DynamicNativeArray<int> evensSquared = array
            .Where
                <
                    int,
                    DynamicNativeArrayEnumerator<int>,
                    EvenPredicate,
                    DynamicNativeArray<int>
                >
                (new EvenPredicate())
            .Select
                <
                    int,
                    DynamicNativeArrayEnumerator<int>,
                    SquareSelector,
                    DynamicNativeArray<int>
                >
                (new SquareSelector());
 
        for (int i = 0; i < evensSquared.Length; ++i)
        {
            print(evensSquared[i]);
        }
    }
}

This is a fair amount longer than the LINQ version due to the need to explicitly specify all of the type parameters. Without them, the middle part is just this:

DynamicNativeArray<int> evensSquared = array
    .Where(new EvenPredicate())
    .Select(new SquareSelector());

One area where this version is shorter is that there’s no need to end with .ToList(). That’s because the ICollection.Allocate method already returns another instance of the same type of collection. This is less flexible, but more terse.

Conclusions

The end result here is a pretty close match to the original LINQ version, but it’s compatible with Burst, allocates zero garbage, and skips a lot of the expensive parts like virtual function calls and delegates. It shows the power and flexibility of the C# generics system and Burst that supports it. We can pick and choose how we want to implement our generic algorithms depending on our own use cases and design sensibilities. Do we want a lot of abstraction and flexibility? Iterators probably make more sense. Do we want to ease the transition from LINQ? This approach might fit the bill.