Merge pull request #25 from rameel/add-cancellation-token

Add CancellationToken parameter to FileTreeAsyncEnumerable

Author: rameel

Ramstack.Globbing/Traversal/FileTreeAsyncEnumerable.cs

@@ -13,6 +13,7 @@ namespace Ramstack.Globbing.Traversal;
 public sealed class FileTreeAsyncEnumerable<TEntry, TResult> : IAsyncEnumerable<TResult>
 {
     private readonly TEntry _directory;
+    private readonly CancellationToken _cancellationToken;
 
     /// <summary>
     /// Gets or sets the glob patterns to include in the enumeration.
@@ -59,40 +60,59 @@ public sealed class FileTreeAsyncEnumerable<TEntry, TResult> : IAsyncEnumerable<
     /// Initializes a new instance of the <see cref="FileTreeAsyncEnumerable{TEntry, TResult}"/> class.
     /// </summary>
     /// <param name="directory">The root directory to start the enumeration from.</param>
-    public FileTreeAsyncEnumerable(TEntry directory) =>
+    /// <param name="cancellationToken">An optional cancellation token that may be used to cancel the asynchronous iteration.</param>
+    public FileTreeAsyncEnumerable(TEntry directory, CancellationToken cancellationToken = default)
+    {
         _directory = directory;
+        _cancellationToken = cancellationToken;
+    }
 
     /// <inheritdoc />
-    IAsyncEnumerator<TResult> IAsyncEnumerable<TResult>.GetAsyncEnumerator(CancellationToken cancellationToken) =>
-        EnumerateAsync(cancellationToken).GetAsyncEnumerator(cancellationToken);
-
-    private async IAsyncEnumerable<TResult> EnumerateAsync([EnumeratorCancellation] CancellationToken cancellationToken)
+    IAsyncEnumerator<TResult> IAsyncEnumerable<TResult>.GetAsyncEnumerator(CancellationToken cancellationToken)
     {
-        var chars = ArrayPool<char>.Shared.Rent(512);
+        CancellationTokenSource? source = null;
 
-        var stack = new Stack<(TEntry Directory, string Path)>();
-        stack.Push((_directory, ""));
+        if (_cancellationToken != default)
+            cancellationToken = cancellationToken != default
+                ? (source = CancellationTokenSource.CreateLinkedTokenSource(_cancellationToken, cancellationToken)).Token
+                : _cancellationToken;
 
-        while (stack.TryPop(out var e))
-        {
-            await foreach (var entry in ChildrenSelector(e.Directory, cancellationToken))
-            {
-                var name = FileNameSelector(entry);
-                var fullName = FileTreeHelper.GetFullName(ref chars, e.Path, name);
+        return EnumerateAsync(source, cancellationToken).GetAsyncEnumerator(cancellationToken);
+    }
 
-                if (PathHelper.IsMatch(fullName, Excludes, Flags))
-                    continue;
+    private async IAsyncEnumerable<TResult> EnumerateAsync(CancellationTokenSource? source, [EnumeratorCancellation] CancellationToken cancellationToken)
+    {
+        var chars = ArrayPool<char>.Shared.Rent(512);
 
-                if (ShouldRecursePredicate == null || ShouldRecursePredicate(entry))
-                    if (PathHelper.IsPartialMatch(fullName, Patterns, Flags))
-                        stack.Push((entry, fullName.ToString()));
+        try
+        {
+            var stack = new Stack<(TEntry Directory, string Path)>();
+            stack.Push((_directory, ""));
 
-                if (ShouldIncludePredicate == null || ShouldIncludePredicate(entry))
-                    if (PathHelper.IsMatch(fullName, Patterns, Flags))
-                        yield return ResultSelector(entry);
+            while (stack.TryPop(out var e))
+            {
+                await foreach (var entry in ChildrenSelector(e.Directory, cancellationToken))
+                {
+                    var name = FileNameSelector(entry);
+                    var fullName = FileTreeHelper.GetFullName(ref chars, e.Path, name);
+
+                    if (PathHelper.IsMatch(fullName, Excludes, Flags))
+                        continue;
+
+                    if (ShouldRecursePredicate == null || ShouldRecursePredicate(entry))
+                        if (PathHelper.IsPartialMatch(fullName, Patterns, Flags))
+                            stack.Push((entry, fullName.ToString()));
+
+                    if (ShouldIncludePredicate == null || ShouldIncludePredicate(entry))
+                        if (PathHelper.IsMatch(fullName, Patterns, Flags))
+                            yield return ResultSelector(entry);
+                }
             }
         }
-
-        ArrayPool<char>.Shared.Return(chars);
+        finally
+        {
+            ArrayPool<char>.Shared.Return(chars);
+            source?.Dispose();
+        }
     }
 }