CPS like Async Monad

cpslike.lean

@[inline]
def emptyTask : Task (Except IO.Error Unit) :=
  Task.pure (Except.pure ())

@[inline]
def block (io : Task α) : IO α :=
  pure io.get

@[inline]
def nextTask {α} (t : Task α) (f : α → BaseIO Unit) : BaseIO Unit :=
  discard <| BaseIO.bindTask t (f · *> pure emptyTask)

-- Async CPS like

namespace CPS

structure Async (α : Type) where
  run : (AsyncTask α → BaseIO Unit) → BaseIO Unit

def Async.toTask {α : Type} [Nonempty α] (x : Async α) : IO (Task (Except IO.Error α)) := do
  let promise ← IO.Promise.new
  x.run (nextTask · promise.resolve)
  pure promise.result!

def Async.exec {α : Type} [Nonempty α] (x : Async α) : IO α := do
  match (← block =<< x.toTask) with
  | .ok res => pure res
  | .error err => throw err

def await {α : Type} (task : IO (AsyncTask α)) : Async α :=
  ⟨fun k => do
    match ← task.toBaseIO with
    | .ok result => k result
    | .error err => k (Task.pure (.error err))⟩

def parallel (task : Async α) : Async Unit :=
  ⟨fun k => do
    discard <| IO.asTask (task.run (fun _ => pure ()))
    k emptyTask⟩

instance : Functor Async where
  map f x := ⟨fun k => x.run (fun t => k (t.map f))⟩

instance : Monad Async where
  pure a := ⟨fun k => k (AsyncTask.pure a)⟩
  bind x f := ⟨fun k => x.run fun t =>
    nextTask t (fun
      | .ok res => (f res).run k
      | .error err => k (Task.pure (.error err)))⟩

instance : MonadLift IO Async where
  monadLift io := ⟨fun k => do
    match ← io.toBaseIO with
    | .ok res => k (AsyncTask.pure res)
    | .error err => k (Task.pure (.error err))⟩

instance : MonadExcept IO.Error Async where
  throw x := ⟨fun t => t (Task.pure (.error x))⟩
  tryCatch x f := ⟨fun k => x.run fun t =>
    nextTask t fun
      | .ok res => k (Task.pure (.ok res))
      | .error err => (f err).run k⟩

end CPS