Concatenate.cs 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359
  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 AsyncEnumerable
  11. {
  12. public static IAsyncEnumerable<TSource> Concat<TSource>(this IAsyncEnumerable<TSource> first, IAsyncEnumerable<TSource> second)
  13. {
  14. if (first == null)
  15. throw new ArgumentNullException(nameof(first));
  16. if (second == null)
  17. throw new ArgumentNullException(nameof(second));
  18. var concatFirst = first as ConcatAsyncIterator<TSource>;
  19. return concatFirst != null ?
  20. concatFirst.Concat(second) :
  21. new Concat2AsyncIterator<TSource>(first, second);
  22. }
  23. public static IAsyncEnumerable<TSource> Concat<TSource>(this IEnumerable<IAsyncEnumerable<TSource>> sources)
  24. {
  25. if (sources == null)
  26. throw new ArgumentNullException(nameof(sources));
  27. return sources.Concat_();
  28. }
  29. public static IAsyncEnumerable<TSource> Concat<TSource>(params IAsyncEnumerable<TSource>[] sources)
  30. {
  31. if (sources == null)
  32. throw new ArgumentNullException(nameof(sources));
  33. return sources.Concat_();
  34. }
  35. private static IAsyncEnumerable<TSource> Concat_<TSource>(this IEnumerable<IAsyncEnumerable<TSource>> sources)
  36. {
  37. return new ConcatEnumerableAsyncIterator<TSource>(sources);
  38. }
  39. private sealed class ConcatEnumerableAsyncIterator<TSource> : AsyncIterator<TSource>
  40. {
  41. private readonly IEnumerable<IAsyncEnumerable<TSource>> source;
  42. public ConcatEnumerableAsyncIterator(IEnumerable<IAsyncEnumerable<TSource>> source)
  43. {
  44. Debug.Assert(source != null);
  45. this.source = source;
  46. }
  47. public override AsyncIterator<TSource> Clone()
  48. {
  49. return new ConcatEnumerableAsyncIterator<TSource>(source);
  50. }
  51. public override void Dispose()
  52. {
  53. if (outerEnumerator != null)
  54. {
  55. outerEnumerator.Dispose();
  56. outerEnumerator = null;
  57. }
  58. if (currentEnumerator != null)
  59. {
  60. currentEnumerator.Dispose();
  61. currentEnumerator = null;
  62. }
  63. base.Dispose();
  64. }
  65. // State machine vars
  66. private IEnumerator<IAsyncEnumerable<TSource>> outerEnumerator;
  67. private IAsyncEnumerator<TSource> currentEnumerator;
  68. private int mode;
  69. const int State_OuterNext = 1;
  70. const int State_While = 4;
  71. protected override async Task<bool> MoveNextCore(CancellationToken cancellationToken)
  72. {
  73. switch (state)
  74. {
  75. case AsyncIteratorState.Allocated:
  76. outerEnumerator = source.GetEnumerator();
  77. mode = State_OuterNext;
  78. state = AsyncIteratorState.Iterating;
  79. goto case AsyncIteratorState.Iterating;
  80. case AsyncIteratorState.Iterating:
  81. switch (mode)
  82. {
  83. case State_OuterNext:
  84. if (outerEnumerator.MoveNext())
  85. {
  86. // make sure we dispose the previous one if we're about to replace it
  87. currentEnumerator?.Dispose();
  88. currentEnumerator = outerEnumerator.Current.GetEnumerator();
  89. mode = State_While;
  90. goto case State_While;
  91. }
  92. break;
  93. case State_While:
  94. if (await currentEnumerator.MoveNext(cancellationToken)
  95. .ConfigureAwait(false))
  96. {
  97. current = currentEnumerator.Current;
  98. return true;
  99. }
  100. // No more on the inner enumerator, move to the next outer
  101. goto case State_OuterNext;
  102. }
  103. Dispose();
  104. break;
  105. }
  106. return false;
  107. }
  108. }
  109. private sealed class Concat2AsyncIterator<TSource> : ConcatAsyncIterator<TSource>
  110. {
  111. private readonly IAsyncEnumerable<TSource> first;
  112. private readonly IAsyncEnumerable<TSource> second;
  113. internal Concat2AsyncIterator(IAsyncEnumerable<TSource> first, IAsyncEnumerable<TSource> second)
  114. {
  115. Debug.Assert(first != null);
  116. Debug.Assert(second != null);
  117. this.first = first;
  118. this.second = second;
  119. }
  120. public override AsyncIterator<TSource> Clone()
  121. {
  122. return new Concat2AsyncIterator<TSource>(first, second);
  123. }
  124. internal override ConcatAsyncIterator<TSource> Concat(IAsyncEnumerable<TSource> next)
  125. {
  126. return new ConcatNAsyncIterator<TSource>(this, next, 2);
  127. }
  128. internal override IAsyncEnumerable<TSource> GetAsyncEnumerable(int index)
  129. {
  130. switch (index)
  131. {
  132. case 0:
  133. return first;
  134. case 1:
  135. return second;
  136. default:
  137. return null;
  138. }
  139. }
  140. }
  141. private abstract class ConcatAsyncIterator<TSource> : AsyncIterator<TSource>, IIListProvider<TSource>
  142. {
  143. private int counter;
  144. private IAsyncEnumerator<TSource> enumerator;
  145. public Task<TSource[]> ToArrayAsync(CancellationToken cancellationToken)
  146. {
  147. return AsyncEnumerableHelpers.ToArray(this, cancellationToken);
  148. }
  149. public async Task<List<TSource>> ToListAsync(CancellationToken cancellationToken)
  150. {
  151. var list = new List<TSource>();
  152. for (var i = 0;; i++)
  153. {
  154. var source = GetAsyncEnumerable(i);
  155. if (source == null)
  156. {
  157. break;
  158. }
  159. using (var e = source.GetEnumerator())
  160. {
  161. while (await e.MoveNext(cancellationToken)
  162. .ConfigureAwait(false))
  163. {
  164. list.Add(e.Current);
  165. }
  166. }
  167. }
  168. return list;
  169. }
  170. public async Task<int> GetCountAsync(bool onlyIfCheap, CancellationToken cancellationToken)
  171. {
  172. if (onlyIfCheap)
  173. {
  174. return -1;
  175. }
  176. var count = 0;
  177. for (var i = 0;; i++)
  178. {
  179. var source = GetAsyncEnumerable(i);
  180. if (source == null)
  181. {
  182. break;
  183. }
  184. checked
  185. {
  186. count += await source.Count(cancellationToken).ConfigureAwait(false);
  187. }
  188. }
  189. return count;
  190. }
  191. public override void Dispose()
  192. {
  193. if (enumerator != null)
  194. {
  195. enumerator.Dispose();
  196. enumerator = null;
  197. }
  198. base.Dispose();
  199. }
  200. protected override async Task<bool> MoveNextCore(CancellationToken cancellationToken)
  201. {
  202. if (state == AsyncIteratorState.Allocated)
  203. {
  204. enumerator = GetAsyncEnumerable(0)
  205. .GetEnumerator();
  206. state = AsyncIteratorState.Iterating;
  207. counter = 2;
  208. }
  209. if (state == AsyncIteratorState.Iterating)
  210. {
  211. while (true)
  212. {
  213. if (await enumerator.MoveNext(cancellationToken)
  214. .ConfigureAwait(false))
  215. {
  216. current = enumerator.Current;
  217. return true;
  218. }
  219. // note, this is simply to match the logic of
  220. // https://github.com/dotnet/corefx/blob/ec2685715b01d12f16b08d0dfa326649b12db8ec/src/system.linq/src/system/linq/concatenate.cs#L173-L173
  221. var next = GetAsyncEnumerable(counter++ - 1);
  222. if (next != null)
  223. {
  224. enumerator.Dispose();
  225. enumerator = next.GetEnumerator();
  226. continue;
  227. }
  228. Dispose();
  229. break;
  230. }
  231. }
  232. return false;
  233. }
  234. internal abstract ConcatAsyncIterator<TSource> Concat(IAsyncEnumerable<TSource> next);
  235. internal abstract IAsyncEnumerable<TSource> GetAsyncEnumerable(int index);
  236. }
  237. // To handle chains of >= 3 sources, we chain the concat iterators together and allow
  238. // GetEnumerable to fetch enumerables from the previous sources. This means that rather
  239. // than each MoveNext/Current calls having to traverse all of the previous sources, we
  240. // only have to traverse all of the previous sources once per chained enumerable. An
  241. // alternative would be to use an array to store all of the enumerables, but this has
  242. // a much better memory profile and without much additional run-time cost.
  243. private sealed class ConcatNAsyncIterator<TSource> : ConcatAsyncIterator<TSource>
  244. {
  245. private readonly IAsyncEnumerable<TSource> next;
  246. private readonly int nextIndex;
  247. private readonly ConcatAsyncIterator<TSource> previousConcat;
  248. internal ConcatNAsyncIterator(ConcatAsyncIterator<TSource> previousConcat, IAsyncEnumerable<TSource> next, int nextIndex)
  249. {
  250. Debug.Assert(previousConcat != null);
  251. Debug.Assert(next != null);
  252. Debug.Assert(nextIndex >= 2);
  253. this.previousConcat = previousConcat;
  254. this.next = next;
  255. this.nextIndex = nextIndex;
  256. }
  257. public override AsyncIterator<TSource> Clone()
  258. {
  259. return new ConcatNAsyncIterator<TSource>(previousConcat, next, nextIndex);
  260. }
  261. internal override ConcatAsyncIterator<TSource> Concat(IAsyncEnumerable<TSource> next)
  262. {
  263. if (nextIndex == int.MaxValue - 2)
  264. {
  265. // In the unlikely case of this many concatenations, if we produced a ConcatNIterator
  266. // with int.MaxValue then state would overflow before it matched it's index.
  267. // So we use the naïve approach of just having a left and right sequence.
  268. return new Concat2AsyncIterator<TSource>(this, next);
  269. }
  270. return new ConcatNAsyncIterator<TSource>(this, next, nextIndex + 1);
  271. }
  272. internal override IAsyncEnumerable<TSource> GetAsyncEnumerable(int index)
  273. {
  274. if (index > nextIndex)
  275. {
  276. return null;
  277. }
  278. // Walk back through the chain of ConcatNIterators looking for the one
  279. // that has its _nextIndex equal to index. If we don't find one, then it
  280. // must be prior to any of them, so call GetEnumerable on the previous
  281. // Concat2Iterator. This avoids a deep recursive call chain.
  282. var current = this;
  283. while (true)
  284. {
  285. if (index == current.nextIndex)
  286. {
  287. return current.next;
  288. }
  289. if (current.previousConcat is ConcatNAsyncIterator<TSource> prevN)
  290. {
  291. current = prevN;
  292. continue;
  293. }
  294. Debug.Assert(current.previousConcat is Concat2AsyncIterator<TSource>);
  295. Debug.Assert(index == 0 || index == 1);
  296. return current.previousConcat.GetAsyncEnumerable(index);
  297. }
  298. }
  299. }
  300. }
  301. }