Buffer.cs 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151
  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.Collections.Generic;
  5. using System.Diagnostics;
  6. using System.Threading;
  7. using System.Threading.Tasks;
  8. namespace System.Linq
  9. {
  10. public static partial class AsyncEnumerableEx
  11. {
  12. public static IAsyncEnumerable<IList<TSource>> Buffer<TSource>(this IAsyncEnumerable<TSource> source, int count)
  13. {
  14. if (source == null)
  15. {
  16. throw new ArgumentNullException(nameof(source));
  17. }
  18. if (count <= 0)
  19. {
  20. throw new ArgumentOutOfRangeException(nameof(count));
  21. }
  22. return new BufferAsyncIterator<TSource>(source, count, count);
  23. }
  24. public static IAsyncEnumerable<IList<TSource>> Buffer<TSource>(this IAsyncEnumerable<TSource> source, int count, int skip)
  25. {
  26. if (source == null)
  27. {
  28. throw new ArgumentNullException(nameof(source));
  29. }
  30. if (count <= 0)
  31. {
  32. throw new ArgumentOutOfRangeException(nameof(count));
  33. }
  34. if (skip <= 0)
  35. {
  36. throw new ArgumentOutOfRangeException(nameof(skip));
  37. }
  38. return new BufferAsyncIterator<TSource>(source, count, skip);
  39. }
  40. private sealed class BufferAsyncIterator<TSource> : AsyncIterator<IList<TSource>>
  41. {
  42. private readonly int count;
  43. private readonly int skip;
  44. private readonly IAsyncEnumerable<TSource> source;
  45. private Queue<IList<TSource>> buffers;
  46. private IAsyncEnumerator<TSource> enumerator;
  47. private int index;
  48. private bool stopped;
  49. public BufferAsyncIterator(IAsyncEnumerable<TSource> source, int count, int skip)
  50. {
  51. Debug.Assert(source != null);
  52. this.source = source;
  53. this.count = count;
  54. this.skip = skip;
  55. }
  56. public override AsyncIterator<IList<TSource>> Clone()
  57. {
  58. return new BufferAsyncIterator<TSource>(source, count, skip);
  59. }
  60. public override async ValueTask DisposeAsync()
  61. {
  62. if (enumerator != null)
  63. {
  64. await enumerator.DisposeAsync().ConfigureAwait(false);
  65. enumerator = null;
  66. }
  67. buffers = null;
  68. await base.DisposeAsync().ConfigureAwait(false);
  69. }
  70. protected override async ValueTask<bool> MoveNextCore(CancellationToken cancellationToken)
  71. {
  72. switch (state)
  73. {
  74. case AsyncIteratorState.Allocated:
  75. enumerator = source.GetAsyncEnumerator(cancellationToken);
  76. buffers = new Queue<IList<TSource>>();
  77. index = 0;
  78. stopped = false;
  79. state = AsyncIteratorState.Iterating;
  80. goto case AsyncIteratorState.Iterating;
  81. case AsyncIteratorState.Iterating:
  82. while (true)
  83. {
  84. if (!stopped)
  85. {
  86. if (await enumerator.MoveNextAsync().ConfigureAwait(false))
  87. {
  88. var item = enumerator.Current;
  89. if (index++ % skip == 0)
  90. {
  91. buffers.Enqueue(new List<TSource>(count));
  92. }
  93. foreach (var buffer in buffers)
  94. {
  95. buffer.Add(item);
  96. }
  97. if (buffers.Count > 0 && buffers.Peek().Count == count)
  98. {
  99. current = buffers.Dequeue();
  100. return true;
  101. }
  102. continue; // loop
  103. }
  104. stopped = true;
  105. await enumerator.DisposeAsync().ConfigureAwait(false);
  106. enumerator = null;
  107. continue; // loop
  108. }
  109. if (buffers.Count > 0)
  110. {
  111. current = buffers.Dequeue();
  112. return true;
  113. }
  114. break; // exit the while
  115. }
  116. break; // case
  117. }
  118. await DisposeAsync().ConfigureAwait(false);
  119. return false;
  120. }
  121. }
  122. }
  123. }