{-# LANGUAGE Arrows #-}
module Main
    (main)
    where

import System.IO
import Data.Array
import Control.Concurrent
import Control.Concurrent.STM
import Control.Monad
import Control.Applicative
import Control.Arrow
import qualified BlockingTransactions.BlockingTransactions as BT
import System.Environment
import System.Random

-- A benchmark that creates a number (param_accounts) of bank accounts
-- and then performs a large number (param_transactions) of transactions on a
-- large number (param_threads) of threads, running in parallel.  Transactions
-- always move one dollar from a larger account to a smaller account.

-- This benchmark is written (hopefull identically) using STM (conventional
-- software transactional memory), BTM (blocking transactions monad),
-- and BTA (blocking transactions arrow).

-- Only the STM implementation is documented, since all three implementations
-- should have the same form.

param_threads :: Int
param_threads = 100

param_accounts :: Int
param_accounts = 100

param_transactions :: Int
param_transactions = 1000

main :: IO ()
main =
    do args <- getArgs
       case args of
           ["stm"] -> benchmarkSTM
           ["btm"] -> benchmarkBTM
           ["bta"] -> benchmarkBTA

benchmarkSTM :: IO ()
benchmarkSTM =
    do -- a flag indicating when we can start
       start <- newTVarIO False
       -- variable indicating when we are finished
       threads_remaining <- newTVarIO param_threads
       -- populate the accounts
       accounts <- liftM (listArray (1,param_accounts)) $
           forM [1..param_accounts] $ \_ -> (newTVarIO =<<) $
               getStdRandom $ randomR (1,param_transactions)
       -- launch worker threads
       forM_ [1..param_threads] $ \_ -> (>> return ()) $ forkIO $
           do -- wait until we can start
              atomically $
                  do can_start <- readTVar start
                     when (not can_start) retry
              -- perform 1-dollar transactions between random accounts
              forM_ [1..param_transactions] $ \_ ->
                  do ac1 <- getStdRandom $ randomR (1,param_accounts)
                     ac2 <- getStdRandom $ randomR (1,param_accounts)
                     atomically $
                         do v1 <- readTVar (accounts ! ac1)
                            v2 <- readTVar (accounts ! ac2)
                            when (v1 > v2) $
                                do writeTVar (accounts ! ac1) $ pred v1
                                   writeTVar (accounts ! ac2) $ succ v2
              -- indicate that we are finished
              atomically $ writeTVar threads_remaining . pred =<<
                               readTVar threads_remaining
              return ()
       -- indicate that we can start (all threads are live)
       atomically $ writeTVar start True
       -- wait until we are finished
       atomically $
           do x <- readTVar threads_remaining
              when (x /= 0) retry
       -- print the answer
       print =<< mapM (atomically . readTVar) (elems accounts)
       return ()

benchmarkBTM :: IO ()
benchmarkBTM =
    do start <- BT.newBVar False
       threads_remaining <- BT.newBVar param_threads
       accounts <- liftM (listArray (1,param_accounts)) $
           forM [1..param_accounts] $ \_ -> (BT.newBVar =<<) $
               getStdRandom $ randomR (1,param_transactions)
       forM_ [1..param_threads] $ \_ -> (>> return ()) $ forkIO $
           do BT.runBTM $
                  do can_start <- BT.readBVar start
                     BT.unless can_start BT.retry
              forM_ [1..param_transactions] $ \_ ->
                  do ac1 <- getStdRandom $ randomR (1,param_accounts)
                     ac2 <- getStdRandom $ randomR (1,param_accounts)
                     BT.runBTM $
                         do v1 <- BT.readBVar (accounts ! ac1)
                            v2 <- BT.readBVar (accounts ! ac2)
                            BT.when ((>) <$> v1 <*> v2) $
                                do BT.writeBVar (accounts ! ac1) $ fmap pred v1
                                   BT.writeBVar (accounts ! ac2) $ fmap succ v2
              BT.runBTM $ BT.writeBVar threads_remaining . fmap pred =<<
                              BT.readBVar threads_remaining
              return ()
       BT.pokeBVar start True
       BT.runBTM $
           do x <- BT.readBVar threads_remaining
              BT.when (fmap (/= 0) x) BT.retry
       print =<< mapM (\v -> BT.runBTM $ BT.readBVar v) (elems accounts)
       return ()

benchmarkBTA :: IO ()
benchmarkBTA =
    do start <- BT.newBVar False
       threads_remaining <- BT.newBVar param_threads
       accounts <- liftM (listArray (1,param_accounts)) $
           forM [1..param_accounts] $ \_ -> (BT.newBVar =<<) $
               getStdRandom $ randomR (1,param_transactions)
       forM_ [1..param_threads] $ \_ -> (>> return ()) $ forkIO $
           do flip BT.runBTA () $ proc () ->
                  do can_start <- BT.fetchBVar start -< ()
                     BT.retryUnless -< can_start
              forM_ [1..param_transactions] $ \_ ->
                  do ac1 <- getStdRandom $ randomR (1,param_accounts)
                     ac2 <- getStdRandom $ randomR (1,param_accounts)
                     flip BT.runBTA () $ proc () ->
                         do v1 <- BT.fetchBVar (accounts ! ac1) -< ()
                            v2 <- BT.fetchBVar (accounts ! ac2) -< ()
                            case () of
                                () | v1 > v2 ->
                                    do BT.storeBVar (accounts ! ac1) -< pred v1
                                       BT.storeBVar (accounts ! ac2) -< succ v2
                                () | otherwise -> returnA -< ()
              flip BT.runBTA () $ BT.storeBVar threads_remaining <<< arr pred
                                       <<< BT.fetchBVar threads_remaining
              return ()
       BT.pokeBVar start True
       flip BT.runBTA () $
            BT.retryWhen <<< arr (/= 0) <<< BT.fetchBVar threads_remaining
       print =<< mapM (flip BT.runBTA () . BT.fetchBVar) (elems accounts)
       return ()