Merge pull request #7 from codibre/fix/minor-bugs

fix: fixing minor bugs
codibre · Nov 18, 2024 · b706980 · b706980
2 parents 945c617 + dce2360
commit b706980
Show file tree

Hide file tree

Showing 7 changed files with 45 additions and 25 deletions.
diff --git a/src/Codibre.EnumerableExtensions.Branching/BaseBranchingBuilder.cs b/src/Codibre.EnumerableExtensions.Branching/BaseBranchingBuilder.cs
@@ -22,5 +22,5 @@ public async Task Run(BranchRunOptions? options = null)
         await Task.WhenAll(_branches.Select(x => x(node.GetBranchedIterable())));
     }
 
-    internal abstract LinkedNode<T> Iterate(BranchRunOptions limit);
+    internal abstract LinkedNode<T> Iterate(BranchRunOptions options);
 }
diff --git a/src/Codibre.EnumerableExtensions.Branching/Internal/AsyncBranchContext.cs b/src/Codibre.EnumerableExtensions.Branching/Internal/AsyncBranchContext.cs
@@ -6,10 +6,5 @@ namespace Codibre.EnumerableExtensions.Branching.Internal;
 internal sealed record AsyncBranchContext<T>(Func<IBranchContext<T>, ValueTask<LinkedNode<T>?>> GetNext)
     : IBranchContext<T>
 {
-
-    public ValueTask<LinkedNode<T>?> FillNext()
-    {
-        var result = GetNext(this);
-        return result.IsCompleted ? new(Task.Run(() => result.Result)) : result;
-    }
+    public ValueTask<LinkedNode<T>?> FillNext() => GetNext(this).ResolveAsync();
 }
diff --git a/src/Codibre.EnumerableExtensions.Branching/Internal/BranchContext.cs b/src/Codibre.EnumerableExtensions.Branching/Internal/BranchContext.cs
@@ -11,7 +11,7 @@ internal sealed record BranchContext<T>(Func<IBranchContext<T>, ValueTask<Linked
     public ValueTask<LinkedNode<T>?> FillNext()
     {
         var result = GetNext(this);
-        if (result.IsCompleted) return ++_count <= options.Limit ? result : GetYielded(result.Result);
+        if (result.IsCompletedSuccessfully) return ++_count <= options.Limit ? result : GetYielded(result.Result);
         _count = 0;
         return result;
     }

diff --git a/src/Codibre.EnumerableExtensions.Branching/Internal/BranchedEnumerable.cs b/src/Codibre.EnumerableExtensions.Branching/Internal/BranchedEnumerable.cs
@@ -12,5 +12,9 @@ internal class BranchedEnumerable<T> : IAsyncEnumerable<T>
     private readonly BranchedEnumerator<T> _enumerator;
     public BranchedEnumerable(LinkedNode<T> root) => _enumerator = new(root);
 
-    public IAsyncEnumerator<T> GetAsyncEnumerator(CancellationToken cancellationToken = default) => _enumerator;
+    public IAsyncEnumerator<T> GetAsyncEnumerator(CancellationToken cancellationToken = default)
+    {
+        _enumerator.CancellationToken = cancellationToken;
+        return _enumerator;
+    }
 }
diff --git a/src/Codibre.EnumerableExtensions.Branching/Internal/BranchedEnumerator.cs b/src/Codibre.EnumerableExtensions.Branching/Internal/BranchedEnumerator.cs
@@ -10,6 +10,7 @@ namespace Codibre.EnumerableExtensions.Branching.Internal;
 internal class BranchedEnumerator<T> : IAsyncEnumerator<T>
 {
     private LinkedNode<T>? _node;
+    internal CancellationToken CancellationToken { get; set; }
 #pragma warning disable CS8618 // Non-nullable field must contain a non-null value when exiting constructor. Consider adding the 'required' modifier or declaring as nullable.
     public BranchedEnumerator(LinkedNode<T> root) => _node = root;
 #pragma warning restore CS8618 // Non-nullable field must contain a non-null value when exiting constructor. Consider adding the 'required' modifier or declaring as nullable.
@@ -19,6 +20,7 @@ internal class BranchedEnumerator<T> : IAsyncEnumerator<T>
 
     public async ValueTask<bool> MoveNextAsync()
     {
+        CancellationToken.ThrowIfCancellationRequested();
         if (_node is null) return false;
         _node = await _node.Next.Value;
         if (_node is null) return false;

diff --git a/src/Codibre.EnumerableExtensions.Branching/Internal/LinkedNode.cs b/src/Codibre.EnumerableExtensions.Branching/Internal/LinkedNode.cs
@@ -21,7 +21,12 @@ public LinkedNode(T value, IBranchContext<T> context)
         var preload = options.Limit;
         LinkedNode<T> root = new(source.Current);
         var node = root;
-        while (preload-- > 0 && source.MoveNext()) node = (node.Next = new(ValueTask.FromResult<LinkedNode<T>?>(new(source.Current)))).Value.Result!;
+        while (preload-- > 0 && source.MoveNext())
+        {
+            LinkedNode<T> current = new(source.Current);
+            node.Next = new(ValueTask.FromResult<LinkedNode<T>?>(current));
+            node = current;
+        }
         node.Next = new(context.FillNext, LazyThreadSafetyMode.ExecutionAndPublication);
         return root;
     }
@@ -32,7 +37,12 @@ public LinkedNode(T value, IBranchContext<T> context)
         var preload = options.Limit;
         LinkedNode<T> root = new(source.Current);
         var node = root;
-        while (preload-- > 0 && await source.MoveNextAsync()) node = (node.Next = new(ValueTask.FromResult<LinkedNode<T>?>(new(source.Current)))).Value.Result!;
+        while (preload-- > 0 && await source.MoveNextAsync())
+        {
+            LinkedNode<T> current = new(source.Current);
+            node.Next = new(ValueTask.FromResult<LinkedNode<T>?>(current));
+            node = current;
+        }
         node.Next = new(context.FillNext, LazyThreadSafetyMode.ExecutionAndPublication);
         return root;
     }
@@ -41,30 +51,30 @@ public static LinkedNode<T> Root(IEnumerator<T> source, BranchRunOptions options
         => new(default(T)!)
         {
             Next = new(
-                    ValueTask.FromResult(
-                        New(
-                            source,
-                            options,
-                            new AsyncBranchContext<T>(
-                                (c) => ValueTask.FromResult(LinkedNode<T>.New(source, options, c))
-                            )
+                ValueTask.FromResult(
+                    New(
+                        source,
+                        options,
+                        new AsyncBranchContext<T>(
+                            (c) => ValueTask.FromResult(LinkedNode<T>.New(source, options, c))
                         )
                     )
                 )
+            )
         };
 
     public static LinkedNode<T> Root(IAsyncEnumerator<T> source, BranchRunOptions options)
         => new(default(T)!)
         {
             Next = new(
-                    New(
-                        source,
-                        BranchRunOptions.Yielder,
-                        new BranchContext<T>(
-                            (c) => LinkedNode<T>.New(source, BranchRunOptions.Yielder, c),
-                            options
-                        )
+                New(
+                    source,
+                    BranchRunOptions.Yielder,
+                    new BranchContext<T>(
+                        (c) => LinkedNode<T>.New(source, BranchRunOptions.Yielder, c),
+                        options
                     )
                 )
+            )
         };
 }
diff --git a/src/Codibre.EnumerableExtensions.Branching/Internal/TaskExtensions.cs b/src/Codibre.EnumerableExtensions.Branching/Internal/TaskExtensions.cs
@@ -0,0 +1,9 @@
+namespace Codibre.EnumerableExtensions.Branching.Internal;
+
+internal static class TaskExtensions
+{
+    public static ValueTask<T> ResolveAsync<T>(this ValueTask<T> valueTask)
+        => valueTask.IsCompletedSuccessfully
+            ? new(Task.Run(() => valueTask!.Result))
+            : valueTask;
+}