Expand.cs 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596
  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.Threading;
  6. using System.Threading.Tasks;
  7. namespace System.Linq
  8. {
  9. public static partial class AsyncEnumerableEx
  10. {
  11. public static IAsyncEnumerable<TSource> Expand<TSource>(this IAsyncEnumerable<TSource> source, Func<TSource, IAsyncEnumerable<TSource>> selector)
  12. {
  13. if (source == null)
  14. throw Error.ArgumentNull(nameof(source));
  15. if (selector == null)
  16. throw Error.ArgumentNull(nameof(selector));
  17. return AsyncEnumerable.Create(Core);
  18. async IAsyncEnumerator<TSource> Core(CancellationToken cancellationToken)
  19. {
  20. var queue = new Queue<IAsyncEnumerable<TSource>>();
  21. queue.Enqueue(source);
  22. while (queue.Count > 0)
  23. {
  24. await foreach (var item in queue.Dequeue().WithCancellation(cancellationToken).ConfigureAwait(false))
  25. {
  26. queue.Enqueue(selector(item));
  27. yield return item;
  28. }
  29. }
  30. }
  31. }
  32. public static IAsyncEnumerable<TSource> Expand<TSource>(this IAsyncEnumerable<TSource> source, Func<TSource, ValueTask<IAsyncEnumerable<TSource>>> selector)
  33. {
  34. if (source == null)
  35. throw Error.ArgumentNull(nameof(source));
  36. if (selector == null)
  37. throw Error.ArgumentNull(nameof(selector));
  38. return AsyncEnumerable.Create(Core);
  39. async IAsyncEnumerator<TSource> Core(CancellationToken cancellationToken)
  40. {
  41. var queue = new Queue<IAsyncEnumerable<TSource>>();
  42. queue.Enqueue(source);
  43. while (queue.Count > 0)
  44. {
  45. await foreach (var item in queue.Dequeue().WithCancellation(cancellationToken).ConfigureAwait(false))
  46. {
  47. queue.Enqueue(await selector(item).ConfigureAwait(false));
  48. yield return item;
  49. }
  50. }
  51. }
  52. }
  53. #if !NO_DEEP_CANCELLATION
  54. public static IAsyncEnumerable<TSource> Expand<TSource>(this IAsyncEnumerable<TSource> source, Func<TSource, CancellationToken, ValueTask<IAsyncEnumerable<TSource>>> selector)
  55. {
  56. if (source == null)
  57. throw Error.ArgumentNull(nameof(source));
  58. if (selector == null)
  59. throw Error.ArgumentNull(nameof(selector));
  60. return AsyncEnumerable.Create(Core);
  61. async IAsyncEnumerator<TSource> Core(CancellationToken cancellationToken)
  62. {
  63. var queue = new Queue<IAsyncEnumerable<TSource>>();
  64. queue.Enqueue(source);
  65. while (queue.Count > 0)
  66. {
  67. await foreach (var item in queue.Dequeue().WithCancellation(cancellationToken).ConfigureAwait(false))
  68. {
  69. queue.Enqueue(await selector(item, cancellationToken).ConfigureAwait(false));
  70. yield return item;
  71. }
  72. }
  73. }
  74. }
  75. #endif
  76. }
  77. }