oecs-sharp/OECS/EntityIterator.cs

using System.Runtime.CompilerServices;

namespace OECS;

/// <summary>
/// Provides zero-allocation ref struct iterators for querying entities.
/// Use with <c>foreach</c> or manual <c>while (iter.MoveNext())</c>.
/// Supports <c>break</c>, <c>continue</c>, and early returns.
/// </summary>
public static class EntityIterator
{
    /// <summary>
    /// Returns an iterator over entities matching the query, with ref access
    /// to one component type.
    /// </summary>
    public static Select1<T1> Select<T1>(
        this World world, QueryDescriptor query)
        where T1 : struct
    {
        return new Select1<T1>(world, query);
    }

    /// <summary>
    /// Returns an iterator over entities that have a component of type
    /// <typeparamref name="T1"/>. No <c>Without</c> filter is applied.
    /// </summary>
    public static Select1<T1> Select<T1>(
        this World world)
        where T1 : struct
    {
        return new Select1<T1>(world, new QueryDescriptor(new HashSet<Type>(), new HashSet<Type>()));
    }

    /// <summary>
    /// Returns an iterator over entities matching the query, with ref access
    /// to two component types.
    /// </summary>
    public static Select2<T1, T2> Select<T1, T2>(
        this World world, QueryDescriptor query)
        where T1 : struct where T2 : struct
    {
        return new Select2<T1, T2>(world, query);
    }

    /// <summary>
    /// Returns an iterator over entities that have both components of type
    /// <typeparamref name="T1"/> and <typeparamref name="T2"/>.
    /// No <c>Without</c> filter is applied.
    /// </summary>
    public static Select2<T1, T2> Select<T1, T2>(
        this World world)
        where T1 : struct where T2 : struct
    {
        return new Select2<T1, T2>(world, new QueryDescriptor(new HashSet<Type>(), new HashSet<Type>()));
    }

    /// <summary>
    /// Returns an iterator over entities matching the query, with ref access
    /// to three component types.
    /// </summary>
    public static Select3<T1, T2, T3> Select<T1, T2, T3>(
        this World world, QueryDescriptor query)
        where T1 : struct where T2 : struct where T3 : struct
    {
        return new Select3<T1, T2, T3>(world, query);
    }

    /// <summary>
    /// Returns an iterator over entities that have all three components of type
    /// <typeparamref name="T1"/>, <typeparamref name="T2"/>, and
    /// <typeparamref name="T3"/>. No <c>Without</c> filter is applied.
    /// </summary>
    public static Select3<T1, T2, T3> Select<T1, T2, T3>(
        this World world)
        where T1 : struct where T2 : struct where T3 : struct
    {
        return new Select3<T1, T2, T3>(world, new QueryDescriptor(new HashSet<Type>(), new HashSet<Type>()));
    }
}

/// <summary>
/// Ref struct iterator for queries with one component type.
/// </summary>
public ref struct Select1<T1> where T1 : struct
{
    private readonly World _world;
    private readonly SparseSet<T1>? _set;
    private readonly ComponentStore _store;
    private readonly IReadOnlySet<Type> _without;
    private int _index;
    private readonly int _count;

    internal Select1(World world, QueryDescriptor query)
    {
        _world = world;
        _store = world.Components;
        _without = query.Without;
        _set = _store.GetSet(typeof(T1)) as SparseSet<T1>;
        _count = _set?.Count ?? 0;
        _index = -1;

        _world.BeginIteration();
    }

    public Entity CurrentEntity { get; private set; }
    public ref T1 Current1 => ref _set!.Get(CurrentEntity);

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public bool MoveNext()
    {
        if (_set == null) return false;
        while (++_index < _count)
        {
            CurrentEntity = _set.DenseEntities[_index];
            if (CurrentEntity.Id == 1) continue; // Skip singleton.
            if (!PassesWithout()) continue;
            return true;
        }
        return false;
    }

    public Select1<T1> GetEnumerator() => this;

    public void Dispose()
    {
        _world.EndIteration();
    }

    private bool PassesWithout()
    {
        foreach (var type in _without)
        {
            var set = _store.GetSet(type);
            if (set != null && set.Contains(CurrentEntity))
                return false;
        }
        return true;
    }
}

/// <summary>
/// Ref struct iterator for queries with two component types.
/// Drives iteration from the smaller sparse set to minimize probes.
/// </summary>
public ref struct Select2<T1, T2>
    where T1 : struct where T2 : struct
{
    private readonly World _world;
    private readonly SparseSet<T1>? _set1;
    private readonly SparseSet<T2>? _set2;
    private readonly ComponentStore _store;
    private readonly IReadOnlySet<Type> _without;
    private int _index;
    private readonly int _count;
    private readonly bool _swapped; // true when driving from _set2

    internal Select2(World world, QueryDescriptor query)
    {
        _world = world;
        _store = world.Components;
        _without = query.Without;
        _set1 = _store.GetSet(typeof(T1)) as SparseSet<T1>;
        _set2 = _store.GetSet(typeof(T2)) as SparseSet<T2>;

        int count1 = _set1?.Count ?? 0;
        int count2 = _set2?.Count ?? 0;
        _swapped = count2 < count1;
        _count = _swapped ? count2 : count1;
        _index = -1;

        _world.BeginIteration();
    }

    public Entity CurrentEntity { get; private set; }
    public ref T1 Current1 => ref _set1!.Get(CurrentEntity);
    public ref T2 Current2 => ref _set2!.Get(CurrentEntity);

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public bool MoveNext()
    {
        if (_set1 == null || _set2 == null) return false;
        while (++_index < _count)
        {
            if (_swapped)
            {
                CurrentEntity = _set2!.DenseEntities[_index];
                if (CurrentEntity.Id == 1) continue;
                if (!_set1.Contains(CurrentEntity)) continue;
            }
            else
            {
                CurrentEntity = _set1.DenseEntities[_index];
                if (CurrentEntity.Id == 1) continue;
                if (!_set2.Contains(CurrentEntity)) continue;
            }
            if (!PassesWithout()) continue;
            return true;
        }
        return false;
    }

    public Select2<T1, T2> GetEnumerator() => this;

    public void Dispose()
    {
        _world.EndIteration();
    }

    private bool PassesWithout()
    {
        foreach (var type in _without)
        {
            var set = _store.GetSet(type);
            if (set != null && set.Contains(CurrentEntity))
                return false;
        }
        return true;
    }
}

/// <summary>
/// Ref struct iterator for queries with three component types.
/// Drives iteration from the smallest sparse set to minimize probes.
/// </summary>
public ref struct Select3<T1, T2, T3>
    where T1 : struct where T2 : struct where T3 : struct
{
    private readonly World _world;
    private readonly SparseSet<T1>? _set1;
    private readonly SparseSet<T2>? _set2;
    private readonly SparseSet<T3>? _set3;
    private readonly ComponentStore _store;
    private readonly IReadOnlySet<Type> _without;
    private int _index;
    private readonly int _count;
    private readonly int _driver; // 0 = _set1, 1 = _set2, 2 = _set3

    internal Select3(World world, QueryDescriptor query)
    {
        _world = world;
        _store = world.Components;
        _without = query.Without;
        _set1 = _store.GetSet(typeof(T1)) as SparseSet<T1>;
        _set2 = _store.GetSet(typeof(T2)) as SparseSet<T2>;
        _set3 = _store.GetSet(typeof(T3)) as SparseSet<T3>;

        int count1 = _set1?.Count ?? 0;
        int count2 = _set2?.Count ?? 0;
        int count3 = _set3?.Count ?? 0;

        if (count2 <= count1 && count2 <= count3)
        {
            _driver = 1;
            _count = count2;
        }
        else if (count3 <= count1 && count3 <= count2)
        {
            _driver = 2;
            _count = count3;
        }
        else
        {
            _driver = 0;
            _count = count1;
        }

        _index = -1;

        _world.BeginIteration();
    }

    public Entity CurrentEntity { get; private set; }
    public ref T1 Current1 => ref _set1!.Get(CurrentEntity);
    public ref T2 Current2 => ref _set2!.Get(CurrentEntity);
    public ref T3 Current3 => ref _set3!.Get(CurrentEntity);

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public bool MoveNext()
    {
        if (_set1 == null || _set2 == null || _set3 == null) return false;
        while (++_index < _count)
        {
            switch (_driver)
            {
                case 0:
                    CurrentEntity = _set1.DenseEntities[_index];
                    if (CurrentEntity.Id == 1) continue;
                    if (!_set2.Contains(CurrentEntity)) continue;
                    if (!_set3.Contains(CurrentEntity)) continue;
                    break;
                case 1:
                    CurrentEntity = _set2!.DenseEntities[_index];
                    if (CurrentEntity.Id == 1) continue;
                    if (!_set1.Contains(CurrentEntity)) continue;
                    if (!_set3.Contains(CurrentEntity)) continue;
                    break;
                case 2:
                    CurrentEntity = _set3!.DenseEntities[_index];
                    if (CurrentEntity.Id == 1) continue;
                    if (!_set1.Contains(CurrentEntity)) continue;
                    if (!_set2.Contains(CurrentEntity)) continue;
                    break;
            }
            if (!PassesWithout()) continue;
            return true;
        }
        return false;
    }

    public Select3<T1, T2, T3> GetEnumerator() => this;

    public void Dispose()
    {
        _world.EndIteration();
    }

    private bool PassesWithout()
    {
        foreach (var type in _without)
        {
            var set = _store.GetSet(type);
            if (set != null && set.Contains(CurrentEntity))
                return false;
        }
        return true;
    }
}