oecs-sharp/src/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 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 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>
/// 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.
/// </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;

    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>;
        _count = _set1?.Count ?? 0;
        _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)
        {
            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.
/// </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;

    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>;
        _count = _set1?.Count ?? 0;
        _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)
        {
            CurrentEntity = _set1.DenseEntities[_index];
            if (CurrentEntity.Id == 1) continue;
            if (!_set2.Contains(CurrentEntity)) continue;
            if (!_set3.Contains(CurrentEntity)) continue;
            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;
    }
}