diff --git a/src/libraries/System.Linq/src/System/Linq/Take.cs b/src/libraries/System.Linq/src/System/Linq/Take.cs index 8d6ad9acb3998d..9d1931dad25b41 100644 --- a/src/libraries/System.Linq/src/System/Linq/Take.cs +++ b/src/libraries/System.Linq/src/System/Linq/Take.cs @@ -1,8 +1,10 @@ // Licensed to the .NET Foundation under one or more agreements. // The .NET Foundation licenses this file to you under the MIT license. +using System.Buffers; using System.Collections.Generic; using System.Diagnostics; +using System.Runtime.CompilerServices; namespace System.Linq { @@ -104,51 +106,94 @@ private static IEnumerable TakeRangeFromEndIterator(IEnumerabl yield break; } - Queue queue; - if (isStartIndexFromEnd) { - // TakeLast compat: enumerator should be disposed before yielding the first element. - using (IEnumerator e = source.GetEnumerator()) + // Buffer the last `startIndex` elements using a pooled array as a ring buffer, avoiding the + // allocation and per-element enqueue/dequeue overhead of Queue. The pooled array is + // returned to the pool when enumeration completes or is abandoned. + Debug.Assert(startIndex > 0); + + // The pooled window retains at most `startIndex` elements, so the ring buffer wraps at this + // capacity. Capture it now because `startIndex` is repurposed below as the recomputed start index. + int capacity = startIndex; + + TSource[]? buffer = null; + int bufferCount = 0; // number of valid elements currently retained + int head = 0; // logical index of the oldest retained element (only advances once the ring is full) + + // Own the rented buffer for the whole operation with a single finally. Buffering, the source + // enumerator's disposal, and yielding can each throw; this returns the buffer exactly once on every + // path (including abandoned enumeration) and never double-returns, because the buffer variable always + // identifies the single array currently rented from the pool. + try { - if (!e.MoveNext()) - { - yield break; - } - - queue = new Queue(); - queue.Enqueue(e.Current); - count = 1; - - while (e.MoveNext()) + // TakeLast compat: enumerator should be disposed before yielding the first element. + using (IEnumerator e = source.GetEnumerator()) { - if (count < startIndex) + if (!e.MoveNext()) { - queue.Enqueue(e.Current); - ++count; + yield break; } - else + + buffer = ArrayPool.Shared.Rent(Math.Min(startIndex, 4)); + buffer[0] = e.Current; + bufferCount = 1; + head = 0; + count = 1; + + while (e.MoveNext()) { - do + if (count < startIndex) { - queue.Dequeue(); - queue.Enqueue(e.Current); - checked { ++count; } - } while (e.MoveNext()); - break; + // The retained window isn't full yet; append, growing the pooled buffer if needed. + if (bufferCount == buffer.Length) + { + int newSize = (int)Math.Min((uint)capacity, 2 * (uint)buffer.Length); + TSource[] newBuffer = ArrayPool.Shared.Rent(newSize); + Array.Copy(buffer, newBuffer, bufferCount); + ReturnToPool(buffer); + buffer = newBuffer; + } + + buffer[bufferCount++] = e.Current; + ++count; + } + else + { + // The window is full; overwrite the oldest element (ring buffer of capacity `startIndex`). + do + { + buffer[head] = e.Current; + head = head + 1 == capacity ? 0 : head + 1; + checked { ++count; } + } while (e.MoveNext()); + break; + } } } - Debug.Assert(queue.Count == Math.Min(count, startIndex)); - } + Debug.Assert(bufferCount == Math.Min(count, startIndex)); - startIndex = CalculateStartIndex(isStartIndexFromEnd: true, startIndex, count); - endIndex = CalculateEndIndex(isEndIndexFromEnd, endIndex, count); - Debug.Assert(endIndex - startIndex <= queue.Count); + startIndex = CalculateStartIndex(isStartIndexFromEnd: true, startIndex, count); + endIndex = CalculateEndIndex(isEndIndexFromEnd, endIndex, count); + Debug.Assert(endIndex - startIndex <= bufferCount); - for (int rangeIndex = startIndex; rangeIndex < endIndex; rangeIndex++) + // The retained window holds `bufferCount` elements in order starting at `head`. Its first + // element corresponds to the recomputed original index `startIndex`, so the elements to yield + // are simply the first (endIndex - startIndex) of the window. + int index = head; + for (int rangeIndex = startIndex; rangeIndex < endIndex; rangeIndex++) + { + yield return buffer[index]; + index = index + 1 == capacity ? 0 : index + 1; + } + } + finally { - yield return queue.Dequeue(); + if (buffer is not null) + { + ReturnToPool(buffer); + } } } else @@ -166,7 +211,7 @@ private static IEnumerable TakeRangeFromEndIterator(IEnumerabl if (count == startIndex) { - queue = new Queue(); + Queue queue = new Queue(); while (e.MoveNext()) { if (queue.Count == endIndex) @@ -192,6 +237,9 @@ static int CalculateStartIndex(bool isStartIndexFromEnd, int startIndex, int cou static int CalculateEndIndex(bool isEndIndexFromEnd, int endIndex, int count) => Math.Min(count, isEndIndexFromEnd ? count - endIndex : endIndex); + + static void ReturnToPool(TSource[] buffer) => + ArrayPool.Shared.Return(buffer, clearArray: RuntimeHelpers.IsReferenceOrContainsReferences()); } public static IEnumerable TakeWhile(this IEnumerable source, Func predicate)