{-# LANGUAGE ScopedTypeVariables #-}

module Control.Concurrent.Parallel3(parallel_) where

import GHC.Conc
import Control.Concurrent.Chan
import Control.Concurrent.MVar
import Control.Concurrent
import Control.Monad
import Control.Exception as E
import System.IO.Unsafe


-- initialise on the main thread, and keep
mainThread :: ThreadId
mainThread = unsafePerformIO $ myThreadId

-- True = kill the thread after it finishes
{-# NOINLINE queue #-}
queue :: Chan (IO a, Bool)
queue = seq mainThread $ unsafePerformIO $ do
    chan <- newChan
    replicateM_ (numCapabilities-1) $ addWorker
    return chan


addWorker :: IO ()
addWorker = do
    forkIO $ f `E.catch` \(e :: SomeException) ->
        throwTo mainThread $ ErrorCall "Control.Concurrent.Parallel: parallel thread died."
    return ()
    where
        f :: IO ()
        f = do
            (act, kill) <- readChan queue
            act
            unless kill f


-- | Run the list of computations in parallel
--   Rule: No thread should get pre-empted (although not a guarantee)
--         On return all actions have been performed
parallel_ :: [IO a] -> IO ()
parallel_ xs | numCapabilities <= 1 = sequence_ xs
parallel_ [] = return ()
parallel_ [x] = x >> return ()
parallel_ (x1:xs) = do
    sem <- newQSem $ 1 - length xs
    forM_ xs $ \x ->
        writeChan queue (x >> signalQSem sem, False)
    x1
    addWorker
    waitQSem sem
    writeChan queue (signalQSem sem, True)
    waitQSem sem