Eio Task abstraction

Task.ml

(*https://package.elm-lang.org/packages/elm/core/latest/Task*)
type never = |

module type IO = sig
  type 'a t

  val return : 'a -> 'a t
  val bind : 'a t -> ('a -> 'b t) -> 'b t
  val of_thunk : (unit -> 'a) -> 'a t
end

module type S = sig
  type 'a io
  type ('x, 'a) task

  val perform : ('a -> 'msg) -> (never, 'a) task -> 'msg io
  val attempt : (('a, 'x) Result.t -> 'msg) -> ('x, 'a) task -> 'msg io
  val from_io : 'a io -> ('x, 'a) task
  val run : (unit -> 'a) -> ('x, 'a) task

  val bind : ('a -> ('x, 'b) task) -> ('x, 'a) task -> ('x, 'b) task
  val succeed : 'a -> ('x, 'a) task
  val fail : 'x -> ('x, 'a) task

  val map : ('a -> 'b) -> ('x, 'a) task -> ('x, 'b) task

  val on_error : ('x -> ('y, 'a) task) -> ('x, 'a) task -> ('y, 'a) task
  val map_error : ('x -> 'y) -> ('x, 'a) task -> ('y, 'a) task

  val ( let* ) : ('x, 'a) task -> ('a -> ('x, 'b) task) -> ('x, 'b) task
  val ( let+ ) : ('x, 'a) task -> ('a -> 'b) -> ('x, 'b) task
  val ( and+ ) : ('x, 'a) task -> ('x, 'b) task -> ('x, 'a * 'b) task
  val ( and* ) : ('x, 'a) task -> ('x, 'b) task -> ('x, 'a * 'b) task
end

module Make (Io : IO) : S with type 'a io = 'a Io.t = struct
  type 'a io = 'a Io.t

  type ('x, 'a) task = {task: 'r. ('a -> 'r Io.t) -> ('x -> 'r Io.t) -> 'r Io.t}

  let succeed : 'a -> ('x, 'a) task = fun a -> {task = (fun v _ -> v a)}
  let fail : 'x -> ('x, 'a) task = fun x -> {task = (fun _ e -> e x)}

  let map : ('a -> 'b) -> ('x, 'a) task -> ('x, 'b) task =
   fun f {task} -> {task = (fun ok err -> task (fun a -> ok (f a)) err)}

  let map_error : ('x -> 'y) -> ('x, 'a) task -> ('y, 'a) task =
   fun f {task} -> {task = (fun ok err -> task ok (fun x -> err (f x)))}

  let bind : ('a -> ('x, 'b) task) -> ('x, 'a) task -> ('x, 'b) task =
   fun f {task} ->
    {task = (fun ok err -> task (fun a -> (f a).task ok err) err)}

  let ( let* ) task f = bind f task
  let ( let+ ) task f = map f task

  let ( and+ ) : ('x, 'a) task -> ('x, 'b) task -> ('x, 'a * 'b) task =
   fun {task = ta} {task = tb} ->
    {task = (fun ok err -> ta (fun a -> tb (fun b -> ok (a, b)) err) err)}

  let ( and* ) = ( and+ )

  let on_error : ('x -> ('y, 'a) task) -> ('x, 'a) task -> ('y, 'a) task =
   fun f {task} -> {task = (fun ok err -> task ok (fun x -> (f x).task ok err))}

  let perform : ('a -> 'msg) -> (never, 'a) task -> 'msg Io.t =
   fun f {task} ->
    task (fun a -> Io.return (f a)) (fun (n : never) -> match n with _ -> .)

  let attempt : (('a, 'x) Result.t -> 'msg) -> ('x, 'a) task -> 'msg Io.t =
   fun f {task} ->
    task (fun a -> Io.return (f (Ok a))) (fun x -> Io.return (f (Error x)))

  let from_io : 'a Io.t -> ('x, 'a) task =
   fun io -> {task = (fun ok _ -> Io.bind io ok)}

  let run : (unit -> 'a) -> ('x, 'a) task =
   fun thunk -> from_io (Io.of_thunk thunk)
end

module Eio_io : IO with type 'a t = 'a Eio.Promise.t = struct
  type 'a t = 'a Eio.Promise.t

  let return = Eio.Promise.create_resolved
  let bind t f = f (Eio.Promise.await t)
  let of_thunk f = Eio.Promise.create_resolved (f ())
end

module Eio_task = Make (Eio_io)

module Make_pool_io (Cfg : sig
  val sw : Eio.Switch.t
  val pool : Eio.Executor_pool.t
end) : IO with type 'a t = 'a Eio.Promise.or_exn = struct
  type 'a t = 'a Eio.Promise.or_exn

  let return a = Eio.Promise.create_resolved (Ok a)

  let bind t f =
    match Eio.Promise.await t with
    | Ok a -> f a
    | Error exn -> Eio.Promise.create_resolved (Error exn)

  let of_thunk f =
    Eio.Executor_pool.submit_fork ~sw:Cfg.sw Cfg.pool ~weight:1.0 f
end

module Make_pool_task (Cfg : sig
  val sw : Eio.Switch.t
  val pool : Eio.Executor_pool.t
end) =
  Make (Make_pool_io (Cfg))