Add pick related functions to SeqT

fsprojects · Oct 6, 2024 · 465832c · 465832c
1 parent 43050e6
commit 465832c
Showing 1 changed file with 38 additions and 0 deletions.
diff --git a/src/FSharpPlus/Data/Seq.fs b/src/FSharpPlus/Data/Seq.fs
@@ -797,6 +797,44 @@ 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 tryPickM<'T, 'U, .. > (f: 'T -> '``Monad<'U option>``) (source: SeqT<'``Monad<bool>``, 'T>) : '``Monad<'U option>`` = 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 res }
+
+    let inline pickM<'T, 'U, .. > (f: 'T -> '``Monad<'U option>``) (source: SeqT<'``Monad<bool>``, 'T>) : '``Monad<'U>`` =
+        source |> tryPickM<_, 'U, '``Monad<'U option>``, _, '``Monad<unit>``> f |> Map.Invoke (function Some v -> (v: 'U) | _ -> raise (KeyNotFoundException ()))
+
+    let inline tryPick<'T, 'U, .. > (f: 'T -> 'U option) (source: SeqT<'``Monad<bool>``, 'T>) : '``Monad<'U option>`` =
+        tryPickM<_, _, _ , _, '``Monad<unit>``> (f >> Return.Invoke) source
+
+    let inline pick (f: 'T -> 'U option) (source: SeqT<'``Monad<bool>``, 'T>) : '``Monad<'U>`` =
+        (source |> tryPick<_, _, _, '``Monad<unit>``, _> f : '``Monad<'U option>``) |> Map.Invoke (function Some v -> (v: 'U) | _ -> raise (KeyNotFoundException ()))
+
+    let inline contains value (source: SeqT<'``Monad<bool>``, 'T>) : '``Monad<bool>`` =
+        source |> tryPick<_, _, _, '``Monad<unit>``, '``Monad<unit option>``> (fun v -> if v = value then Some () else None) |> Map.Invoke Option.isSome
+
+    let inline tryFind<'T, .. > f (source: SeqT<'``Monad<bool>``, 'T>) : '``Monad<'T option>`` =
+        source |> tryPick<_, _, _, '``Monad<unit>``, _> (fun v -> if f v then Some v else None)
+
+    let inline find f (source: SeqT<'``Monad<bool>``, 'T>) : '``Monad<'T>`` =
+        source |> tryFind<_, _, '``Monad<'T option>``, '``Monad<unit>``> f |> Map.Invoke (function Some v -> (v: 'T) | _ -> raise (KeyNotFoundException ()))
+
+    let inline exists f (source: SeqT<'``Monad<bool>``, 'T>) : '``Monad<bool>`` =
+        source |> tryFind<_, _, '``Monad<'T option>``, '``Monad<unit>``> f |> Map.Invoke Option.isSome
+
+    let inline forall f (source: SeqT<'``Monad<bool>``, 'T>) : '``Monad<bool>`` =
+        source |> tryFind<_, _, '``Monad<'T option>``, '``Monad<unit>``> (f >> not) |> Map.Invoke Option.isNone
+
     [<RequireQualifiedAccess; EditorBrowsable(EditorBrowsableState.Never)>]
     type TryWithState<'``Monad<bool>``, 'T> =
        | NotStarted of SeqT<'``Monad<bool>``, 'T>