Add pick related functions to SeqT (#608)

fsprojects · Oct 6, 2024 · f17a734 · f17a734
1 parent 43050e6
commit f17a734
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 0 deletions.
diff --git a/src/FSharpPlus/Data/Seq.fs b/src/FSharpPlus/Data/Seq.fs
@@ -797,6 +797,61 @@ module SeqT_V2 =
     let inline iteri<'T, .. > (f: int -> 'T -> unit)      (source: SeqT<'``Monad<bool>``, 'T>) : '``Monad<unit>`` = iteriM (fun i x -> result (f i x)) source
     let inline iter<'T, .. > f (source: SeqT<'``Monad<bool>``, 'T>) : '``Monad<unit>`` = iterM (f >> result) source
 
+    let inline internal tryPickMAndMap<'T, 'U, .. > (f: 'T -> '``Monad<'U option>``) (source: SeqT<'``Monad<bool>``, 'T>) (postMap: 'U option -> 'V) : '``Monad<'V>`` = innerMonad2<_, '``Monad<unit>``> () {
+        use ie = (source :> IEnumerableM<'``Monad<bool>``, 'T>).GetEnumerator ()
+        let! (move: bool) = ie.MoveNext ()
+        let mutable b = move
+        let mutable res = None
+        while b && res.IsNone do
+            let! (fv: 'U option) = f ie.Current
+            match fv with
+            | Some _ as r -> res <- r
+            | None ->
+                let! moven = ie.MoveNext ()
+                b <- moven
+        return postMap res }
+
+    let inline internal tryPickAndMap<'T, 'U, .. > (f: 'T -> 'U option) (source: SeqT<'``Monad<bool>``, 'T>) (postMap: 'U option -> 'V) : '``Monad<'V>`` = innerMonad2<_, '``Monad<unit>``> () {
+        use ie = (source :> IEnumerableM<'``Monad<bool>``, 'T>).GetEnumerator ()
+        let! (move: bool) = ie.MoveNext ()
+        let mutable b = move
+        let mutable res = None
+        while b && res.IsNone do
+            let (fv: 'U option) = f ie.Current
+            match fv with
+            | Some _ as r -> res <- r
+            | None ->
+                let! moven = ie.MoveNext ()
+                b <- moven
+        return postMap res }
+
+    let inline tryPickM<'T, 'U, .. > (f: 'T -> '``Monad<'U option>``) (source: SeqT<'``Monad<bool>``, 'T>) : '``Monad<'U option>`` =
+        tryPickMAndMap<_, 'U, _, '``Monad<unit>``, _, _, _> f source id
+
+    let inline pickM<'T, 'U, .. > (f: 'T -> '``Monad<'U option>``) (source: SeqT<'``Monad<bool>``, 'T>) : '``Monad<'U>`` =
+        tryPickMAndMap<_, 'U, '``Monad<'U option>``, '``Monad<unit>``, _, _, _> f source (function Some v -> (v: 'U) | _ -> raise (KeyNotFoundException ()))
+
+    let inline tryPick<'T, 'U, .. > (f: 'T -> 'U option) (source: SeqT<'``Monad<bool>``, 'T>) : '``Monad<'U option>`` =
+        tryPickAndMap<_, _, _, _, '``Monad<unit>``, _> f source id
+
+    let inline pick (f: 'T -> 'U option) (source: SeqT<'``Monad<bool>``, 'T>) : '``Monad<'U>`` =
+        tryPickAndMap<_, _, _, _, '``Monad<unit>``, _> f source (function Some v -> (v: 'U) | _ -> raise (KeyNotFoundException ()))
+
+    let inline contains value (source: SeqT<'``Monad<bool>``, 'T>) : '``Monad<bool>`` =
+        tryPickAndMap<_, _, _, _, '``Monad<unit>``, _> (fun v -> if v = value then Some () else None) source Option.isSome
+
+    let inline tryFind<'T, .. > f (source: SeqT<'``Monad<bool>``, 'T>) : '``Monad<'T option>`` =
+        tryPickAndMap<_, _, _, _, '``Monad<unit>``, _> (fun v -> if f v then Some v else None) source id
+
+    let inline find f (source: SeqT<'``Monad<bool>``, 'T>) : '``Monad<'T>`` =
+        tryPickAndMap<_, _, _, _, '``Monad<unit>``, _> (fun v -> if f v then Some v else None) source (function Some v -> (v: 'T) | _ -> raise (KeyNotFoundException ()))
+
+    let inline exists f (source: SeqT<'``Monad<bool>``, 'T>) : '``Monad<bool>`` =
+        tryPickAndMap<_, _, _, _, '``Monad<unit>``, _> (fun v -> if f v then Some v else None) source Option.isSome
+
+    let inline forall f (source: SeqT<'``Monad<bool>``, 'T>) : '``Monad<bool>`` =
+        tryPickAndMap<_, _, _, _, '``Monad<unit>``, _> (fun v -> if f v then None else Some v) source Option.isNone
+
     [<RequireQualifiedAccess; EditorBrowsable(EditorBrowsableState.Never)>]
     type TryWithState<'``Monad<bool>``, 'T> =
        | NotStarted of SeqT<'``Monad<bool>``, 'T>

diff --git a/tests/FSharpPlus.Tests/SeqT.fs b/tests/FSharpPlus.Tests/SeqT.fs
@@ -64,6 +64,13 @@ module BasicTests =
         let y3 = SeqT.run x3 |> extract |> toList
         CollectionAssert.AreEqual (y3, [("0", 0, 0); ("1", 10, 1); ("2", 20, 2)])
 
+    [<Test>]
+    let picks () =
+        let infinite: SeqT<Async<_>, _> = SeqT.unfold (fun x -> Some (x, x + 1)) 0
+        let five = SeqT.find ((<) 4) infinite |> Async.RunSynchronously
+        Assert.AreEqual (5, five)
+
+
     // Compile tests
     let binds () =
         let res1 = SeqT [|seq [1..4] |] >>= fun x -> SeqT [|seq [x * 2] |]