Benchmarks Game/Parallel/Chameneos

{- The Computer Language Benchmarks Game
   http://shootout.alioth.debian.org/
   Written by Tom Pledger, 13 Nov 2006. modified by Don Stewart
   Updated for chameneos-redux by Spencer Janssen, 27 Nov 2007
   Improved concurrency by Spencer Janssen, 29 Sep 2007 -}
 
import Control.Concurrent
import Control.Monad
import Data.Char
import System.Environment
import System.IO
import GHC.Conc (forkOnIO)
 
default(Int)
 
data Colour = Blue | Red | Yellow deriving (Show, Eq, Enum)
 
complement a b | a == b = a
complement a b = case a of
    Blue   -> case b of Red  -> Yellow; _ -> Red
    Red    -> case b of Blue -> Yellow; _ -> Blue
    Yellow -> case b of Blue -> Red;    _ -> Blue
 
data MP = MP !Int !(Maybe (Colour, ThreadId, MVar (Colour, ThreadId)))
 
arrive mpv finish c0 = do
    tid <- myThreadId
    let inc x = (fromEnum (tid == x) +)
        go !c !t !b = do
            MP q w <- takeMVar mpv
            case w of
                _ | q == 0 -> do
                    putMVar mpv $ MP 0 w
                    putMVar finish (t, b)
 
                Nothing -> do
                     waker <- newEmptyMVar
                     putMVar mpv $ MP q (Just (c, tid, waker))
                     (c', tid') <- takeMVar waker
                     go c' (t+1) $ inc tid' b
 
                Just (k, tid', waker) -> do
                    putMVar mpv $ MP (q-1) Nothing
                    let !c' = complement k c
                    putMVar waker (c', tid)
                    go c' (t+1) $ inc tid' b
    go c0 0 0
 
showN = unwords . map ((digits !!) . digitToInt) . show
 
digits = words "zero one two three four five six seven eight nine"
 
run cpu n cs = do
    fs    <- replicateM (length cs) newEmptyMVar
    mpv   <- newMVar $ MP n Nothing
    zipWithM ((forkOnIO cpu .) . arrive mpv) fs cs
    return $ do
        ns    <- mapM takeMVar fs
 
        putStrLn . map toLower . unwords . ([]:) . map show $ cs
        putStr . map toLower . unlines $ [unwords [show n, showN b] | (n, b) <- ns]
        putStrLn . (" "++) . showN . sum . map fst $ ns
        putStrLn ""
 
main = do
    putStrLn . map toLower . unlines $
        [unwords [show a, "+", show b, "->", show $ complement a b]
            | a <- [Blue ..], b <- [Blue ..]]
 
    n <- readIO . head =<< getArgs
    collect1 <- run 1 n [Blue ..]
    collect2 <- run 2 n [Blue, Red, Yellow, Red, Yellow, Blue, Red, Yellow, Red, Blue]
    collect1
    collect2