mealy.hs

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE StaticPointers #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -Wall #-}

import Control.Concurrent
import Control.Distributed.Closure
import Control.Monad
import Data.Binary
import qualified Data.ByteString.Lazy as LBS
import Data.Functor
import GHC.Generics
import System.IO.Unsafe

-- | 'Log' represents a state value which may change in a clock cycle.
newtype Log
  = Log Int
  deriving (Generic, Show)

instance Binary Log

-- | 'Ping' represents an input event value in each clock cycle.
newtype Ping
  = Ping Int
  deriving (Generic, Show)

instance Binary Ping

-- | 'Pong' represents an output value in each clock cycle.
newtype Pong
  = Pong Int
  deriving (Generic, Show)

instance Binary Pong

-- | 'Mealy' models a Mealy machine. In each clock cycle, a Mealy machine
-- consumes a state value and an input value, and produces an updated state
-- value and an output value.
--
-- Our 'Mealy' runs in 'IO' since it needs side effects like logging for our
-- demo. The state value is not managed by 'IO'.
--
-- The 'Mealy' values we work with should be CAFs, which mean they capture no
-- free variables. Thus we can construct their static pointers and 'Closure's.
newtype Mealy
  = Mealy
      { runMealy :: Log -> Ping -> IO (Log, Pong)
      }

-- | For our demo, simply passing a 'Closure' 'Mealy' value would be too
-- trivial. So we pass 'MealyPack's instead.
--
-- A 'MealyPack' is a serializable Mealy machine. But the state value is
-- captured into the 'Closure' and not reflected in the 'MealyPack' type
-- definition. By passing 'MealyPack's in messages, we demonstrate the ability
-- to serialize function closures.
newtype MealyPack
  = MealyPack
      { unMealyPack :: Closure (Ping -> IO (MealyPack, Pong))
      }

-- | Given the 'Closure's of a 'Mealy' machine and a state value, this function
-- constructs a 'MealyPack' which captures the state value in the closure.
makeMealyPack :: Closure Mealy -> Closure Log -> MealyPack
makeMealyPack cm cs =
  MealyPack
    ( ( closure
          ( static
              ( \t f i ->
                  (\(s', o) -> (t (cpure (closure (static Dict)) s'), o)) <$> f i
              )
          )
          `cap` (closure (static makeMealyPack) `cap` cduplicate cm)
      )
        `cap` ((closure (static runMealy) `cap` cm) `cap` cs)
    )

-- | Given the incoming/outgoing message channels, this function forks a thread
-- which repeatedly:
--
-- 1. Reads an incoming message, decodes a Mealy machine closure and an output
--    value.
-- 2. From the output value, computes an input value for the current clock
--    cycle.
-- 3. Invokes the Mealy machine closure, and produces a new closure and a new
--    output value.
-- 4. Encodes and sends the new closure and the new output value via the
--    outgoing message channel.
--
-- 'forkMealy' has no knowledge about the actual 'Mealy' values in the program
-- which implement our Mealy machine logic. All it knows is the type definition
-- of 'MealyPack' and how to use it to serialize/deserialize closures and run
-- the deserialized functions.
forkMealy :: MVar LBS.ByteString -> MVar LBS.ByteString -> IO ()
forkMealy mi mo =
  void
    $ forkIO
    $ forever
    $ do
      (c, Pong x) <- decode <$> takeMVar mi
      (unMealyPack -> c', Pong y) <- unclosure c (Ping x)
      putMVar mo (encode (c', Pong y))

-- | Here is the actual Mealy machine implementation for the demo. It modifies
-- the state value by adding the input value, and produces the output value by
-- incrementing the input value by 1. It also logs the current 'ThreadId' and
-- the state/input values to the console.
mealy :: Mealy
mealy =
  Mealy
    ( \(Log l) (Ping i) -> do
        tid <- myThreadId
        withMVar printLock ($ show (tid, Log l, Ping i))
          $> (Log (l + 1), Pong (i + l))
    )

-- | We create a pair of message channels and a pair of threads, each thread
-- will repeatedly deserialize and run a Mealy machine closure. After putting an
-- initial message, the whole ping/pong loop will be activated.
main :: IO ()
main = do
  a2b <- newEmptyMVar
  b2a <- newEmptyMVar
  forkMealy a2b b2a
  forkMealy b2a a2b
  putMVar
    a2b
    ( encode
        ( unMealyPack
            ( makeMealyPack
                (closure (static mealy))
                (cpure (closure (static Dict)) (Log 0))
            ),
          Pong 1
        )
    )

-- | This is a locked version of 'putStrLn', to ensure different threads don't
-- mess up the console.
{-# NOINLINE printLock #-}
printLock :: MVar (String -> IO ())
printLock = unsafePerformIO $ newMVar putStrLn