diff --git a/src/libraries/System.Linq/src/System/Linq/Chunk.cs b/src/libraries/System.Linq/src/System/Linq/Chunk.cs index 2a2ddd44961bcf..b3e406f2676ff0 100644 --- a/src/libraries/System.Linq/src/System/Linq/Chunk.cs +++ b/src/libraries/System.Linq/src/System/Linq/Chunk.cs @@ -2,6 +2,7 @@ // The .NET Foundation licenses this file to you under the MIT license. using System.Collections.Generic; +using System.Diagnostics.CodeAnalysis; namespace System.Linq { @@ -50,26 +51,25 @@ public static IEnumerable Chunk(this IEnumerable so private static IEnumerable ChunkIterator(IEnumerable source, int size) { using IEnumerator e = source.GetEnumerator(); - while (e.MoveNext()) - { - TSource[] chunk = new TSource[size]; - chunk[0] = e.Current; - int i = 1; - for (; i < chunk.Length && e.MoveNext(); i++) + if (e.MoveNext()) + { + List chunkBuilder = new(); + while (true) { - chunk[i] = e.Current; - } + do + { + chunkBuilder.Add(e.Current); + } + while (chunkBuilder.Count < size && e.MoveNext()); - if (i == chunk.Length) - { - yield return chunk; - } - else - { - Array.Resize(ref chunk, i); - yield return chunk; - yield break; + yield return chunkBuilder.ToArray(); + + if (chunkBuilder.Count < size || !e.MoveNext()) + { + yield break; + } + chunkBuilder.Clear(); } } } diff --git a/src/libraries/System.Linq/tests/ChunkTests.cs b/src/libraries/System.Linq/tests/ChunkTests.cs index 89cdcecc81342d..ee348604192792 100644 --- a/src/libraries/System.Linq/tests/ChunkTests.cs +++ b/src/libraries/System.Linq/tests/ChunkTests.cs @@ -141,5 +141,14 @@ public void AddingToSourceBeforeIterating() Assert.Equal(new[] {new[] {9999, 0, 888}, new[] {-1, 66, -777}, new[] {1, 2, -12345}, new[] {10}}, chunks); } + + // reproduces https://github.com/dotnet/runtime/issues/67132 + [Fact] + public void DoesNotPrematurelyAllocateHugeArray() + { + int[][] chunks = Enumerable.Range(0, 10).Chunk(int.MaxValue).ToArray(); + + Assert.Equal(new[] { Enumerable.Range(0, 10).ToArray() }, chunks); + } } -} \ No newline at end of file +}