Skip to content

Commit 0e09ab3

Browse files
authored
Merge pull request #284 from fsprojects/repo-assist/improve-traverse-funcs-2026-03-17-64ee7add01bad603
[Repo Assist] Fix traverseOptionAsync/traverseChoiceAsync — avoid wasteful MoveNext on failure
2 parents cc1a858 + dbd030d commit 0e09ab3

File tree

2 files changed

+82
-31
lines changed

2 files changed

+82
-31
lines changed

src/FSharp.Control.AsyncSeq/AsyncSeq.fs

Lines changed: 35 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -2449,43 +2449,47 @@ module AsyncSeq =
24492449

24502450
let traverseOptionAsync (f:'T -> Async<'U option>) (source:AsyncSeq<'T>) : Async<AsyncSeq<'U> option> = async {
24512451
use ie = source.GetEnumerator()
2452-
let! move = ie.MoveNext()
2453-
let b = ref move
2452+
let! first = ie.MoveNext()
2453+
let mutable current = first
2454+
let mutable failed = false
24542455
let buffer = ResizeArray<_>()
2455-
let fail = ref false
2456-
while b.Value.IsSome && not fail.Value do
2457-
let! vOpt = f b.Value.Value
2458-
match vOpt with
2459-
| Some v -> buffer.Add v
2460-
| None -> b := None; fail := true
2461-
let! moven = ie.MoveNext()
2462-
b := moven
2463-
if fail.Value then
2464-
return None
2456+
while current.IsSome && not failed do
2457+
let! vOpt = f current.Value
2458+
match vOpt with
2459+
| Some v ->
2460+
buffer.Add v
2461+
let! next = ie.MoveNext()
2462+
current <- next
2463+
| None ->
2464+
failed <- true
2465+
if failed then
2466+
return None
24652467
else
2466-
let res = buffer.ToArray()
2467-
return Some (asyncSeq { for v in res do yield v })
2468-
}
2468+
let res = buffer.ToArray()
2469+
return Some (asyncSeq { for v in res do yield v })
2470+
}
24692471

24702472
let traverseChoiceAsync (f:'T -> Async<Choice<'U, 'e>>) (source:AsyncSeq<'T>) : Async<Choice<AsyncSeq<'U>, 'e>> = async {
24712473
use ie = source.GetEnumerator()
2472-
let! move = ie.MoveNext()
2473-
let b = ref move
2474+
let! first = ie.MoveNext()
2475+
let mutable current = first
2476+
let mutable failWith = ValueNone
24742477
let buffer = ResizeArray<_>()
2475-
let fail = ref None
2476-
while b.Value.IsSome && fail.Value.IsNone do
2477-
let! vOpt = f b.Value.Value
2478-
match vOpt with
2479-
| Choice1Of2 v -> buffer.Add v
2480-
| Choice2Of2 err -> b := None; fail := Some err
2481-
let! moven = ie.MoveNext()
2482-
b := moven
2483-
match fail.Value with
2484-
| Some err -> return Choice2Of2 err
2485-
| None ->
2486-
let res = buffer.ToArray()
2487-
return Choice1Of2 (asyncSeq { for v in res do yield v })
2488-
}
2478+
while current.IsSome && failWith.IsNone do
2479+
let! vOpt = f current.Value
2480+
match vOpt with
2481+
| Choice1Of2 v ->
2482+
buffer.Add v
2483+
let! next = ie.MoveNext()
2484+
current <- next
2485+
| Choice2Of2 err ->
2486+
failWith <- ValueSome err
2487+
match failWith with
2488+
| ValueSome err -> return Choice2Of2 err
2489+
| ValueNone ->
2490+
let res = buffer.ToArray()
2491+
return Choice1Of2 (asyncSeq { for v in res do yield v })
2492+
}
24892493

24902494
#if (NETSTANDARD || NET)
24912495
#if !FABLE_COMPILER

tests/FSharp.Control.AsyncSeq.Tests/AsyncSeqTests.fs

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1211,6 +1211,53 @@ let ``AsyncSeq.traverseChoiceAsync``() =
12111211
Assert.AreEqual("oh no", e)
12121212
Assert.True(([1;2] = (seen |> List.ofSeq)))
12131213

1214+
[<Test>]
1215+
let ``AsyncSeq.traverseOptionAsync returns Some sequence when all elements succeed``() =
1216+
let s = [1;2;3] |> AsyncSeq.ofSeq
1217+
let f i = Some (i * 10) |> async.Return
1218+
let r = AsyncSeq.traverseOptionAsync f s |> Async.RunSynchronously
1219+
match r with
1220+
| None -> Assert.Fail("Expected Some")
1221+
| Some result ->
1222+
let values = result |> AsyncSeq.toListAsync |> Async.RunSynchronously
1223+
Assert.AreEqual([10;20;30], values)
1224+
1225+
[<Test>]
1226+
let ``AsyncSeq.traverseOptionAsync does not read past failing element``() =
1227+
let readCount = ref 0
1228+
let s = asyncSeq {
1229+
for i in 1..10 do
1230+
incr readCount
1231+
yield i
1232+
}
1233+
let f i = (if i <= 3 then Some i else None) |> async.Return
1234+
let _r = AsyncSeq.traverseOptionAsync f s |> Async.RunSynchronously
1235+
// f returns None on element 4; only elements 1..4 should be read from source
1236+
Assert.AreEqual(4, readCount.Value)
1237+
1238+
[<Test>]
1239+
let ``AsyncSeq.traverseChoiceAsync returns Choice1Of2 sequence when all elements succeed``() =
1240+
let s = [1;2;3] |> AsyncSeq.ofSeq
1241+
let f i = Choice1Of2 (i * 10) |> async.Return
1242+
let r = AsyncSeq.traverseChoiceAsync f s |> Async.RunSynchronously
1243+
match r with
1244+
| Choice2Of2 _ -> Assert.Fail("Expected Choice1Of2")
1245+
| Choice1Of2 result ->
1246+
let values = result |> AsyncSeq.toListAsync |> Async.RunSynchronously
1247+
Assert.AreEqual([10;20;30], values)
1248+
1249+
[<Test>]
1250+
let ``AsyncSeq.traverseChoiceAsync does not read past failing element``() =
1251+
let readCount = ref 0
1252+
let s = asyncSeq {
1253+
for i in 1..10 do
1254+
incr readCount
1255+
yield i
1256+
}
1257+
let f i = (if i <= 3 then Choice1Of2 i else Choice2Of2 "stop") |> async.Return
1258+
let _r = AsyncSeq.traverseChoiceAsync f s |> Async.RunSynchronously
1259+
// f returns Choice2Of2 on element 4; only elements 1..4 should be read from source
1260+
Assert.AreEqual(4, readCount.Value)
12141261

12151262
[<Test>]
12161263
let ``AsyncSeq.toBlockingSeq does not hung forever and rethrows exception``() =

0 commit comments

Comments
 (0)