Cooperative multitasking in haskell

Scheduling.hs

{-# LANGUAGE RankNTypes #-}

module Scheduling where

newtype Fiber f a = Fiber (f (Maybe (Maybe a, Fiber f a)))

noop :: Applicative f => Fiber f a
noop = Fiber $ pure Nothing

singleton :: Applicative f => f (Maybe a) -> Fiber f a
singleton task = Fiber $ fmap (\step -> Just (step, noop)) task

stitch :: Monad f => Fiber f a -> Fiber f a -> Fiber f a
stitch (Fiber f1) (Fiber f2) = Fiber $ f1 >>= stitch'
                             where stitch' Nothing             = f2
                                   stitch' (Just (next, tail)) = pure $ Just (next, stitch tail (Fiber f2))

cons :: Monad f => f (Maybe a) -> Fiber f a -> Fiber f a
cons task fiber = stitch (singleton task) fiber

snoc :: Monad f => f (Maybe a) -> Fiber f a -> Fiber f a
snoc task fiber = stitch fiber (singleton task)

consAction :: Monad f => f () -> Fiber f a -> Fiber f a
consAction action fiber = cons (fmap (const Nothing) action) fiber

consEmission :: Monad f => f a -> Fiber f a -> Fiber f a
consEmission task fiber = cons (fmap Just task) fiber

snocEmission :: Monad f => f a -> Fiber f a -> Fiber f a
snocEmission task fiber = snoc (fmap Just task) fiber

snocAction :: Monad f => f () -> Fiber f a -> Fiber f a
snocAction action fiber = snoc (fmap (const Nothing) action) fiber

unfoldFiber :: Monad f => (b -> f (Maybe (Maybe a, b))) -> f b -> Fiber f a
unfoldFiber f = Fiber . (fmap $ fmap $ fmap $ unfoldFiber f . pure) . (=<<) f

forever :: Monad f => f a -> Fiber f a
forever task = unfoldFiber (\_ -> fmap (\a -> Just (Just a, ())) task)  (pure ())

times :: Monad f => Int -> f a -> Fiber f a
times n task = unfoldFiber (\x -> if x > 0 then fmap (\a -> Just (Just a, x - 1)) task else pure Nothing) $ pure n

type Scheduler f = forall a. [Fiber f a] -> f ()

roundRobin :: Monad f => Scheduler f
roundRobin [] = pure ()
roundRobin ((Fiber f):fs) = f >>= reschedule
                    where reschedule Nothing             = roundRobin fs
                          reschedule (Just (next, rest)) = roundRobin $ fs ++ [rest]

main :: IO ()
main = roundRobin [ snocAction (putStrLn "done pinging") $ times 3 $ putStrLn "ping"
                  , snocAction (putStrLn "done ponging") $ times 2 $ putStrLn "pong"
                  ]