feat: creating branch options

codibre · Nov 15, 2024 · ec2ee99 · ec2ee99
1 parent 0e3a7a1
commit ec2ee99
Show file tree

Hide file tree

Showing 10 changed files with 168 additions and 18 deletions.
diff --git a/src/Codibre.EnumerableExtensions.Branching/AsyncBranchingBuilder.cs b/src/Codibre.EnumerableExtensions.Branching/AsyncBranchingBuilder.cs
@@ -9,12 +9,12 @@ namespace Codibre.EnumerableExtensions.Branching;
 public sealed class AsyncBranchingBuilder<T>(IAsyncEnumerable<T> source) : BaseBranchingBuilder<T>
 {
     private static readonly LinkedNode<T>? _null = null;
-    internal override LinkedNode<T> Iterate(int branchCount)
+    internal override LinkedNode<T> Iterate(BranchRunOptions options)
     {
         var enumerator = source.GetAsyncEnumerator();
-        return new(enumerator.Current, new(
+        return LinkedNode<T>.Root(
             async (c) => await enumerator.MoveNextAsync() ? new(enumerator.Current, c) : _null,
-            branchCount
-        ));
+            options
+        );
     }
 }
diff --git a/src/Codibre.EnumerableExtensions.Branching/BaseBranchingBuilder.cs b/src/Codibre.EnumerableExtensions.Branching/BaseBranchingBuilder.cs
@@ -16,11 +16,11 @@ internal BaseBranchingBuilder<T> Add(Func<IAsyncEnumerable<T>, Task> branch)
         return this;
     }
 
-    public async Task Run()
+    public async Task Run(BranchRunOptions? options = null)
     {
-        var node = Iterate(_branches.Count);
+        var node = Iterate(options ?? BranchRunOptions.Default);
         await Task.WhenAll(_branches.Select(x => x(node.GetBranchedIterable())));
     }
 
-    internal abstract LinkedNode<T> Iterate(int branchCount);
+    internal abstract LinkedNode<T> Iterate(BranchRunOptions limit);
 }
diff --git a/src/Codibre.EnumerableExtensions.Branching/BranchRunOptions.cs b/src/Codibre.EnumerableExtensions.Branching/BranchRunOptions.cs
@@ -0,0 +1,8 @@
+
+namespace Codibre.EnumerableExtensions.Branching;
+
+public readonly struct BranchRunOptions(int limit)
+{
+    public static readonly BranchRunOptions Default = new(ushort.MaxValue / 4);
+    public int Limit => limit;
+}
diff --git a/src/Codibre.EnumerableExtensions.Branching/BranchingBuilder.cs b/src/Codibre.EnumerableExtensions.Branching/BranchingBuilder.cs
@@ -8,12 +8,12 @@ namespace Codibre.EnumerableExtensions.Branching;
 
 public sealed class BranchingBuilder<T>(IEnumerable<T> source) : BaseBranchingBuilder<T>
 {
-    internal override LinkedNode<T> Iterate(int branchCount)
+    internal override LinkedNode<T> Iterate(BranchRunOptions options)
     {
         var enumerator = source.GetEnumerator();
-        return new(default!, new(
+        return LinkedNode<T>.Root(
             (c) => new(enumerator.MoveNext() ? new LinkedNode<T>(enumerator.Current, c) : null),
-            branchCount
-        ));
+            options
+        );
     }
 }
diff --git a/src/Codibre.EnumerableExtensions.Branching/Internal/AsyncBranchContext.cs b/src/Codibre.EnumerableExtensions.Branching/Internal/AsyncBranchContext.cs
@@ -0,0 +1,15 @@
+using System.Runtime.CompilerServices;
+using System.Xml.XPath;
+
+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;
+    }
+}
diff --git a/src/Codibre.EnumerableExtensions.Branching/Internal/BranchContext.cs b/src/Codibre.EnumerableExtensions.Branching/Internal/BranchContext.cs
@@ -3,18 +3,23 @@
 
 namespace Codibre.EnumerableExtensions.Branching.Internal;
 
-internal sealed record BranchContext<T>(Func<BranchContext<T>, ValueTask<LinkedNode<T>?>> GetNext, int _branchCount)
+internal sealed record BranchContext<T>(Func<IBranchContext<T>, ValueTask<LinkedNode<T>?>> GetNext, BranchRunOptions options)
+    : IBranchContext<T>
 {
     private ushort _count = 0;
-    private readonly ushort _limit = ushort.MaxValue / 4;
 
-    internal ValueTask<LinkedNode<T>?> FillNext()
-        => ++_count <= _limit ? GetNext(this) : GetYielded();
+    public ValueTask<LinkedNode<T>?> FillNext()
+    {
+        var result = GetNext(this);
+        if (result.IsCompleted) return ++_count <= options.Limit ? result : GetYielded(result.Result);
+        _count = 0;
+        return result;
+    }
 
-    private async ValueTask<LinkedNode<T>?> GetYielded()
+    private async ValueTask<LinkedNode<T>?> GetYielded(LinkedNode<T>? result)
     {
         _count = 0;
         await Task.Yield();
-        return await GetNext(this);
+        return result;
     }
 }
diff --git a/src/Codibre.EnumerableExtensions.Branching/Internal/IBranchContext.cs b/src/Codibre.EnumerableExtensions.Branching/Internal/IBranchContext.cs
@@ -0,0 +1,9 @@
+using System.Runtime.CompilerServices;
+using System.Xml.XPath;
+
+namespace Codibre.EnumerableExtensions.Branching.Internal;
+
+internal interface IBranchContext<T>
+{
+    ValueTask<LinkedNode<T>?> FillNext();
+}
diff --git a/src/Codibre.EnumerableExtensions.Branching/Internal/LinkedNode.cs b/src/Codibre.EnumerableExtensions.Branching/Internal/LinkedNode.cs
@@ -1,6 +1,14 @@
 namespace Codibre.EnumerableExtensions.Branching.Internal;
 
-internal sealed record LinkedNode<T>(T Value, BranchContext<T> Context)
+internal sealed record LinkedNode<T>(T Value, IBranchContext<T> Context)
 {
     public Lazy<ValueTask<LinkedNode<T>?>> Next { get; } = new(Context.FillNext, LazyThreadSafetyMode.ExecutionAndPublication);
+
+    public static LinkedNode<T> Root(Func<IBranchContext<T>, ValueTask<LinkedNode<T>?>> getNext, BranchRunOptions options)
+    {
+        IBranchContext<T> context = options.Limit <= 1
+            ? new AsyncBranchContext<T>(getNext)
+            : new BranchContext<T>(getNext, options);
+        return new(default!, context);
+    }
 }
diff --git a/test/Codibre.EnumerableExtensions.Branching.Test/AsyncEnumerableExtensionsTest.cs b/test/Codibre.EnumerableExtensions.Branching.Test/AsyncEnumerableExtensionsTest.cs
@@ -102,4 +102,76 @@ await enumerable.Branch()
         var refValue = steps[0];
         steps.TakeWhile((x) => x == refValue).Count().Should().BeLessThan(total);
     }
+
+    [Fact]
+    public async Task Should_Intercalate_The_Steps_Between_Every_Branch_When_EnumerableIsTrulyAsync()
+    {
+        // Arrange
+        var total = 1000;
+        var list = Enumerable.Range(0, total)
+            .ToAsyncEnumerable()
+            .SelectAwait(async (x) =>
+            {
+                await Task.Yield();
+                return await Task.Run(() => x);
+            });
+        List<int> steps = [];
+        var enumerable = Op(list);
+
+        // Act
+        await enumerable.Branch()
+            .Add(x => x.Select((x) =>
+                {
+                    steps.Add(1);
+                    return x;
+                }).ToArrayAsync(), out var a)
+            .Add(x => x.Select((x) =>
+                {
+                    steps.Add(2);
+                    return x;
+                }).ToArrayAsync(), out var b)
+            .Add(x => x.Select((x) =>
+                {
+                    steps.Add(3);
+                    return x;
+                }).ToArrayAsync(), out var c)
+            .Run();
+
+        // Assert
+        var refValue = steps[0];
+        steps.TakeWhile((x) => x == refValue).Count().Should().BeLessThan(total);
+    }
+
+    [Fact]
+    public async Task Should_Intercalate_The_Steps_Between_Every_Branch_When_Using_AsyncContextBranch()
+    {
+        // Arrange
+        var total = 100;
+        var list = Enumerable.Range(0, total).ToAsyncEnumerable();
+        List<int> steps = [];
+        var enumerable = Op(list);
+
+        // Act
+        await enumerable.Branch()
+            .Add(x => x.Select((x) =>
+                {
+                    steps.Add(1);
+                    return x;
+                }).ToArrayAsync(), out var a)
+            .Add(x => x.Select((x) =>
+                {
+                    steps.Add(2);
+                    return x;
+                }).ToArrayAsync(), out var b)
+            .Add(x => x.Select((x) =>
+                {
+                    steps.Add(3);
+                    return x;
+                }).ToArrayAsync(), out var c)
+            .Run(new(1));
+
+        // Assert
+        var refValue = steps[0];
+        steps.TakeWhile((x) => x == refValue).Count().Should().BeLessThan(total);
+    }
 }
diff --git a/test/Codibre.EnumerableExtensions.Branching.Test/EnumerableExtensionsTest.cs b/test/Codibre.EnumerableExtensions.Branching.Test/EnumerableExtensionsTest.cs
@@ -102,4 +102,37 @@ await enumerable.Branch()
         var refValue = steps[0];
         steps.TakeWhile((x) => x == refValue).Count().Should().BeLessThan(total);
     }
+
+    [Fact]
+    public async Task Should_Intercalate_The_Steps_Between_Every_Branch_When_Using_AsyncContextBranch()
+    {
+        // Arrange
+        var total = 100000;
+        var list = Enumerable.Range(0, total);
+        List<int> steps = [];
+        var enumerable = Op(list);
+
+        // Act
+        await enumerable.Branch()
+            .Add(x => x.Select((x) =>
+                {
+                    steps.Add(1);
+                    return x;
+                }).ToArrayAsync(), out var a)
+            .Add(x => x.Select((x) =>
+                {
+                    steps.Add(2);
+                    return x;
+                }).ToArrayAsync(), out var b)
+            .Add(x => x.Select((x) =>
+                {
+                    steps.Add(3);
+                    return x;
+                }).ToArrayAsync(), out var c)
+            .Run(new(1));
+
+        // Assert
+        var refValue = steps[0];
+        steps.TakeWhile((x) => x == refValue).Count().Should().BeLessThan(total);
+    }
 }