-- Examples of Software Transactional Memory
--
-- See also the section in the GHC Users Guide:
--  http://www.haskell.org/ghc/docs/latest/html/users_guide
-- Simon PJ's list of papers at
--  http://research.microsoft.com/en-us/um/people/simonpj/papers/stm/
-- and Chapters 24 and 28 of the book "Real World Haskell":
--  http://book.realworldhaskell.org/read/
-----------------------------------------------------------------------------

module Main where

import Array
import Control.Parallel
import Control.Parallel.Strategies
import Control.Monad
import Control.Concurrent
import Control.Concurrent.STM
import System(getArgs)

-- MVars: -------------------------------------------------------

threadA :: Int -> MVar Int -> MVar Float  -> MVar () -> IO ()
threadA m valueToSendMVar valueReceiveMVar a_done
 = do -- some work
    -- now perform rendezvous by sending 72
    putMVar valueToSendMVar m -- send value
    v <- takeMVar valueReceiveMVar
    putStrLn $ (show v) ++ "; should be "++(show $ (fromIntegral m)*1.2)
    putMVar a_done ()

threadB :: Int -> MVar Int -> MVar Float -> MVar () -> IO ()
threadB n valueToReceiveMVar valueToSendMVar b_done
 = do -- some work
      -- now perform rendezvous by waiting on value
    z <- takeMVar valueToReceiveMVar
    putMVar valueToSendMVar (1.2 * fromIntegral z)
    -- continue with other work
    putMVar b_done ()

sync_on :: MVar () -> MVar () -> IO ()
sync_on a b = do
               takeMVar a
               takeMVar b
               
test0 :: Int -> Int -> IO ()
test0 m n = do aMVar <- newEmptyMVar
               bMVar <- newEmptyMVar
               a_done <- newEmptyMVar
               b_done <- newEmptyMVar
               forkIO (threadA m aMVar bMVar a_done)
               forkIO (threadB n aMVar bMVar b_done)
               -- threadDelay 1000 -- wait for threadA and threadB to finish (sleazy)
               sync_on a_done b_done

-- TVars: -------------------------------------------------------

withdraw :: TVar Int -> Int -> STM ()
withdraw acc n = do
                  bal <- readTVar acc
                  if bal<n 
                    then retry
                    else writeTVar acc (bal-n)

deposit :: TVar Int -> Int -> STM ()
deposit acc n = do
                 v <- readTVar acc
                 writeTVar acc (v+n) 

balance :: TVar Int -> STM Int
balance acc = do { b <- readTVar acc ; return b }

-- start 2 threads, doing withdraw and balance concurrently
test1 m n = do 
             my_acc <- (atomically $ newTVar 0)
             fid1_done <- newEmptyMVar
             fid2_done <- newEmptyMVar
             fid1 <- forkIO $ atomically (withdraw my_acc m) >> putMVar fid1_done ()
             fid2 <- forkIO $ atomically (deposit  my_acc n) >> putMVar fid2_done ()
             takeMVar fid1_done
             takeMVar fid2_done
             b <- atomically $ balance my_acc
             putStrLn $ "Done with balance: "++(show b)++"; should be "++(show (n-m))

test2 m n = error "test2 undefined"

-- good old nfib

nfib :: Int -> Int
nfib 0 = 1
nfib 1 = 1
nfib n = nfib (n-1) + nfib (n-2) + 1

-- STM parallelism

test4 m n = do
              putStrLn $ "Calculating nfib "++(show m)++" and nfib "++(show n)++" in parallel, using STM"
              fid1_done <- newEmptyMVar
              fid2_done <- newEmptyMVar
              fid1 <- forkIO (do { let x = (nfib m) in x `seq` putMVar fid1_done x } )
              fid2 <- forkIO (do { let x = (nfib n) in x `seq` putMVar fid2_done x } )
              x1 <- readMVar fid1_done
              x2 <- readMVar fid2_done
              putStrLn $ "nfib m = "++(show x1)
              putStrLn $ "nfib n = "++(show x2)

-- GpH parallelism (with par's)

pfib :: Int -> Int -> Int
pfib 0 _ = 1
pfib 1 _ = 1
pfib n t | n<=t      = nfib n
         | otherwise = x `par` y `seq` x + y + 1
         	       where x = pfib (n-1) t
         		     y = pfib (n-2) t

test5 m n = do
              putStrLn $ "Calculating pfib "++(show m)++" with threshold "++(show n)++" in parallel, using par's"
              let x1 = pfib m n
              putStrLn $ "pfib m n = "++(show x1)

-----------------------------------------------------------------------------
-- test wrapper

main = do args <- getArgs
          let 
            v = read (args!!0) -- version
            m = read (args!!1) 
            n = read (args!!2)
          let (str, exe) = 
           	 case v of 
           	  0 -> ("Rendevous example", test0 m n)
           	  1 -> ("STM-style withdraw/deposit example", test1 m n)
           	  2 -> ("undefined", test2 m n)
           	  -- 3 -> ("STM-style queues", test3 m n)
           	  4 -> ("STM-style nfib on 2 threads", test4 m n)
           	  5 -> ("GpH-style parfib", test5 m n)
                  _ -> error "Unknown version. Usage: ./STM <version> <m> <n> where <version> must be between 0 and 5"
          putStrLn str
          exe
          putStrLn "Done"