UniTask/src/UniTask.NetCore/IAsyncEnumerable.cs

using System.Runtime.InteropServices;
using System.Threading;

namespace Cysharp.Threading.Tasks
{
    public interface IUniTaskAsyncEnumerable<out T>
    {
        IUniTaskAsyncEnumerator<T> GetAsyncEnumerator(CancellationToken cancellationToken = default);
    }

    public interface IUniTaskAsyncEnumerator<out T> : IUniTaskAsyncDisposable
    {
        T Current { get; }
        UniTask<bool> MoveNextAsync();
    }

    public interface IUniTaskAsyncDisposable
    {
        UniTask DisposeAsync();
    }

    //public interface IUniTaskAsyncGrouping<out TKey, out TElement> : IUniTaskAsyncEnumerable<TElement>
    //{
    //    TKey Key { get; }
    //}

    public static class UniTaskAsyncEnumerableExtensions
    {
        public static UniTaskCancelableAsyncEnumerable<T> WithCancellation<T>(this IUniTaskAsyncEnumerable<T> source, CancellationToken cancellationToken)
        {
            return new UniTaskCancelableAsyncEnumerable<T>(source, cancellationToken);
        }
    }

    [StructLayout(LayoutKind.Auto)]
    public readonly struct UniTaskCancelableAsyncEnumerable<T>
    {
        private readonly IUniTaskAsyncEnumerable<T> enumerable;
        private readonly CancellationToken cancellationToken;

        internal UniTaskCancelableAsyncEnumerable(IUniTaskAsyncEnumerable<T> enumerable, CancellationToken cancellationToken)
        {
            this.enumerable = enumerable;
            this.cancellationToken = cancellationToken;
        }

        public Enumerator GetAsyncEnumerator()
        {
            return new Enumerator(enumerable.GetAsyncEnumerator(cancellationToken));
        }

        [StructLayout(LayoutKind.Auto)]
        public readonly struct Enumerator
        {
            private readonly IUniTaskAsyncEnumerator<T> enumerator;

            internal Enumerator(IUniTaskAsyncEnumerator<T> enumerator)
            {
                this.enumerator = enumerator;
            }

            public T Current => enumerator.Current;

            public UniTask<bool> MoveNextAsync()
            {
                return enumerator.MoveNextAsync();
            }


            public UniTask DisposeAsync()
            {
                return enumerator.DisposeAsync();
            }
        }
    }
}