Recording and replaying arbitrary FreeT computation. For what happened next see https://github.com/fizruk/replay-free

Replay.hs

{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE DeriveFunctor #-}
module Main where

import Control.Monad.Trans
import Control.Monad.Free.Class
import qualified Control.Monad.Trans.Free as FT
import qualified Control.Monad.Free as F

import Control.Concurrent
import Control.Concurrent.STM

-- | Capturing the notion that @g@ encodes all necessary information to
-- replay @f@ action.
class Replay f g where
  replay :: f a -> g b -> Maybe (g (a, b))

-- | Replay @FreeT f m a@ computation given @Free g ()@ computation log tree.
-- The result is new @FreeT g m@ log tree with leftover @FreeT f m@ computations and
-- unmatched @Free g ()@ log subtrees in leaves.
replayFreeT :: (Replay f g, Functor f, Functor g, Monad m)
            => F.Free g ()                                            -- ^ The log tree.
            -> FT.FreeT f m a                                         -- ^ The computation to replay.
            -> FT.FreeT g m (FT.FreeT f m a, F.Free g ())
replayFreeT r@(F.Pure _) ft = return (ft, r)
replayFreeT r@(F.Free g) (FT.FreeT m) = do
  f <- lift m
  case f of
    FT.Pure x -> return (return x, r)
    FT.Free h ->
      case replay h g of
        Nothing -> return (wrap h, r)
        Just k -> wrap $ fmap (\(ft', r') -> replayFreeT r' ft') k

-- | Run @FreeT f m a@ computation and record actions.
-- The result is a @FreeT g m a@ computation.
--
-- This function is analogous to 'iterT'.
recordFreeT :: (Functor f, Functor g, Monad m)
            => (forall x. f (m x) -> m (g x))   -- ^ How to record each layer of computation.
            -> FT.FreeT f m a                   -- ^ Computation to record.
            -> m (F.Free g a)                   -- ^ The computation log tree.
recordFreeT mapF (FT.FreeT m) = do
  f <- m
  case fmap (recordFreeT mapF) f of
    FT.Pure x -> return (return x)
    FT.Free g -> mapF g >>= return . wrap

-- ===============================================================================
-- Example
-- ===============================================================================

-- | This is our base functor, which describes the list of actions.
data F a
  = Ask (String -> a)       -- ^ Get some input.
  | Fork a (ThreadId -> a)  -- ^ Fork computation.
  | Halt                    -- ^ Abort computation.
  deriving (Functor)

-- | This is a derived data structure which retains 'F' tree structure and
-- stores recorded values for functions in 'F'.
data F' a
  = Ask' (String, a)        -- ^ Recorded input.
  | Fork' a (ThreadId, a)   -- ^ Recorded child ThreadId.
  | Halt'                   -- ^ We don't record anything for halt.
  | Save'                   -- ^ Here we paused our computation to collect the log.
  deriving (Show, Functor)

instance Replay F F' where
  replay (Ask f) (Ask' (s, x)) = Just $ Ask' (s, (f s, x))
  replay (Fork c p) (Fork' c' (pid, p')) = Just $ Fork' (c, c') (pid, (p pid, p'))
  replay Halt Halt' = Just Halt'
  replay _ _ = Nothing

-- DSL commands for F functor
--
-- Note: these can be actually derived automatically using $(makeFree ''F)

-- | Ask user for input.
ask :: MonadFree F m => m String
ask = liftF $ Ask id

-- | Fork computation.
fork :: MonadFree F m => m (Maybe ThreadId)
fork = liftF $ Fork Nothing Just

-- | Halt computation.
halt :: MonadFree F m => m a
halt = liftF Halt

-- | Perform and record an F action in IO monad.
recordF :: F (IO a) -> IO (F' a)
recordF (Ask g) = do
  s <- getLine
  case s of
    "save" -> return Save'
    _ -> do
      x <- g s
      return (Ask' (s, x))
recordF (Fork c p) = do
  v <- atomically newEmptyTMVar
  pid <- forkIO $ c >>= atomically . putTMVar v
  px <- p pid
  cx <- atomically $ takeTMVar v
  return (Fork' cx (pid, px))
recordF Halt = return Halt'

-- | Perform recorded actions (simplified).
evalF' :: F' (IO a) -> IO a
evalF' (Ask' (_, m)) = m
evalF' (Fork' mc (_, mp)) = do
  v <- atomically newEmptyTMVar
  _ <- forkIO $ mc >> atomically (putTMVar v ())
  x <- mp
  _ <- atomically $ takeTMVar v -- this is simply waiting for another thread to finish
  return x
evalF' Halt' = error "halt"
evalF' Save' = error "save"

-- | Sample program.
test :: (MonadFree F m, MonadIO m) => m ()
test = do
  name <- prompt "What's your name?"
  liftIO $ putStrLn ("Hello, " ++ name ++ "!")

  x <- prompt "What do you want to do (save/halt/continue)?"
  case x of
    "halt" -> halt
    _ -> liftIO $ putStrLn "Continuing..."

  mpid <- fork
  liftIO . putStrLn $
    case mpid of
      Nothing   -> "I am child!"
      Just pid  -> "I am parent! My child is " ++ show pid ++ "."

  y <- prompt $ show mpid ++ ": And the final input!"
  liftIO $ putStrLn y
  where
    prompt s = do
      liftIO $ putStrLn s
      ask

main :: IO ()
main = do
  putStrLn "========================================"
  putStrLn "  Recording"
  putStrLn "========================================"

  logTree <- recordFreeT recordF test

  putStrLn "========================================"
  print logTree

  putStrLn "========================================"
  putStrLn "  Replaying"
  putStrLn "========================================"

  let -- build a replayed computation
      replayed   = replayFreeT logTree test
      -- evaluate leftover computations at the leaves of replayed computation
      -- we ignore unmatched logTree subtrees
      replayed'  = fmap (recordFreeT recordF . fst) replayed
      -- attach computations at leaves to the computation tree
      replayed'' = do
        m <- replayed'
        lift $ do
          putStrLn "========================================"
          putStrLn "  Continuing"
          putStrLn "========================================"
          m
  -- replay and continue computation
  _logTree <- FT.iterT evalF' replayed''

  putStrLn "========================================"

  return ()