fix: refsactoring

codibre · Nov 16, 2024 · baf07d8 · baf07d8
1 parent 3e63ac4
commit baf07d8
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 25 deletions.
diff --git a/src/Codibre.EnumerableExtensions.Branching/AsyncBranchingBuilder.cs b/src/Codibre.EnumerableExtensions.Branching/AsyncBranchingBuilder.cs
@@ -8,13 +8,6 @@
 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(BranchRunOptions options)
-    {
-        var enumerator = source.GetAsyncEnumerator();
-        return LinkedNode<T>.Root(
-            async (c) => await enumerator.MoveNextAsync() ? new(enumerator.Current, c) : _null,
-            options
-        );
-    }
+        => LinkedNode<T>.Root(source.GetAsyncEnumerator(), options);
 }
diff --git a/src/Codibre.EnumerableExtensions.Branching/BranchingBuilder.cs b/src/Codibre.EnumerableExtensions.Branching/BranchingBuilder.cs
@@ -8,11 +8,5 @@ namespace Codibre.EnumerableExtensions.Branching;
 
 public sealed class BranchingBuilder<T>(IEnumerable<T> source) : BaseBranchingBuilder<T>
 {
-    internal override LinkedNode<T> Iterate(BranchRunOptions options)
-    {
-        var enumerator = source.GetEnumerator();
-        var getNext = (IBranchContext<T> c) => ValueTask.FromResult(LinkedNode<T>.New(enumerator, options, c));
-        BranchContext<T> context = new(getNext, BranchRunOptions.Yielder);
-        return LinkedNode<T>.Root(enumerator, options, context);
-    }
+    internal override LinkedNode<T> Iterate(BranchRunOptions options) => LinkedNode<T>.Root(source.GetEnumerator(), options);
 }
diff --git a/src/Codibre.EnumerableExtensions.Branching/Internal/LinkedNode.cs b/src/Codibre.EnumerableExtensions.Branching/Internal/LinkedNode.cs
@@ -3,15 +3,17 @@
 internal sealed record LinkedNode<T>
 {
     public T Value { get; }
-    public Lazy<ValueTask<LinkedNode<T>?>>? Next { get; private set; }
+    public Lazy<ValueTask<LinkedNode<T>?>> Next { get; private set; }
 
     public LinkedNode(T value, IBranchContext<T> context)
     {
         Value = value;
         Next = new(context.FillNext, LazyThreadSafetyMode.ExecutionAndPublication);
     }
 
+#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.
     private LinkedNode(T value) => Value = value;
+#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.
 
     public static LinkedNode<T>? New(IEnumerator<T> source, BranchRunOptions options, IBranchContext<T> context)
     {
@@ -24,17 +26,45 @@ public LinkedNode(T value, IBranchContext<T> context)
         return root;
     }
 
-    public static LinkedNode<T> Root(IEnumerator<T> source, BranchRunOptions options, IBranchContext<T> context)
+    public static async ValueTask<LinkedNode<T>?> New(IAsyncEnumerator<T> source, BranchRunOptions options, IBranchContext<T> context)
+    {
+        if (!await source.MoveNextAsync()) return null;
+        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!;
+        node.Next = new(context.FillNext, LazyThreadSafetyMode.ExecutionAndPublication);
+        return root;
+    }
+
+    public static LinkedNode<T> Root(IEnumerator<T> source, BranchRunOptions options)
         => new(default(T)!)
         {
-            Next = new(ValueTask.FromResult(New(source, options, context)))
+            Next = new(
+                    ValueTask.FromResult(
+                        New(
+                            source,
+                            options,
+                            new AsyncBranchContext<T>(
+                                (c) => ValueTask.FromResult(LinkedNode<T>.New(source, options, c))
+                            )
+                        )
+                    )
+                )
         };
 
-    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);
-    }
+    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
+                        )
+                    )
+                )
+        };
 }