Optimize AsyncEnumerable.ChunkAsync

Keep up-to-date with https://github.com/dotnet/runtime/blob/8ec62773f573f0915f2b6e1eea306fd4a84f4eb4/src/libraries/System.Linq/src/System/Linq/Chunk.cs

Author: JasonWei512

Ix.NET/Source/System.Linq.Async/System/Linq/Operators/Chunk.cs

@@ -48,25 +48,56 @@ public static async IAsyncEnumerable<TSource[]> ChunkAsync<TSource>(this IAsyncE
 
             await using var e = source.GetConfiguredAsyncEnumerator(cancellationToken, false);
 
+            // Before allocating anything, make sure there's at least one element.
             if (await e.MoveNextAsync())
             {
-                var chunkBuilder = new List<TSource>();
-                while (true)
+                // Now that we know we have at least one item, allocate an initial storage array. This is not
+                // the array we'll yield.  It starts out small in order to avoid significantly overallocating
+                // when the source has many fewer elements than the chunk size.
+                var arraySize = Math.Min(size, 4);
+                int i;
+                do
                 {
-                    do
+                    var array = new TSource[arraySize];
+
+                    // Store the first item.
+                    array[0] = e.Current;
+                    i = 1;
+
+                    if (size != array.Length)
                     {
-                        chunkBuilder.Add(e.Current);
-                    }
-                    while (chunkBuilder.Count < size && await e.MoveNextAsync());
+                        // This is the first chunk. As we fill the array, grow it as needed.
+                        for (; i < size && await e.MoveNextAsync(); i++)
+                        {
+                            if (i >= array.Length)
+                            {
+                                arraySize = (int)Math.Min((uint)size, 2 * (uint)array.Length);
+                                Array.Resize(ref array, arraySize);
+                            }
 
-                    yield return chunkBuilder.ToArray();
+                            array[i] = e.Current;
+                        }
+                    }
+                    else
+                    {
+                        // For all but the first chunk, the array will already be correctly sized.
+                        // We can just store into it until either it's full or MoveNext returns false.
+                        var local = array; // avoid bounds checks by using cached local (`array` is lifted to iterator object as a field)
+                        Debug.Assert(local.Length == size);
+                        for (; (uint)i < (uint)local.Length && await e.MoveNextAsync(); i++)
+                        {
+                            local[i] = e.Current;
+                        }
+                    }
 
-                    if (chunkBuilder.Count < size || !await e.MoveNextAsync())
+                    if (i != array.Length)
                     {
-                        yield break;
+                        Array.Resize(ref array, i);
                     }
-                    chunkBuilder.Clear();
+
+                    yield return array;
                 }
+                while (i >= size && await e.MoveNextAsync());
             }
         }
     }