Lib.hs

module Lib (module Lib) where

import Control.Monad (ap, (>=>))

----------------------------------------------------------------------------------------------------

data Coroutine i o a
  = Await (i -> Coroutine i o a)
  | Yield o (Coroutine i o a)
  | Done a

await :: Coroutine i o i
await = Await Done

yield :: o -> Coroutine i o ()
yield o = Yield o (Done ())

instance Functor (Coroutine i o) where
  fmap f = \case
    Await k -> Await \i -> fmap f (k i)
    Yield o k -> Yield o (fmap f k)
    Done x -> Done (f x)

instance Applicative (Coroutine i o) where
  pure = Done
  (<*>) = ap

instance Monad (Coroutine i o) where
  Await next >>= k = Await \i -> next i >>= k
  Yield o next >>= k = Yield o (next >>= k)
  Done a >>= k = k a

interp :: [i] -> Coroutine i o a -> ([o], a)
interp inputs c = go inputs [] c
 where
  go (i : is) os (Await next) = go is os (next i)
  go [] _ (Await _) = error "interp: no input for Await"
  go is os (Yield o next) = go is (o : os) next
  go _ os (Done a) = (os, a)

example :: Coroutine Int String Char
example = do
  x <- await
  yield . show $ x + 3
  yield . show $ x + 6
  pure . head . show $ x + 10

test :: ([String], Char)
test = interp [1 ..] example

----------------------------------------------------------------------------------------------------

newtype CoroutineCPS i o a = CoroutineCPS
  { runCoroutineCPS
      :: forall b
       . ((i -> CoroutineCPS i o a) -> b)
      -> (o -> CoroutineCPS i o a -> b)
      -> (a -> b)
      -> b
  }

awaitCPS :: CoroutineCPS i o i
awaitCPS = CoroutineCPS \kAwait _ _ -> kAwait \i -> CoroutineCPS \_ _ kDone -> kDone i

yieldCPS :: o -> CoroutineCPS i o ()
yieldCPS o = CoroutineCPS \_ kYield _ -> kYield o (CoroutineCPS \_ _ kDone -> kDone ())

instance Functor (CoroutineCPS i o) where
  fmap f k = CoroutineCPS \kAwait kYield kDone ->
    (runCoroutineCPS k)
      (\next -> kAwait \i -> fmap f (next i))
      (\o next -> kYield o (fmap f next))
      (\a -> kDone (f a))

instance Applicative (CoroutineCPS i o) where
  pure a = CoroutineCPS \_ _ kDone -> kDone a
  (<*>) = ap

instance Monad (CoroutineCPS i o) where
  k >>= next = CoroutineCPS \kAwait kYield kDone ->
    (runCoroutineCPS k)
      (\next' -> kAwait \i -> next' i >>= next)
      (\o next' -> kYield o (next' >>= next))
      (\a -> runCoroutineCPS (next a) kAwait kYield kDone)

interpCPS :: [i] -> CoroutineCPS i o a -> ([o], a)
interpCPS input c = go input [] c
 where
  go is os k =
    (runCoroutineCPS k)
      (\next -> case is of [] -> error "interpCPS: no input for Await"; (i : is') -> go is' os (next i))
      (\o next -> go is (o : os) next)
      (\a -> (os, a))

exampleCPS :: CoroutineCPS Int String Char
exampleCPS = do
  x <- awaitCPS
  yieldCPS . show $ x + 3
  yieldCPS . show $ x + 6
  pure . head $ show $ x + 10

testCPS :: ([String], Char)
testCPS = interpCPS [1 ..] exampleCPS

----------------------------------------------------------------------------------------------------

newtype Codensity m a = Codensity {runCodensity :: forall b. (a -> m b) -> m b}

lowerCodensity :: (Monad m) => Codensity m a -> m a
lowerCodensity ma = (runCodensity ma) return

instance Functor (Codensity m) where
  fmap f fa = Codensity \k -> (runCodensity fa) \a -> k (f a)

instance Applicative (Codensity m) where
  pure a = Codensity \k -> k a
  mf <*> ma = Codensity \k -> (runCodensity mf) \f -> (runCodensity ma) \a -> k (f a)

instance Monad (Codensity m) where
  ma >>= f = Codensity \k -> (runCodensity ma) \a -> runCodensity (f a) k

newtype CoroutineCodensity i o a = CoroutineCodensity {runCoroutineCodensity :: Codensity (Coroutine i o) a}
  deriving newtype (Functor, Applicative, Monad)

toCoroutineCodensity :: Coroutine i o a -> CoroutineCodensity i o a
toCoroutineCodensity c = CoroutineCodensity case c of
  Await next -> Codensity \k -> Await \i -> next i >>= k
  Yield o next -> Codensity \k -> Yield o (next >>= k)
  Done a -> Codensity \k -> k a

fromCoroutineCodensity :: CoroutineCodensity i o a -> Coroutine i o a
fromCoroutineCodensity c = runCodensity (runCoroutineCodensity c) pure

awaitCodensity :: CoroutineCodensity i o i
awaitCodensity = CoroutineCodensity (Codensity Await)

yieldCodensity :: o -> CoroutineCodensity i o ()
yieldCodensity o = CoroutineCodensity (Codensity \k -> Yield o (k ()))

interpCodensity :: [i] -> CoroutineCodensity i o a -> ([o], a)
interpCodensity input c = interp input (fromCoroutineCodensity c)

exampleCodensity :: CoroutineCodensity Int String Char
exampleCodensity = do
  x <- awaitCodensity
  yieldCodensity . show $ x + 3
  yieldCodensity . show $ x + 6
  pure . head . show $ x + 10

testCodensity :: ([String], Char)
testCodensity = interpCodensity [1 ..] exampleCodensity

----------------------------------------------------------------------------------------------------

sumInput :: Int -> Coroutine Int o Int
sumInput n = Await \i -> foldl (>=>) pure (replicate (n - 1) (\x -> await >>= pure . (+ x))) i

sumInputCPS :: Int -> CoroutineCPS Int o Int
sumInputCPS n = CoroutineCPS \kAwait _ _ -> kAwait \i ->
  foldl (>=>) pure (replicate (n - 1) (\x -> awaitCPS >>= pure . (+ x))) i

sumInputCodensity :: Int -> CoroutineCodensity Int o Int
sumInputCodensity n = CoroutineCodensity $ Codensity \k -> Await \i ->
  let c = foldl (>=>) pure (replicate (n - 1) (\x -> awaitCodensity >>= pure . (+ x))) i
   in runCodensity (runCoroutineCodensity c) k