Productive coprogramming in the Par monad

Productive.hs

{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE InstanceSigs #-}
module Productive where

-- | Productive coprogramming in the Par monad.  Based off of
-- "Productive Coprogramming with Guarded Recursion" by Robert
-- Atkey and Conor McBride.

import Prelude hiding (fix, take)
import Unsafe.Coerce

-- Your favorite implementation of monad-par
import MiniPar (Par, IVar, fork, get, put, new, newFull)

-- | Clock-variable @k@ annotated guardedness type constructor;
-- this says that a value of type @a@ was constructed on time stream
-- @k@.
data Clock
newtype D (k :: Clock) a = D { unD :: IVar a }

-- D is a functor, but not in Hask: it is a functor in the
-- Kleisli category for some monad.  (BTW, FunctorM used to exist in
-- base, but it was replaced by Foldable and Traversable.  But those
-- classes are not right for this case.)

-- | Functor in the Kleisli category @m@.
class Monad m => FunctorM m f where
    fmapM :: (a -> m b) -> f a -> m (f b)

-- | Applicative functor in the Kleisli category @m@.
class Monad m => ApplicativeM m f where
    pureM :: a -> m (f a)
    apM   :: f (a -> m b) -> f a -> m (f b)

instance FunctorM Par (D k) where
    fmapM :: (a -> Par b) -> D k a -> Par (D k b)
    fmapM f (D va) = do
        vb <- new
        fork $ do
            a <- get va
            put vb =<< f a
        return (D vb)

instance ApplicativeM Par (D k) where
    pureM :: a -> Par (D k a)
    pureM = fmap D . newFull

    apM :: D k (a -> Par b) -> D k a -> Par (D k b)
    apM (D vf) (D va) = do
        vb <- new
        fork $ do
            f <- get vf
            a <- get va
            put vb =<< f a
        return (D vb)

-- | Helper function. Technically you can build this with pureM and apM
-- but it is a pain.
apM2 :: D k (a -> b -> Par c) -> D k a -> D k b -> Par (D k c)
apM2 (D vf) (D va) (D vb) = do
    vc <- new
    fork $ do
        f <- get vf
        a <- get va
        b <- get vb
        put vc =<< f a b
    return (D vc)

-- | Fixpoint operator lets you refer to the result of your computation,
-- but the knot tying variable may only be used in a guarded way.
fix :: (D k a -> Par a) -> Par a
fix f = do
    v <- new
    a <- f (D v)
    put v a
    return a

-- | A guarded value can be observed today, if it was constructed from
-- a universally quantified time stream (allowing us to make any finite
-- observation.)
force0 :: (forall (k :: Clock). D k a) -> Par a
force0 (D v) = get v

-- Some helpers to deal with Haskell's lack of type-level lambda.

newtype ForAll  f     = ForAll  { unForAll  :: forall (k :: Clock). f k     }
newtype ForAll1 f a   = ForAll1 { unForAll1 :: forall (k :: Clock). f k a   }
newtype ForAll2 f a b = ForAll2 { unForAll2 :: forall (k :: Clock). f k a b }

-- When the clock variable can occur in the type (i.e., f is
-- parametrized by k), we must ensure that once we force it, the
-- universal quantifier is *inside* the monad; otherwise, when we
-- bind the value, the quantifier will get skolemized.  I don't
-- know if there is a way to implement these functions without
-- unsafeCoerce.  Since Clock is an uninhabited kind, it's easy
-- to see that these types are representationally equal.

force  :: (forall (k :: Clock). D k (f k)    ) -> Par (ForAll  f)
force1 :: (forall (k :: Clock). D k (f k a)  ) -> Par (ForAll1 f a)
force2 :: (forall (k :: Clock). D k (f k a b)) -> Par (ForAll2 f a b)

force    d = get (unsafeCoerce d)
force1   d = get (unsafeCoerce d)
force2   d = get (unsafeCoerce d)

-- Similarly, if we have a universally quantified monadic computation
-- of some type, we also have a monadic computation of a universally
-- quantified type.

commute  :: (forall (k :: Clock). Par (f k))     -> Par (ForAll  f)
commute1 :: (forall (k :: Clock). Par (f k a))   -> Par (ForAll1 f a)
commute2 :: (forall (k :: Clock). Par (f k a b)) -> Par (ForAll2 f a b)

commute  m = unsafeCoerce m
commute1 m = unsafeCoerce m
commute2 m = unsafeCoerce m

-- Here are stream examples from the paper.

data Stream k a = Cons { scar :: a, scdr :: D k (Stream k a) }

ones :: Par (Stream k Int)
ones = fix $ \d -> return (Cons 1 d)

mergef :: (a -> a -> D k (Stream k a) -> Par (Stream k a))
       -> Stream k a -> Stream k a -> Par (Stream k a)
mergef f xs ys = do
    go <- fix $ \g -> return $ \(Cons x xs) (Cons y ys) ->
                    f x y =<< apM2 g xs ys
    go xs ys

deStreamCons :: forall a. (forall (k :: Clock). Stream k a) -> Par (a, ForAll1 Stream a)
deStreamCons s = do
    ds <- force1 (scdr s)
    return (scar s, ds)

take :: Int -> (forall (k :: Clock). Stream k a) -> Par [a]
take 0 _ = return []
take n s = do
    (x, s') <- deStreamCons s
    xs      <- take (n - 1) (unForAll1 s')
    return (x : xs)