AsyncEnumerableHelpers.cs 5.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112
  1. // Licensed to the .NET Foundation under one or more agreements.
  2. // The .NET Foundation licenses this file to you under the Apache 2.0 License.
  3. // See the LICENSE file in the project root for more information.
  4. using System.Linq;
  5. using System.Threading;
  6. using System.Threading.Tasks;
  7. namespace System.Collections.Generic
  8. {
  9. // Based on https://github.com/dotnet/corefx/blob/ec2685715b01d12f16b08d0dfa326649b12db8ec/src/Common/src/System/Collections/Generic/EnumerableHelpers.cs
  10. internal static class AsyncEnumerableHelpers
  11. {
  12. internal static async Task<T[]> ToArray<T>(IAsyncEnumerable<T> source, CancellationToken cancellationToken)
  13. {
  14. var result = await ToArrayWithLength(source, cancellationToken)
  15. .ConfigureAwait(false);
  16. Array.Resize(ref result.array, result.length);
  17. return result.array;
  18. }
  19. internal static async Task<ArrayWithLength<T>> ToArrayWithLength<T>(IAsyncEnumerable<T> source, CancellationToken cancellationToken)
  20. {
  21. var result = new ArrayWithLength<T>();
  22. // Check for short circuit optimizations. This one is very unlikely
  23. // but could be here as a group
  24. if (source is ICollection<T> ic)
  25. {
  26. var count = ic.Count;
  27. if (count != 0)
  28. {
  29. // Allocate an array of the desired size, then copy the elements into it. Note that this has the same
  30. // issue regarding concurrency as other existing collections like List<T>. If the collection size
  31. // concurrently changes between the array allocation and the CopyTo, we could end up either getting an
  32. // exception from overrunning the array (if the size went up) or we could end up not filling as many
  33. // items as 'count' suggests (if the size went down). This is only an issue for concurrent collections
  34. // that implement ICollection<T>, which as of .NET 4.6 is just ConcurrentDictionary<TKey, TValue>.
  35. result.array = new T[count];
  36. ic.CopyTo(result.array, 0);
  37. result.length = count;
  38. return result;
  39. }
  40. }
  41. else
  42. {
  43. using (var en = source.GetEnumerator())
  44. {
  45. if (await en.MoveNext(cancellationToken)
  46. .ConfigureAwait(false))
  47. {
  48. const int DefaultCapacity = 4;
  49. var arr = new T[DefaultCapacity];
  50. arr[0] = en.Current;
  51. var count = 1;
  52. while (await en.MoveNext(cancellationToken)
  53. .ConfigureAwait(false))
  54. {
  55. if (count == arr.Length)
  56. {
  57. // MaxArrayLength is defined in Array.MaxArrayLength and in gchelpers in CoreCLR.
  58. // It represents the maximum number of elements that can be in an array where
  59. // the size of the element is greater than one byte; a separate, slightly larger constant,
  60. // is used when the size of the element is one.
  61. const int MaxArrayLength = 0x7FEFFFFF;
  62. // This is the same growth logic as in List<T>:
  63. // If the array is currently empty, we make it a default size. Otherwise, we attempt to
  64. // double the size of the array. Doubling will overflow once the size of the array reaches
  65. // 2^30, since doubling to 2^31 is 1 larger than Int32.MaxValue. In that case, we instead
  66. // constrain the length to be MaxArrayLength (this overflow check works because of the
  67. // cast to uint). Because a slightly larger constant is used when T is one byte in size, we
  68. // could then end up in a situation where arr.Length is MaxArrayLength or slightly larger, such
  69. // that we constrain newLength to be MaxArrayLength but the needed number of elements is actually
  70. // larger than that. For that case, we then ensure that the newLength is large enough to hold
  71. // the desired capacity. This does mean that in the very rare case where we've grown to such a
  72. // large size, each new element added after MaxArrayLength will end up doing a resize.
  73. var newLength = count << 1;
  74. if ((uint)newLength > MaxArrayLength)
  75. {
  76. newLength = MaxArrayLength <= count ? count + 1 : MaxArrayLength;
  77. }
  78. Array.Resize(ref arr, newLength);
  79. }
  80. arr[count++] = en.Current;
  81. }
  82. result.length = count;
  83. result.array = arr;
  84. return result;
  85. }
  86. }
  87. }
  88. result.length = 0;
  89. #if NO_ARRAY_EMPTY
  90. result.array = EmptyArray<T>.Value;
  91. #else
  92. result.array = Array.Empty<T>();
  93. #endif
  94. return result;
  95. }
  96. internal struct ArrayWithLength<T>
  97. {
  98. public T[] array;
  99. public int length;
  100. }
  101. }
  102. }