oecs-sharp/OECS/ChangeBuffer.cs

using R3;

namespace OECS;

/// <summary>
/// Buffers changes during system execution and posts them to R3 subjects
/// when <see cref="Post"/> is called.
///
/// Subscribers receive batched changes: all changes accumulated since the
/// last <see cref="Post"/> call are pushed in a single burst.
/// </summary>
internal class ChangeBuffer
{
    private readonly ChangeSet _pending = new();

    private readonly Subject<EntityChange> _entitySubject = new();
    private readonly Dictionary<Type, TrackedSubject> _componentSubjects = new();
    private readonly Dictionary<QueryKey, TrackedSubject> _querySubjects = new();

    /// <summary>
    /// Wraps a <see cref="Subject{T}"/> with a subscriber count so that
    /// subjects with no remaining subscribers can be cleaned up.
    /// </summary>
    private sealed class TrackedSubject
    {
        public readonly Subject<EntityChange> Subject = new();
        public int SubscriberCount;
    }

    /// <summary>
    /// Identifies a query subscription for cleanup tracking.
    /// </summary>
    private readonly struct QueryKey : IEquatable<QueryKey>
    {
        public readonly Type[] WithTypes;
        public readonly Type[] WithoutTypes;

        public QueryKey(Type[] withTypes, Type[] withoutTypes)
        {
            WithTypes = withTypes;
            WithoutTypes = withoutTypes;
        }

        public bool Equals(QueryKey other)
        {
            if (WithTypes.Length != other.WithTypes.Length) return false;
            if (WithoutTypes.Length != other.WithoutTypes.Length) return false;
            for (int i = 0; i < WithTypes.Length; i++)
                if (WithTypes[i] != other.WithTypes[i]) return false;
            for (int i = 0; i < WithoutTypes.Length; i++)
                if (WithoutTypes[i] != other.WithoutTypes[i]) return false;
            return true;
        }

        public override bool Equals(object? obj) => obj is QueryKey other && Equals(other);

        public override int GetHashCode()
        {
            var hash = new HashCode();
            foreach (var t in WithTypes) hash.Add(t);
            foreach (var t in WithoutTypes) hash.Add(t);
            return hash.ToHashCode();
        }
    }

    /// <summary>
    /// The change set currently accumulating. Cleared after each <see cref="Post"/>.
    /// </summary>
    public ChangeSet Pending => _pending;

    /// <summary>
    /// Pushes all pending changes to R3 subjects, then clears the pending set.
    /// </summary>
    public void Post()
    {
        if (_pending.Count == 0)
            return;

        var changes = _pending.Changes;

        foreach (var change in changes)
        {
            // Push to the global entity subject.
            _entitySubject.OnNext(change);

            // Push to component-specific subjects.
            if (change.ComponentType != null)
            {
                if (_componentSubjects.TryGetValue(change.ComponentType, out var compTracked))
                {
                    compTracked.Subject.OnNext(change);
                }
            }

            // Push to matching query subjects.
            foreach (var (query, tracked) in _querySubjects)
            {
                if (ChangeMatchesQuery(change, query))
                {
                    tracked.Subject.OnNext(change);
                }
            }
        }

        _pending.Clear();
    }

    /// <summary>
    /// Returns an observable that emits all entity changes.
    /// </summary>
    public Observable<EntityChange> ObserveEntityChanges()
    {
        return _entitySubject;
    }

    /// <summary>
    /// Returns an observable that emits changes for a specific component type.
    /// </summary>
    public Observable<EntityChange> ObserveComponentChanges(Type componentType)
    {
        if (!_componentSubjects.TryGetValue(componentType, out var tracked))
        {
            tracked = new TrackedSubject();
            _componentSubjects[componentType] = tracked;
        }

        tracked.SubscriberCount++;
        return WrapWithCleanup(tracked.Subject, () =>
        {
            tracked.SubscriberCount--;
            if (tracked.SubscriberCount <= 0)
            {
                tracked.Subject.Dispose();
                _componentSubjects.Remove(componentType);
            }
        });
    }

    /// <summary>
    /// Returns an observable that emits component-level changes matching
    /// the given query's With types and excluding its Without types.
    /// </summary>
    public Observable<EntityChange> ObserveQuery<T1>(Query<T1> query)
        where T1 : struct
    {
        var key = MakeKey(query);
        return SubscribeQuery(key);
    }

    /// <summary>
    /// Returns an observable that emits component-level changes matching
    /// the given query.
    /// </summary>
    public Observable<EntityChange> ObserveQuery<T1, T2>(Query<T1, T2> query)
        where T1 : struct where T2 : struct
    {
        var key = MakeKey(query);
        return SubscribeQuery(key);
    }

    private Observable<EntityChange> SubscribeQuery(QueryKey key)
    {
        if (!_querySubjects.TryGetValue(key, out var tracked))
        {
            tracked = new TrackedSubject();
            _querySubjects[key] = tracked;
        }

        tracked.SubscriberCount++;
        return WrapWithCleanup(tracked.Subject, () =>
        {
            tracked.SubscriberCount--;
            if (tracked.SubscriberCount <= 0)
            {
                tracked.Subject.Dispose();
                _querySubjects.Remove(key);
            }
        });
    }

    private static QueryKey MakeKey<T1>(Query<T1> query) where T1 : struct
    {
        var without = query.WithoutTypes;
        return new QueryKey(
            [typeof(T1)],
            without != null ? [.. without] : []);
    }

    private static QueryKey MakeKey<T1, T2>(Query<T1, T2> query)
        where T1 : struct where T2 : struct
    {
        var without = query.WithoutTypes;
        return new QueryKey(
            [typeof(T1), typeof(T2)],
            without != null ? [.. without] : []);
    }

    /// <summary>
    /// Wraps an observable so that <paramref name="onLastDispose"/> is called
    /// when the last subscriber disposes.
    /// </summary>
    private static Observable<EntityChange> WrapWithCleanup(
        Observable<EntityChange> source, Action onLastDispose)
    {
        return Observable.Create<EntityChange>(observer =>
        {
            var subscription = source.Subscribe(observer);
            return Disposable.Create(() =>
            {
                subscription.Dispose();
                onLastDispose();
            });
        });
    }

    /// <summary>
    /// Disposes all R3 subjects.
    /// </summary>
    public void Dispose()
    {
        _entitySubject.Dispose();
        foreach (var tracked in _componentSubjects.Values)
        {
            tracked.Subject.Dispose();
        }
        _componentSubjects.Clear();
        foreach (var tracked in _querySubjects.Values)
        {
            tracked.Subject.Dispose();
        }
        _querySubjects.Clear();
    }

    private static bool ChangeMatchesQuery(EntityChange change, QueryKey query)
    {
        // Entity-level changes don't have a component type.
        if (change.ComponentType == null)
            return false;

        // Must be in With types.
        bool inWith = false;
        foreach (var t in query.WithTypes)
        {
            if (t == change.ComponentType) { inWith = true; break; }
        }
        if (!inWith) return false;

        // Must not be in Without types.
        foreach (var t in query.WithoutTypes)
        {
            if (t == change.ComponentType) return false;
        }

        return true;
    }
}