AsyncEnumerable optimizations

Package/Core/Cancelations/Internal/CancelationInternal.cs

@@ -26,7 +26,7 @@
 
 namespace Proto.Promises
 {
-    internal static partial class Internal
+    partial class Internal
     {
 #if !PROTO_PROMISE_DEVELOPER_MODE
         [DebuggerNonUserCode, StackTraceHidden]

Package/Core/InternalShared/InterfacesInternal.cs

@@ -6,7 +6,7 @@
 
 namespace Proto.Promises
 {
-    internal static partial class Internal
+    partial class Internal
     {
         internal partial interface ITraceable { }
 

Package/Core/InternalShared/ValueStopwatch.cs

@@ -18,7 +18,7 @@
 
 namespace Proto.Promises
 {
-    internal static partial class Internal
+    partial class Internal
     {
         // Idea from https://www.meziantou.net/how-to-measure-elapsed-time-without-allocating-a-stopwatch.htm
 #pragma warning disable IDE0250 // Make struct 'readonly'

Package/Core/Linq/CompilerServices/AsyncStreamWriter.cs

@@ -13,11 +13,11 @@ namespace Proto.Promises.Async.CompilerServices
 #endif
     public readonly struct AsyncStreamWriter<T>
     {
-        private readonly Internal.PromiseRefBase.AsyncEnumerableBase<T> _target;
+        private readonly Internal.PromiseRefBase.AsyncEnumerableWithIterator<T> _target;
         private readonly int _id;
 
         [MethodImpl(Internal.InlineOption)]
-        internal AsyncStreamWriter(Internal.PromiseRefBase.AsyncEnumerableBase<T> target, int id)
+        internal AsyncStreamWriter(Internal.PromiseRefBase.AsyncEnumerableWithIterator<T> target, int id)
         {
             _target = target;
             _id = id;
@@ -41,11 +41,11 @@ public AsyncStreamYielder<T> YieldAsync(T value)
 #endif
     public readonly partial struct AsyncStreamYielder<T> : ICriticalNotifyCompletion, Internal.IPromiseAwaiter
     {
-        private readonly Internal.PromiseRefBase.AsyncEnumerableBase<T> _target;
+        private readonly Internal.PromiseRefBase.AsyncEnumerableWithIterator<T> _target;
         private readonly int _enumerableId;
 
         [MethodImpl(Internal.InlineOption)]
-        internal AsyncStreamYielder(Internal.PromiseRefBase.AsyncEnumerableBase<T> target, int enumerableId)
+        internal AsyncStreamYielder(Internal.PromiseRefBase.AsyncEnumerableWithIterator<T> target, int enumerableId)
         {
             _target = target;
             _enumerableId = enumerableId;

Package/Core/Linq/Generators/Canceled.cs

@@ -12,7 +12,7 @@
 namespace Proto.Promises.Linq
 {
 #if CSHARP_7_3_OR_NEWER // We only expose AsyncEnumerable where custom async method builders are supported.
-    public static partial class AsyncEnumerable
+    partial class AsyncEnumerable
     {
         /// <summary>
         /// Generates an async-enumerable sequence that will be immediately canceled when iterated.
@@ -72,8 +72,6 @@ internal override Promise DisposeAsync(int id)
                 // Do nothing, just return a resolved promise.
                 => Promise.Resolved();
 
-            protected override void Start(int enumerableId) { throw new System.InvalidOperationException(); }
-            protected override void DisposeAndReturnToPool() { throw new System.InvalidOperationException(); }
             internal override void MaybeDispose() { throw new System.InvalidOperationException(); }
         }
     }

Package/Core/Linq/Generators/Create.cs

@@ -43,7 +43,7 @@ public static AsyncEnumerable<T> Create(Func<AsyncStreamWriter<T>, CancelationTo
         {
             ValidateArgument(asyncIterator, nameof(asyncIterator), 1);
 
-            var enumerable = Internal.AsyncEnumerableImpl<T, Internal.AsyncIterator<T>>.GetOrCreate(new Internal.AsyncIterator<T>(asyncIterator));
+            var enumerable = Internal.AsyncEnumerableCreate<T, Internal.AsyncIterator<T>>.GetOrCreate(new Internal.AsyncIterator<T>(asyncIterator));
             return new AsyncEnumerable<T>(enumerable);
         }
 
@@ -54,7 +54,7 @@ public static AsyncEnumerable<T> Create<TCapture>(TCapture captureValue, Func<TC
         {
             ValidateArgument(asyncIterator, nameof(asyncIterator), 1);
 
-            var enumerable = Internal.AsyncEnumerableImpl<T, Internal.AsyncIterator<T, TCapture>>.GetOrCreate(new Internal.AsyncIterator<T, TCapture>(captureValue, asyncIterator));
+            var enumerable = Internal.AsyncEnumerableCreate<T, Internal.AsyncIterator<T, TCapture>>.GetOrCreate(new Internal.AsyncIterator<T, TCapture>(captureValue, asyncIterator));
             return new AsyncEnumerable<T>(enumerable);
         }
     }

Package/Core/Linq/Generators/Empty.cs

@@ -10,7 +10,7 @@
 namespace Proto.Promises.Linq
 {
 #if CSHARP_7_3_OR_NEWER // We only expose AsyncEnumerable where custom async method builders are supported.
-    public static partial class AsyncEnumerable
+    partial class AsyncEnumerable
     {
         /// <summary>
         /// Returns an empty async-enumerable sequence.
@@ -21,7 +21,7 @@ public static AsyncEnumerable<T> Empty<T>()
             => AsyncEnumerable<T>.Empty();
     }
 
-    public readonly partial struct AsyncEnumerable<T>
+    partial struct AsyncEnumerable<T>
     {
         /// <summary>
         /// Returns an empty async-enumerable sequence.
@@ -69,8 +69,6 @@ internal override Promise DisposeAsync(int id)
                 // Do nothing, just return a resolved promise.
                 => Promise.Resolved();
 
-            protected override void Start(int enumerableId) { throw new System.InvalidOperationException(); }
-            protected override void DisposeAndReturnToPool() { throw new System.InvalidOperationException(); }
             internal override void MaybeDispose() { throw new System.InvalidOperationException(); }
         }
     }

Package/Core/Linq/Generators/Range.cs

@@ -12,7 +12,7 @@
 namespace Proto.Promises.Linq
 {
 #if CSHARP_7_3_OR_NEWER // We only expose AsyncEnumerable where custom async method builders are supported.
-    public static partial class AsyncEnumerable
+    partial class AsyncEnumerable
     {
         /// <summary>
         /// Generates an async-enumerable sequence of integral numbers within a specified range.
@@ -104,23 +104,22 @@ internal override Promise DisposeAsync(int id)
                 if (Interlocked.CompareExchange(ref _enumerableId, id + 1, id) == id)
                 {
                     // This was not already disposed.
-#if PROMISE_DEBUG || PROTO_PROMISE_DEVELOPER_MODE
-                    SetCompletionState(null, Promise.State.Resolved);
-#endif
-                    DisposeAndReturnToPool();
+                    Dispose();
                 }
                 // IAsyncDisposable.DisposeAsync must not throw if it's called multiple times, according to MSDN documentation.
                 return Promise.Resolved();
             }
 
-            protected override void DisposeAndReturnToPool()
+            new private void Dispose()
             {
-                Dispose();
+#if PROMISE_DEBUG || PROTO_PROMISE_DEVELOPER_MODE
+                SetCompletionState(null, Promise.State.Resolved);
+#endif
+                base.Dispose();
                 _disposed = true;
                 ObjectPool.MaybeRepool(this);
             }
 
-            protected override void Start(int enumerableId) { throw new System.InvalidOperationException(); }
             internal override void MaybeDispose() { throw new System.InvalidOperationException(); }
         }
     }

Package/Core/Linq/Generators/Rejected.cs

@@ -12,7 +12,7 @@
 namespace Proto.Promises.Linq
 {
 #if CSHARP_7_3_OR_NEWER // We only expose AsyncEnumerable where custom async method builders are supported.
-    public static partial class AsyncEnumerable
+    partial class AsyncEnumerable
     {
         /// <summary>
         /// Generates an async-enumerable sequence that will be immediately rejected with the provided reason when iterated.

Package/Core/Linq/Generators/Repeat.cs

@@ -12,7 +12,7 @@
 namespace Proto.Promises.Linq
 {
 #if CSHARP_7_3_OR_NEWER // We only expose AsyncEnumerable where custom async method builders are supported.
-    public static partial class AsyncEnumerable
+    partial class AsyncEnumerable
     {
         /// <summary>
         /// Generates an async-enumerable sequence that contains one repeated value.
@@ -107,24 +107,23 @@ internal override Promise DisposeAsync(int id)
                 if (Interlocked.CompareExchange(ref _enumerableId, id + 1, id) == id)
                 {
                     // This was not already disposed.
-#if PROMISE_DEBUG || PROTO_PROMISE_DEVELOPER_MODE
-                    SetCompletionState(null, Promise.State.Resolved);
-#endif
-                    DisposeAndReturnToPool();
+                    Dispose();
                 }
                 // IAsyncDisposable.DisposeAsync must not throw if it's called multiple times, according to MSDN documentation.
                 return Promise.Resolved();
             }
 
-            protected override void DisposeAndReturnToPool()
+            new private void Dispose()
             {
-                Dispose();
+#if PROMISE_DEBUG || PROTO_PROMISE_DEVELOPER_MODE
+                SetCompletionState(null, Promise.State.Resolved);
+#endif
+                base.Dispose();
                 _current = default;
                 _disposed = true;
                 ObjectPool.MaybeRepool(this);
             }
 
-            protected override void Start(int enumerableId) { throw new System.InvalidOperationException(); }
             internal override void MaybeDispose() { throw new System.InvalidOperationException(); }
         }
     }

Package/Core/Linq/Generators/Return.cs

@@ -12,7 +12,7 @@
 namespace Proto.Promises.Linq
 {
 #if CSHARP_7_3_OR_NEWER // We only expose AsyncEnumerable where custom async method builders are supported.
-    public static partial class AsyncEnumerable
+    partial class AsyncEnumerable
     {
         /// <summary>
         /// Generates an async-enumerable sequence that contains a single element.
@@ -83,24 +83,23 @@ internal override Promise DisposeAsync(int id)
                 if (Interlocked.CompareExchange(ref _enumerableId, id + 1, id) == id)
                 {
                     // This was not already disposed.
-#if PROMISE_DEBUG || PROTO_PROMISE_DEVELOPER_MODE
-                    SetCompletionState(null, Promise.State.Resolved);
-#endif
-                    DisposeAndReturnToPool();
+                    Dispose();
                 }
                 // IAsyncDisposable.DisposeAsync must not throw if it's called multiple times, according to MSDN documentation.
                 return Promise.Resolved();
             }
 
-            protected override void DisposeAndReturnToPool()
+            new private void Dispose()
             {
-                Dispose();
+#if PROMISE_DEBUG || PROTO_PROMISE_DEVELOPER_MODE
+                SetCompletionState(null, Promise.State.Resolved);
+#endif
+                base.Dispose();
                 _current = default;
                 _disposed = true;
                 ObjectPool.MaybeRepool(this);
             }
 
-            protected override void Start(int enumerableId) { throw new System.InvalidOperationException(); }
             internal override void MaybeDispose() { throw new System.InvalidOperationException(); }
         }
     }

Package/Core/Linq/Internal/AsyncEnumerableInternal.cs

@@ -76,17 +76,11 @@ partial class PromiseRefBase
 #endif
             internal abstract class AsyncEnumerableBase<T> : PromiseSingleAwait<bool>
             {
-                // This is used as the backing reference to 3 different awaiters. MoveNextAsync (Promise<bool>), DisposeAsync (Promise), and YieldAsync (AsyncStreamYielder<T>).
-                // We use `Interlocked.CompareExchange(ref _enumerableId` to enforce only 1 awaiter uses it at a time, in the correct order.
-                // We use a separate field for AsyncStreamYielder continuation, because using _next for 2 separate async functions (the iterator and the consumer) proves problematic.
-                protected PromiseRefBase _iteratorPromiseRef;
+                internal CancelationToken _cancelationToken;
                 protected T _current;
-                private int _iteratorCompleteExpectedId;
-                private int _iteratorCompleteId;
                 protected int _enumerableId = 1; // Start with Id 1 instead of 0 to reduce risk of false positives.
                 protected bool _disposed;
                 protected bool _isStarted;
-                internal CancelationToken _cancelationToken;
 
                 internal int EnumerableId
                 {
@@ -149,8 +143,23 @@ internal T GetCurrent(int id)
                     return _current;
                 }
 
-                [MethodImpl(InlineOption)]
-                internal virtual Promise<bool> MoveNextAsync(int id)
+                internal abstract Promise<bool> MoveNextAsync(int id);
+                internal abstract Promise DisposeAsync(int id);
+            } // class AsyncEnumerableBase<T>
+
+#if !PROTO_PROMISE_DEVELOPER_MODE
+            [DebuggerNonUserCode, StackTraceHidden]
+#endif
+            internal abstract class AsyncEnumerableWithIterator<T> : AsyncEnumerableBase<T>
+            {
+                // This is used as the backing reference to 3 different awaiters. MoveNextAsync (Promise<bool>), DisposeAsync (Promise), and YieldAsync (AsyncStreamYielder<T>).
+                // We use `Interlocked.CompareExchange(ref _enumerableId` to enforce only 1 awaiter uses it at a time, in the correct order.
+                // We use a separate field for AsyncStreamYielder continuation, because using _next for 2 separate async functions (the iterator and the consumer) proves problematic.
+                protected PromiseRefBase _iteratorPromiseRef;
+                private int _iteratorCompleteExpectedId;
+                private int _iteratorCompleteId;
+
+                internal override Promise<bool> MoveNextAsync(int id)
                 {
                     // We increment by 1 when MoveNextAsync, then decrement by 1 when YieldAsync.
                     int newId = id + 1;
@@ -183,6 +192,16 @@ internal virtual Promise<bool> MoveNextAsync(int id)
                     return new Promise<bool>(this, Id, 0);
                 }
 
+                private void MoveNext()
+                {
+                    // Invalidate the previous awaiter.
+                    IncrementPromiseIdAndClearPrevious();
+                    // Reset for the next awaiter.
+                    ResetWithoutStacktrace();
+                    // Handle iterator promise to move the async state machine forward.
+                    InterlockedExchange(ref _iteratorPromiseRef, null).Handle(this, null, Promise.State.Resolved);
+                }
+
                 [MethodImpl(InlineOption)]
                 internal AsyncStreamYielder<T> YieldAsync(in T value, int id)
                 {
@@ -199,7 +218,7 @@ internal AsyncStreamYielder<T> YieldAsync(in T value, int id)
                     return new AsyncStreamYielder<T>(this, newId);
                 }
 
-                internal virtual Promise DisposeAsync(int id)
+                internal override Promise DisposeAsync(int id)
                 {
                     int newId = id + 3;
                     // When the async iterator function completes before DisposeAsync is called, it's set to id + 2.
@@ -317,43 +336,32 @@ internal void AwaitOnCompletedForAsyncStreamYielder(PromiseRefBase asyncPromiseR
                     HandleNextInternal(null, Promise.State.Resolved);
                 }
 
-                protected void MoveNext()
-                {
-                    // Invalidate the previous awaiter.
-                    IncrementPromiseIdAndClearPrevious();
-                    // Reset for the next awaiter.
-                    ResetWithoutStacktrace();
-                    // Handle iterator promise to move the async state machine forward.
-                    InterlockedExchange(ref _iteratorPromiseRef, null).Handle(this, null, Promise.State.Resolved);
-                }
-
                 protected abstract void Start(int enumerableId);
-
                 protected abstract void DisposeAndReturnToPool();
-            } // class AsyncEnumerableBase<T>
+            } // class AsyncEnumerableWithIterator<TValue>
         } // class PromiseRefBase
 
 #if !PROTO_PROMISE_DEVELOPER_MODE
         [DebuggerNonUserCode, StackTraceHidden]
 #endif
-        internal sealed class AsyncEnumerableImpl<TValue, TIterator> : PromiseRefBase.AsyncEnumerableBase<TValue>
+        internal sealed class AsyncEnumerableCreate<TValue, TIterator> : PromiseRefBase.AsyncEnumerableWithIterator<TValue>
             where TIterator : IAsyncIterator<TValue>
         {
             private TIterator _iterator;
 
-            private AsyncEnumerableImpl() { }
+            private AsyncEnumerableCreate() { }
 
             [MethodImpl(InlineOption)]
-            private static AsyncEnumerableImpl<TValue, TIterator> GetOrCreate()
+            private static AsyncEnumerableCreate<TValue, TIterator> GetOrCreate()
             {
-                var obj = ObjectPool.TryTakeOrInvalid<AsyncEnumerableImpl<TValue, TIterator>>();
+                var obj = ObjectPool.TryTakeOrInvalid<AsyncEnumerableCreate<TValue, TIterator>>();
                 return obj == InvalidAwaitSentinel.s_instance
-                    ? new AsyncEnumerableImpl<TValue, TIterator>()
-                    : obj.UnsafeAs<AsyncEnumerableImpl<TValue, TIterator>>();
+                    ? new AsyncEnumerableCreate<TValue, TIterator>()
+                    : obj.UnsafeAs<AsyncEnumerableCreate<TValue, TIterator>>();
             }
 
             [MethodImpl(InlineOption)]
-            internal static AsyncEnumerableImpl<TValue, TIterator> GetOrCreate(in TIterator iterator)
+            internal static AsyncEnumerableCreate<TValue, TIterator> GetOrCreate(in TIterator iterator)
             {
                 var enumerable = GetOrCreate();
                 enumerable.Reset();

Package/Core/Linq/Internal/MergeInternal.cs

@@ -22,7 +22,7 @@ partial class Internal
 #if !PROTO_PROMISE_DEVELOPER_MODE
         [DebuggerNonUserCode, StackTraceHidden]
 #endif
-        internal abstract class AsyncEnumerableMergerBase<TValue> : PromiseRefBase.AsyncEnumerableBase<TValue>
+        internal abstract class AsyncEnumerableMergerBase<TValue> : PromiseRefBase.AsyncEnumerableWithIterator<TValue>
         {
             // TODO: optimize these collections.
             protected readonly List<AsyncEnumerator<TValue>> _enumerators = new List<AsyncEnumerator<TValue>>();

Package/Core/Promises/Internal/RejectContainersInternal.cs

@@ -22,7 +22,7 @@
 
 namespace Proto.Promises
 {
-    internal static partial class Internal
+    partial class Internal
     {
         // Extension method instead of including on the interface, since old IL2CPP compiler does not support virtual generics with structs.
         internal static bool TryGetValue<TValue>(this IRejectContainer rejectContainer, out TValue converted)