diff --git a/RELEASE_NOTES.md b/RELEASE_NOTES.md index 1544f75..ed2b48e 100644 --- a/RELEASE_NOTES.md +++ b/RELEASE_NOTES.md @@ -1,5 +1,6 @@ ### 4.7.0 +* Added `AsyncSeq.splitAt` — splits a sequence at the given index, returning the first `count` elements as an array and the remaining elements as a new `AsyncSeq`. Mirrors `Seq.splitAt`. The source is enumerated once. * Added `AsyncSeq.removeAt` — returns a new sequence with the element at the specified index removed, mirroring `Seq.removeAt`. * Added `AsyncSeq.updateAt` — returns a new sequence with the element at the specified index replaced by a given value, mirroring `Seq.updateAt`. * Added `AsyncSeq.insertAt` — returns a new sequence with a value inserted before the element at the specified index (or appended if the index equals the sequence length), mirroring `Seq.insertAt`. diff --git a/src/FSharp.Control.AsyncSeq/AsyncSeq.fs b/src/FSharp.Control.AsyncSeq/AsyncSeq.fs index 4ce71be..efc25c4 100644 --- a/src/FSharp.Control.AsyncSeq/AsyncSeq.fs +++ b/src/FSharp.Control.AsyncSeq/AsyncSeq.fs @@ -1841,6 +1841,35 @@ module AsyncSeq = let tail (source : AsyncSeq<'T>) : AsyncSeq<'T> = skip 1 source + /// Splits an async sequence at the given index, returning the first `count` elements as an array + /// and the remaining elements as a new AsyncSeq. The source is enumerated once. + let splitAt (count: int) (source: AsyncSeq<'T>) : Async<'T array * AsyncSeq<'T>> = async { + if count < 0 then invalidArg "count" "must be non-negative" + let ie = source.GetEnumerator() + let ra = ResizeArray<'T>() + let! m = ie.MoveNext() + let b = ref m + while b.Value.IsSome && ra.Count < count do + ra.Add b.Value.Value + let! next = ie.MoveNext() + b := next + let first = ra.ToArray() + let rest = + if b.Value.IsNone then + ie.Dispose() + empty<'T> + else + let cur = ref b.Value + asyncSeq { + try + while cur.Value.IsSome do + yield cur.Value.Value + let! next = ie.MoveNext() + cur := next + finally + ie.Dispose() } + return first, rest } + let toArrayAsync (source : AsyncSeq<'T>) : Async<'T[]> = async { let ra = (new ResizeArray<_>()) use ie = source.GetEnumerator() diff --git a/src/FSharp.Control.AsyncSeq/AsyncSeq.fsi b/src/FSharp.Control.AsyncSeq/AsyncSeq.fsi index 914b9c8..17841a4 100644 --- a/src/FSharp.Control.AsyncSeq/AsyncSeq.fsi +++ b/src/FSharp.Control.AsyncSeq/AsyncSeq.fsi @@ -599,6 +599,11 @@ module AsyncSeq = /// Returns an empty sequence if the source is empty. val tail : source:AsyncSeq<'T> -> AsyncSeq<'T> + /// Splits an async sequence at the given index. Returns an async computation that yields + /// the first `count` elements as an array and the remaining elements as a new AsyncSeq. + /// The source is enumerated once; the returned AsyncSeq lazily produces the remainder. + val splitAt : count:int -> source:AsyncSeq<'T> -> Async<'T array * AsyncSeq<'T>> + /// Creates an async computation which iterates the AsyncSeq and collects the output into an array. val toArrayAsync : source:AsyncSeq<'T> -> Async<'T []> diff --git a/tests/FSharp.Control.AsyncSeq.Tests/AsyncSeqTests.fs b/tests/FSharp.Control.AsyncSeq.Tests/AsyncSeqTests.fs index 7e163c4..f8c956d 100644 --- a/tests/FSharp.Control.AsyncSeq.Tests/AsyncSeqTests.fs +++ b/tests/FSharp.Control.AsyncSeq.Tests/AsyncSeqTests.fs @@ -3409,6 +3409,50 @@ let ``AsyncSeq.sortWith returns empty array for empty sequence`` () = let result = AsyncSeq.sortWith compare AsyncSeq.empty Assert.AreEqual([||], result) +[] +let ``AsyncSeq.splitAt splits a sequence at the given index`` () = + let source = asyncSeq { yield 1; yield 2; yield 3; yield 4; yield 5 } + let first, rest = AsyncSeq.splitAt 3 source |> Async.RunSynchronously + let restArr = rest |> AsyncSeq.toArrayAsync |> Async.RunSynchronously + Assert.AreEqual([| 1; 2; 3 |], first) + Assert.AreEqual([| 4; 5 |], restArr) + +[] +let ``AsyncSeq.splitAt with count=0 returns empty array and full rest`` () = + let source = asyncSeq { yield 10; yield 20 } + let first, rest = AsyncSeq.splitAt 0 source |> Async.RunSynchronously + let restArr = rest |> AsyncSeq.toArrayAsync |> Async.RunSynchronously + Assert.AreEqual([||], first) + Assert.AreEqual([| 10; 20 |], restArr) + +[] +let ``AsyncSeq.splitAt with count >= length returns all elements in first and empty rest`` () = + let source = asyncSeq { yield 1; yield 2; yield 3 } + let first, rest = AsyncSeq.splitAt 10 source |> Async.RunSynchronously + let restArr = rest |> AsyncSeq.toArrayAsync |> Async.RunSynchronously + Assert.AreEqual([| 1; 2; 3 |], first) + Assert.AreEqual([||], restArr) + +[] +let ``AsyncSeq.splitAt on empty sequence returns empty first and empty rest`` () = + let first, rest = AsyncSeq.splitAt 3 AsyncSeq.empty |> Async.RunSynchronously + let restArr = rest |> AsyncSeq.toArrayAsync |> Async.RunSynchronously + Assert.AreEqual([||], first) + Assert.AreEqual([||], restArr) + +[] +let ``AsyncSeq.splitAt with count equal to length returns all in first and empty rest`` () = + let source = asyncSeq { yield 7; yield 8; yield 9 } + let first, rest = AsyncSeq.splitAt 3 source |> Async.RunSynchronously + let restArr = rest |> AsyncSeq.toArrayAsync |> Async.RunSynchronously + Assert.AreEqual([| 7; 8; 9 |], first) + Assert.AreEqual([||], restArr) + +[] +let ``AsyncSeq.splitAt with negative count throws ArgumentException`` () = + Assert.Throws(fun () -> + AsyncSeq.splitAt -1 AsyncSeq.empty |> Async.RunSynchronously |> ignore) |> ignore + // ===== removeAt ===== []