coroutines.hs

{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
module Continuations where

import Control.Monad
import Text.Printf
import Debug.Trace

t :: String -> a -> a
t msg = trace (printf "XXX: %s" msg)

data C r a = C { unC :: (a -> r) -> r }

-- | Show a 'C' computation by running it and showing the result.
instance (Show a) => Show (C a a) where
  show m = printf "C %s" (show (unC m id))

instance Functor (C r) where
  fmap f (C m) = C $ \k -> m (k . f)

instance Applicative (C r) where
  pure x = C $ \k -> k x
  (<*>) = ap

instance Monad (C r) where
  (C m) >>= f = C $ \k -> m (\a -> unC (f a) k)

callcc :: forall a r b. ((a -> C r b) -> C r a) -> C r a
callcc f = C $ \(k :: a -> r) ->
  let k' :: a -> C r b
      -- A computation that ignores it's continuation @_@ and uses the
      -- enclosing continuation @k@ instead.
      k' a = C (\_ -> k a)
  in unC (f k') k

-- | Same code as 'callcc', but more polymorphic, so that the cc can
-- be used at any result type.
callcc' :: forall a r. String -> ((forall b. a -> C r b) -> C r a) -> C r a
callcc' msg f = C $ \(k :: a -> r) ->
  let k' :: a -> C r b
      -- A computation that ignores it's continuation @_@ and uses the
      -- enclosing continuation @k@ instead.
      k' a = t msg $ C (\_ -> k a)
  in unC (f k') k

expose :: C r r -> r
expose (C f) = f id

data ContRec r i o = ContRec (forall b. Result r i o -> C r b) i

data Result r i o = Done o
                  | Yield o (forall b. ContRec r i o -> C r b)

type Coroutine r i o =
  (forall b. Result r i o -> C r b) -> (forall b. C r b)

yield :: (forall b. Result r i o -> C r b) -> o -> C r (ContRec r i o)
yield cont out = callcc' "yield" $ \cc -> cont (Yield out cc)

done :: (forall b. Result r i o -> C r b) -> o -> C r a
done cont out = cont (Done out)

yieldBack :: (forall b. ContRec r i o -> C r b) -> i -> C r (Result r i o)
yieldBack cont i = callcc' "yieldBack" $ \cc -> cont (ContRec cc i)

runningSum :: Coroutine r Int Int
runningSum cc0 = do
  ContRec cc1 i1 <- yield cc0 0
  ContRec cc2 i2 <- yield cc1 1
  done cc2 (i1 + i2)

test1 :: C (String, Int) (String, Int)
test1 = do
  x <- callcc' "running" $ \c -> runningSum c
  case x of
    Done i -> pure ("alpha", i)
    Yield i r -> do
      y <- yieldBack r 8
      case y of
        Done i -> pure ("beta", i)
        Yield i r -> do
          z <- yieldBack r 900
          case z of
            Done i -> pure ("gamma", i)
            _ -> error "foo"

-- | The constantly-done coroutine
coconst :: forall r i o. o -> Coroutine r i o
coconst o cc = cc (Done o)

-- | Generalization: The coroutine that yields from a stream
costream :: forall r i o. [o] -> Coroutine r i o
costream (o:os) cont = cont (Yield o (\(ContRec cont' _) -> costream os cont'))

-- | Feed in a constant value to a coroutine
feedConstant :: forall r i o. i -> Coroutine r i o -> C r [o]
feedConstant i co = callcc' "consume" co >>= helper
  where helper :: Result r i o ->  C r [o]
        helper (Done o)     = pure [o]
        helper (Yield o c0) = do
          r <- yieldBack c0 i
          (o:) <$> helper r

-- | Generalization: Feed in a stream of values to a coroutine
feedStream :: forall r i o. [i] -> Coroutine r i o -> C r [o]
feedStream is co = callcc' "consume" co >>= helper is
  where helper :: [i] -> Result r i o -> C r [o]
        helper _        (Done o)     = pure [o]
        helper (i : is) (Yield o c0) = do
          r <- yieldBack c0 i
          (o:) <$> helper is r

schedule :: forall r i o. Coroutine r i o -> Coroutine r o i -> C r (Either i o)
schedule co1 co2 = do
  result1 <- callcc' "schedule" co1
  result2 <- callcc' "schedule" co2
  helper result1 result2
  where helper :: Result r i o -> Result r o i -> C r (Either i o)
        helper (Done o)      _                 = pure $ Right o
        helper _            (Done i)           = pure $ Left i
        helper r1@(Yield o c1) r2@(Yield i c2) = do
          r2 <- yieldBack c2 o
          r1 <- yieldBack c1 i
          helper r1 r2

main :: IO ()
main = do
  -- let x = expose $ stream [4..] runningSum
  print x