Reduce allocations in {Concat,SelectMany,Append,Prepend}.ToArray by close to 50%

src/Common/src/System/Collections/Generic/ArrayBuilder.cs

@@ -9,6 +9,7 @@ namespace System.Collections.Generic
  /// <summary>
  /// Helper type for avoiding allocations while building arrays.
  /// </summary>
+ /// <typeparam name="T">The element type.</typeparam>
  internal struct ArrayBuilder<T>
  {
  private const int DefaultCapacity = 4;
@@ -74,7 +75,7 @@ public void Add(T item)
  }
 
  /// <summary>
- /// Returns an array with equivalent contents as this builder.
+ /// Creates an array from the contents of this builder.
  /// </summary>
  /// <remarks>
  /// Do not call this method twice on the same builder.

src/Common/src/System/Collections/Generic/EnumerableHelpers.Linq.cs

@@ -0,0 +1,44 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System.Diagnostics;
+using System.Linq;
+
+namespace System.Collections.Generic
+{
+ /// <summary>
+ /// Internal helper functions for working with enumerables.
+ /// </summary>
+ internal static partial class EnumerableHelpers
+ {
+ /// <summary>
+ /// Tries to get the count of the enumerable cheaply.
+ /// </summary>
+ /// <typeparam name="T">The element type of the source enumerable.</typeparam>
+ /// <param name="source">The enumerable to count.</param>
+ /// <param name="count">The count of the enumerable, if it could be obtained cheaply.</param>
+ /// <returns><c>true</c> if the enumerable could be counted cheaply; otherwise, <c>false</c>.</returns>
+ internal static bool TryGetCount<T>(IEnumerable<T> source, out int count)
+ {
+ Debug.Assert(source != null);
+
+ var collection = source as ICollection<T>;
+ if (collection != null)
+ {
+ count = collection.Count;
+ return true;
+ }
+
+ var provider = source as IIListProvider<T>;
+ if (provider != null)
+ {
+ count = provider.GetCount(onlyIfCheap: true);
+ return count >= 0;
+ }
+
+ count = -1;
+ return false;
+ }
+ }
+}

src/Common/src/System/Collections/Generic/EnumerableHelpers.cs

@@ -6,9 +6,61 @@
 
 namespace System.Collections.Generic
 {
- /// <summary>Internal helper functions for working with enumerables.</summary>
- internal static class EnumerableHelpers
+ /// <summary>
+ /// Internal helper functions for working with enumerables.
+ /// </summary>
+ internal static partial class EnumerableHelpers
  {
+ /// <summary>
+ /// Copies items from an enumerable to an array.
+ /// </summary>
+ /// <typeparam name="T">The element type of the enumerable.</typeparam>
+ /// <param name="source">The source enumerable.</param>
+ /// <param name="array">The destination array.</param>
+ /// <param name="arrayIndex">The index in the array to start copying to.</param>
+ /// <param name="count">The number of items in the enumerable.</param>
+ internal static void Copy<T>(IEnumerable<T> source, T[] array, int arrayIndex, int count)
+ {
+ Debug.Assert(source != null);
+ Debug.Assert(arrayIndex >= 0);
+ Debug.Assert(count >= 0);
+ Debug.Assert(array?.Length - arrayIndex >= count);
+
+ var collection = source as ICollection<T>;
+ if (collection != null)
+ {
+ Debug.Assert(collection.Count == count);
+ collection.CopyTo(array, arrayIndex);
+ return;
+ }
+
+ IterativeCopy(source, array, arrayIndex, count);
+ }
+
+ /// <summary>
+ /// Copies items from a non-collection enumerable to an array.
+ /// </summary>
+ /// <typeparam name="T">The element type of the enumerable.</typeparam>
+ /// <param name="source">The source enumerable.</param>
+ /// <param name="array">The destination array.</param>
+ /// <param name="arrayIndex">The index in the array to start copying to.</param>
+ /// <param name="count">The number of items in the enumerable.</param>
+ internal static void IterativeCopy<T>(IEnumerable<T> source, T[] array, int arrayIndex, int count)
+ {
+ Debug.Assert(source != null && !(source is ICollection<T>));
+ Debug.Assert(arrayIndex >= 0);
+ Debug.Assert(count >= 0);
+ Debug.Assert(array?.Length - arrayIndex >= count);
+
+ int endIndex = arrayIndex + count;
+ foreach (T item in source)
+ {
+ array[arrayIndex++] = item;
+ }
+
+ Debug.Assert(arrayIndex == endIndex);
+ }
+
  /// <summary>Converts an enumerable to an array.</summary>
  /// <param name="source">The enumerable to convert.</param>
  /// <returns>The resulting array.</returns>

src/Common/src/System/Collections/Generic/LargeArrayBuilder.cs

@@ -7,9 +7,51 @@
 
 namespace System.Collections.Generic
 {
+ /// <summary>
+ /// Represents a position within a <see cref="LargeArrayBuilder{T}"/>.
+ /// </summary>
+ [DebuggerDisplay("{DebuggerDisplay,nq}")]
+ internal struct CopyPosition
+ {
+ /// <summary>
+ /// Constructs a new <see cref="CopyPosition"/>.
+ /// </summary>
+ /// <param name="row">The index of the buffer to select.</param>
+ /// <param name="column">The index within the buffer to select.</param>
+ internal CopyPosition(int row, int column)
+ {
+ Debug.Assert(row >= 0);
+ Debug.Assert(column >= 0);
+
+ Row = row;
+ Column = column;
+ }
+
+ /// <summary>
+ /// Represents a position at the start of a <see cref="LargeArrayBuilder{T}"/>.
+ /// </summary>
+ public static CopyPosition Start => default(CopyPosition);
+
+ /// <summary>
+ /// The index of the buffer to select.
+ /// </summary>
+ internal int Row { get; }
+
+ /// <summary>
+ /// The index within the buffer to select.
+ /// </summary>
+ internal int Column { get; }
+
+ /// <summary>
+ /// Gets a string suitable for display in the debugger.
+ /// </summary>
+ private string DebuggerDisplay => $"[{Row}, {Column}]";
+ }
+
  /// <summary>
  /// Helper type for building dynamically-sized arrays while minimizing allocations and copying.
  /// </summary>
+ /// <typeparam name="T">The element type.</typeparam>
  internal struct LargeArrayBuilder<T>
  {
  private const int StartingCapacity = 4;
@@ -42,13 +84,11 @@ public LargeArrayBuilder(bool initialize)
  /// Do not add more than <paramref name="maxCapacity"/> items to this builder.
  /// </remarks>
  public LargeArrayBuilder(int maxCapacity)
+ : this()
  {
  Debug.Assert(maxCapacity >= 0);
 
  _first = _current = Array.Empty<T>();
- _buffers = default(ArrayBuilder<T[]>);
- _index = 0;
- _count = 0;
  _maxCapacity = maxCapacity;
  }
 
@@ -124,18 +164,18 @@ public void AddRange(IEnumerable<T> items)
  /// Copies the contents of this builder to the specified array.
  /// </summary>
  /// <param name="array">The destination array.</param>
- /// <param name="arrayIndex">The index in <see cref="array"/> to start copying.</param>
+ /// <param name="arrayIndex">The index in <see cref="array"/> to start copying to.</param>
  /// <param name="count">The number of items to copy.</param>
  public void CopyTo(T[] array, int arrayIndex, int count)
  {
  Debug.Assert(arrayIndex >= 0);
  Debug.Assert(count >= 0 && count <= Count);
  Debug.Assert(array?.Length - arrayIndex >= count);
 
- for (int i = -1; count > 0; i++)
+ for (int i = 0; count > 0; i++)
  {
  // Find the buffer we're copying from.
- T[] buffer = i < 0 ? _first : i < _buffers.Count ? _buffers[i] : _current;
+ T[] buffer = GetBuffer(index: i);
 
  // Copy until we satisfy count, or we reach the end of the buffer.
  int toCopy = Math.Min(count, buffer.Length);
@@ -146,6 +186,69 @@ public void CopyTo(T[] array, int arrayIndex, int count)
  arrayIndex += toCopy;
  }
  }
+
+ /// <summary>
+ /// Copies the contents of this builder to the specified array.
+ /// </summary>
+ /// <param name="position">The position in this builder to start copying from.</param>
+ /// <param name="array">The destination array.</param>
+ /// <param name="arrayIndex">The index in <see cref="array"/> to start copying to.</param>
+ /// <param name="count">The number of items to copy.</param>
+ /// <returns>The position in this builder that was copied up to.</returns>
+ public CopyPosition CopyTo(CopyPosition position, T[] array, int arrayIndex, int count)
+ {
+ Debug.Assert(arrayIndex >= 0);
+ Debug.Assert(count >= 0 && count <= Count);
+ Debug.Assert(array?.Length - arrayIndex >= count);
+
+ // Go through each buffer, which contains one 'row' of items.
+ // The index in each buffer is referred to as the 'column'.
+
+ /*
+ * Visual representation:
+ * 
+ * C0 C1 C2 .. C31 .. C63
+ * R0: [0] [1] [2] .. [31]
+ * R1: [32] [33] [34] .. [63]
+ * R2: [64] [65] [66] .. [95] .. [127]
+ */
+
+ int row = position.Row;
+ int column = position.Column;
+
+ for (; count > 0; row++, column = 0)
+ {
+ T[] buffer = GetBuffer(index: row);
+
+ // During this iteration, copy until we satisfy `count` or reach the
+ // end of the current buffer.
+ int copyCount = Math.Min(buffer.Length, count);
+
+ if (copyCount > 0)
+ {
+ Array.Copy(buffer, column, array, arrayIndex, copyCount);
+
+ arrayIndex += copyCount;
+ count -= copyCount;
+ column += copyCount;
+ }
+ }
+
+ return new CopyPosition(row: row, column: column);
+ }
+
+ /// <summary>
+ /// Retrieves the buffer at the specified index.
+ /// </summary>
+ /// <param name="index">The index of the buffer.</param>
+ public T[] GetBuffer(int index)
+ {
+ Debug.Assert(index >= 0 && index < _buffers.Count + 2);
+
+ return index == 0 ? _first :
+ index <= _buffers.Count ? _buffers[index - 1] :
+ _current;
+ }
 
  /// <summary>
  /// Adds an item to this builder.
@@ -159,21 +262,33 @@ public void CopyTo(T[] array, int arrayIndex, int count)
  public void SlowAdd(T item) => Add(item);
 
  /// <summary>
- /// Returns an array representation of this builder.
+ /// Creates an array from the contents of this builder.
  /// </summary>
  public T[] ToArray()
  {
- if (_count == _first.Length)
+ T[] array;
+ if (TryMove(out array))
  {
  // No resizing to do.
- return _first;
+ return array;
  }
 
- var array = new T[_count];
+ array = new T[_count];
  CopyTo(array, 0, _count);
  return array;
  }
 
+ /// <summary>
+ /// Attempts to transfer this builder into an array without copying.
+ /// </summary>
+ /// <param name="array">The transferred array, if the operation succeeded.</param>
+ /// <returns><c>true</c> if the operation succeeded; otherwise, <c>false</c>.</returns>
+ public bool TryMove(out T[] array)
+ {
+ array = _first;
+ return _count == _first.Length;
+ }
+
  private void AllocateBuffer()
  {
  // - On the first few adds, simply resize _first.