benchmark.hs

{-# LANGUAGE Arrows #-}

module Main where

import Prelude hiding (filter, id, (.), unzip, mapM_)
import Control.Category
import Control.Arrow
import Criterion
import Criterion.Main (defaultMain)
import Data.Time
import System.Random.MWC
import System.IO
import Control.Arrow.Machine
import Data.Foldable
import qualified Data.Vector as V
import qualified Data.IntMap as IM
import Control.Monad.Trans.State
import Control.Monad.Trans


main :: IO ()
main = defaultMain
    -- Benchmark 1
    [ bgroup "benchmark1/100-nodes-100-steps"
          [ bench "machinecell" $ whnfIO $ benchmark1 100 100 ]
    , bgroup "benchmark1/100-nodes-1000-steps"
          [ bench "machinecell" $ whnfIO $ benchmark1 100 1000 ]
    , bgroup "benchmark1/1000-nodes-100-steps"
          [ bench "machinecell" $ whnfIO $ benchmark1 1000 100 ]
    , bgroup "benchmark1/1000-nodes-1000-steps"
          [ bench "machinecell" $ whnfIO $ benchmark1 1000 1000 ]
    ]

unzip :: Functor f => f (a, b) -> (f a, f b)
unzip = fmap fst &&& fmap snd

benchmark1 netsize dur =
  do
    startTime <- getCurrentTime
    randGen <- create
    midTime <- getCurrentTime
    runStateT `flip` benchmark1Arrow netsize $ forM_ [1..dur] $ \step ->
      do
        let str = show step
        evs <- V.replicateM 10 $ liftIO $ uniformR (0,netsize-1) randGen
        forM_ evs $ \ev ->
          do
            pa <- get
            let (ei, pa') = stepRun pa (ev, str)
            liftIO $ mapM_ ePutStrLn (yields ei)
            put pa'
        
    endTime <- getCurrentTime
    return (midTime `diffUTCTime` startTime, endTime `diffUTCTime` midTime)

benchmark1Arrow' :: Int -> ProcessA (->) (Event (Int, String)) (Event String)
benchmark1Arrow' netsize =
    arr unzip >>>
    first (hold 0) >>>
    par unicast (V.replicate netsize id) >>>
    gather

unicast :: (Int, Event String) -> V.Vector sf -> V.Vector (Event String, sf)
unicast (idx, evs) = V.imap $ \i sf -> (if i == idx then evs else noEvent, sf)


benchmark1Arrow :: Int -> ProcessA (->) (Event (Int, String)) (Event String)
benchmark1Arrow netsize = construct $
  do
    let initMap :: IM.IntMap (ProcessA (->) (Event String) (Event String))
        initMap = IM.fromList $ zip [0..] (replicate netsize id)

    runStateT `flip` initMap $ forever $
      do
        (i, str) <- lift await
        mpa <- gets $ IM.lookup i
        maybe (return ()) `flip` mpa $ \pa ->
          do
            let (ei, pa') = stepRun pa str
            mapM_ (lift . yield) (yields ei)
            modify $ IM.insert i pa'


ePutStrLn :: String -> IO ()
ePutStrLn = hPutStrLn stderr