Free (CoYoneda f) ~ Operational

CoYoOperational.hs

{-# LANGUAGE GADTs, LambdaCase, MultiParamTypeClasses, FlexibleInstances #-}
-- |
-- Module      : Control.Operational.Monad
-- Copyright   : (c) Austin Seipp 2013
-- License     : BSD3
--
-- Maintainer  : aseipp@pobox.com
-- Stability   : experimental
-- Portability : unportable (GADTs, MPTCs, etc)
--
-- 'CoYoneda' + 'Free' = 'Operational'.
--
module CoYoOperational
       ( -- * Pure API
         Program     -- :: (* -> *) -> * -> *
       , ProgramView -- :: (* -> *) -> * -> *
       , singleton   -- :: f a -> Program f a
       , view        -- :: Functor f => Program f a -> ProgramView f a

         -- * Transformer API
       , ProgramT         -- :: (* -> *) -> (* -> *) -> * -> *
       , ProgramViewT(..) -- :: (* -> *) -> (* -> *) -> * -> *
       , singletonT       -- :: Monad m => f a -> ProgramT f m a
       , viewT            -- :: (Monad m, Functor f) => ProgramT f m a -> ProgramViewT f m a
       ) where
import Control.Applicative
import Data.Functor.Identity

import Control.Monad.Trans.Free          -- from 'free'
import Data.Functor.Yoneda.Contravariant -- from 'kan-extensions'

-- the following are only needed for examples
import Data.Maybe
import System.Exit (exitSuccess)
import Data.Function
import Data.Sequence hiding (singleton)
import qualified Data.Sequence as Seq(singleton)
import Control.Monad
import Control.Monad.Trans

--------------------------------------------------------------------------------
-- Pure API. Merely the transformer version combined with Identity

-- | The type 'Program f a' represents a *program*, which is defined
-- as a sequence of primitive instructions. The underlying
-- instructions should be 'Functor's, as we embed them into a 'Free'
-- 'Monad'.
--
-- This is defined as 'ProgramT' over the 'Identity' 'Monad'.
type Program f a = ProgramT f Identity a

-- | 'ProgramView' represents a means of representing a 'Program' and
-- inspecting it step-by-step.
--
-- This is defined as 'ProgramViewT' over the 'Identity' 'Monad'.
type ProgramView f a = ProgramViewT f Identity a

-- | Lift a single primitive into a program.
singleton :: f a -> Program f a
singleton f = singletonT f
{-# INLINE singleton #-}

-- | Transform a given 'Program' into a 'ProgramViewT' which will
-- allow you to interpret the underlying 'Free' 'Monad' structure.
view :: Program f a -> ProgramView f a
view f = runIdentity (viewT f)
{-# INLINE view #-}

--------------------------------------------------------------------------------
-- Transformer API

-- | This is a version of 'Program' which is abstracted over the base
-- 'Monad'.
newtype ProgramT f m a where
  ProgramT :: { runProgT :: FreeT (Yoneda f) m a } -> ProgramT f m a

instance Monad m => Functor (ProgramT f m) where
  fmap f (ProgramT g) = ProgramT (fmap f g)
  {-# INLINE fmap #-}

instance Monad m => Applicative (ProgramT f m) where
  pure x    = ProgramT (return x)
  {-# INLINE pure #-}

  (<*>) f g = ProgramT (runProgT f <*> runProgT g)
  {-# INLINE (<*>) #-}

instance Monad m => Monad (ProgramT f m) where
  return             = ProgramT . return
  {-# INLINE return #-}

  (ProgramT f) >>= g = ProgramT (f >>= runProgT . g)
  {-# INLINE (>>=) #-}

instance Monad m => MonadFree (Yoneda f) (ProgramT f m) where
  wrap = ProgramT . wrap . fmap runProgT
  {-# INLINE wrap #-}

instance MonadTrans (ProgramT f) where
  lift = ProgramT . lift
  {-# INLINE lift #-}

instance MonadIO m => MonadIO (ProgramT f m) where
  liftIO = lift . liftIO
  {-# INLINE liftIO #-}

-- | This is a version of 'ProgramView' that is abstracted over the
-- base 'Monad'.
data ProgramViewT f m a where
  Return :: a   -> ProgramViewT f m a
  (:>>=) :: f a -> (a -> ProgramT f m b) -> ProgramViewT f m b

-- | This is a version of 'singleton' that is abstracted over the base
-- 'Monad'.
singletonT :: Monad m => f a -> ProgramT f m a
singletonT f = liftF (liftYoneda f)
{-# INLINE singletonT #-}

-- | This is a version of 'view' that is abstracted over the base
-- 'Monad'.
viewT :: Monad m => ProgramT f m a -> m (ProgramViewT f m a)
viewT (ProgramT x) = view' x
  where view' (FreeT m) = do
          m >>= \case
            Pure k            -> return (Return k)
            Free (Yoneda f b) -> return (b :>>= (ProgramT . f))
{-# INLINE viewT #-}

--------------------------------------------------------------------------------
-- Example #1: a simple teletype monad

-- The TeletypeF functor defines which actions a Teletype program may perform.
data TeletypeF r where
  PutStrLn    :: String  -> r  -> TeletypeF r
  GetLine     :: (String -> r) -> TeletypeF r
  ExitSuccess :: TeletypeF r

-- Natural functor instance.
instance Functor TeletypeF where
  fmap f (PutStrLn s r) = PutStrLn s (f r)
  fmap f (GetLine k)    = GetLine (f . k)
  fmap _ ExitSuccess    = ExitSuccess

-- A Teletype 'Program' is defined as a Program over the
-- TeletypeF functor.
type Teletype m a = ProgramT TeletypeF m a

putStrLn' :: Monad m => String -> Teletype m ()
putStrLn' s = singletonT (PutStrLn s ())

getLine' :: Monad m => Teletype m String
getLine' = singletonT (GetLine id)

exitSuccess' :: Monad m => Teletype m ()
exitSuccess' = singletonT ExitSuccess

echo :: Monad m => Teletype m ()
echo = do
  s <- getLine'
  putStrLn' s
  exitSuccess'

runTeletypeIO :: MonadIO m => Teletype m a -> m a
runTeletypeIO = eval <=< viewT
  where
    eval :: MonadIO m => ProgramViewT TeletypeF m a -> m a
    eval (Return x)            = return x
    eval (ExitSuccess  :>>= _) = liftIO exitSuccess
    eval (PutStrLn s r :>>= k) = liftIO (putStrLn s) >> runTeletypeIO (k r)
    eval (GetLine k    :>>= h) = liftIO getLine >>= runTeletypeIO . (h . k)

--------------------------------------------------------------------------------
-- Example #2: a proper List monad transformer
--
-- Thanks to Heinrich Apfelmus for this example from 'operational'

data MPlus m a where
  MZero :: MPlus m a
  MPlus :: ListT m a -> ListT m a -> MPlus m a
type ListT m a = ProgramT (MPlus m) m a

instance Monad m => Functor (MPlus m) where
  fmap _ MZero       = MZero
  fmap f (MPlus m n) = on MPlus (fmap f) m n

instance Monad m => MonadPlus (ProgramT (MPlus m) m) where
  mzero     = singletonT MZero
  mplus m n = singletonT (MPlus m n)

runListT :: Monad m => ListT m a -> m [a]
runListT = eval <=< viewT
  where
    eval :: Monad m => ProgramViewT (MPlus m) m a -> m [a]
    eval (Return x)         = return [x]
    eval (MZero :>>= _)     = return []
    eval (MPlus m n :>>= k) =
      liftM2 (++) (runListT $ m >>= k) (runListT $ n >>= k)

testListT :: IO ()
testListT = void $ runListT $ do
  n <- choice [1..5::Int]
  lift . putStrLn $ "You chose the number " ++ show n
  where
    choice = foldr1 mplus . map return

-- testing the monad laws, from the Haskellwiki
-- http://www.haskell.org/haskellwiki/ListT_done_right#Order_of_printing
--
-- t1 and t2 have to print the same sequence of letters
testListTLaws :: IO ()
testListTLaws = do
  putStrLn "Test #"
  runListT t1
  putStrLn "\nTest #2"
  runListT t2
  putStrLn "\nDone"
  return ()
  where
    t1 = ((a `mplus` a) >> b) >> c
    t2 = (a `mplus` a) >> (b >> c)

    a,b,c :: ListT IO ()
    [a,b,c] = map (lift . putChar) ['a'..'c']

--------------------------------------------------------------------------------
-- Example #3: Oleg Kiselyov's 'LogicT' transformer
--
-- Thanks to Heinrich Apfelmus for this example from 'operational'

type LogicT m a = ListT m a

    -- msplit  is the lift of a function  split  in the base monad
msplit :: Monad m => LogicT m a -> LogicT m (Maybe (a, LogicT m a))
msplit = lift . split

    -- split  in the base monad
split :: Monad m => LogicT m a -> m (Maybe (a, LogicT m a))
split = eval <=< viewT where
  eval :: Monad m => ProgramViewT (MPlus m) m a -> m (Maybe (a, LogicT m a))
  eval (MZero     :>>= _) = return Nothing
  eval (MPlus m n :>>= k) = do
    ma <- split (m >>= k)
    case ma of
      Nothing     -> split (n >>= k)
      Just (a,m') -> return $ Just (a, m' `mplus` (n >>= k))
                                -- inefficient!
                                -- `mplus` will add another (>>= return)
                                -- to  n  each time it's called.
                                -- Curing this is not easy.

-- main interpreter, section 6 in the paper
-- returns the first result, if any; may fail
observe :: Monad m => LogicT m a -> m a
observe m = (fst . fromJust) `liftM` split m

bagOfN :: Monad m => Maybe Int -> LogicT m a -> LogicT m [a]
bagOfN (Just n) m | n <= 0 = return []
bagOfN n m                 = msplit m >>= bagofN'
    where
      bagofN' Nothing         = return []
      bagofN' (Just (x,m'))   = (x:) `liftM` bagOfN (fmap pred  n) m'
        where pred n = n-1

-- interleave
interleave :: Monad m => LogicT m a -> LogicT m a -> LogicT m a
interleave m1 m2 = do
    r <- msplit m1
    case r of
        Nothing      -> m2
        Just (a,m1') -> return a `mplus` interleave m2 m1'

--------------------------------------------------------------------------------
-- Example #4: Koen Claessen/Gabriel Gonzalez's Concurrency Monad.

data ThreadF t where
  Fork  :: t -> t -> ThreadF t
  Yield :: t -> ThreadF t
  Done  :: ThreadF t

instance Functor ThreadF where
  fmap f (Fork l r) = Fork (f l) (f r)
  fmap f (Yield t)  = Yield (f t)
  fmap _ Done       = Done

type Thread m a = ProgramT ThreadF m a

yield :: Monad m => Thread m ()
yield = singletonT (Yield ())

done :: Monad m => Thread m r
done = singletonT Done

cFork :: Monad m => Thread m Bool
cFork = singletonT (Fork False True)

fork :: (Monad m) => Thread m a -> Thread m ()
fork thread = do
    child <- cFork
    when child $ do
        thread
        done

runThreadT :: Monad m => Thread m a -> m ()
runThreadT = (schedule . Seq.singleton)
  where
    schedule :: Monad m => Seq (Thread m a) -> m ()
    schedule l = case (viewl l) of
      EmptyL    -> return ()
      (x :< ls) -> viewT x >>= eval ls

    eval :: Monad m => Seq (Thread m a) -> ProgramViewT ThreadF m a -> m ()
    -- Fork process
    eval q (Fork l r :>>= k) = schedule (k l <| (q |> k r))

    -- Switch process
    eval q (Yield t  :>>= k) = schedule (q |> k t)

    -- Thread died. Remove from queue.
    eval q (Done     :>>= _) = schedule q
    eval q (Return _)        = schedule q

runThreadTest1 :: IO ()
runThreadTest1 = runThreadT $ do
    lift $ putStrLn "Forking thread #1"
    fork thread1
    yield
    lift $ putStrLn "Forking thread #2"
    fork thread2
    yield

thread1 :: Thread IO ()
thread1 = forM_ [1..10] $ \i -> do
    lift $ print i
    yield

thread2 :: Thread IO ()
thread2 = replicateM_ 3 $ do
    lift $ putStrLn "Hello"
    yield

--------------------------------------------------------------------------------
-- Example #5: State Monad
--
-- Thanks to Heinrich Apfelmus for this example from 'operational'

data StateI s a where
    Get :: StateI s s
    Put :: s -> StateI s ()

type State s a = Program (StateI s) a

evalState :: State s a -> s -> a
evalState = eval . view
    where
    eval :: ProgramView (StateI s) a -> (s -> a)
    eval (Return x)     = const x
    eval (Get   :>>= k) = \s -> evalState (k s ) s
    eval (Put s :>>= k) = \_ -> evalState (k ()) s

put :: Monad m => s -> StateT s m ()
put = singletonT . Put

get :: Monad m => StateT s m s
get = singletonT Get

testState :: Int -> Int
testState = evalState $ do
        x <- get
        put (x+2)
        get

type StateT s m a = ProgramT (StateI s) m a

evalStateT :: Monad m => StateT s m a -> s -> m a
evalStateT m = \s -> viewT m >>= \p -> eval p s
    where
    eval :: Monad m => ProgramViewT (StateI s) m a -> (s -> m a)
    eval (Return x)     = \_ -> return x
    eval (Get   :>>= k) = \s -> evalStateT (k s ) s
    eval (Put s :>>= k) = \_ -> evalStateT (k ()) s

testStateT = evalStateT $ do
    x <- get
    lift $ putStrLn "Hello StateT"
    put (x+1)