    {-# LANGUAGE CPP #-}

    -- | Gang Primitives.
    module Data.Array.Repa.Eval.Gang
            ( theGang
            , Gang, forkGang, gangSize, gangIO, gangST)
    where
    import GHC.IO
    import GHC.ST
    import GHC.Conc                 (forkOn)
    import Control.Concurrent.MVar
    import Control.Exception        (assert)
    import Control.Monad
    import GHC.Conc                 (numCapabilities)
    import System.IO

    -- TheGang --------------------------------------------------------------------
    -- | This globally shared gang is auto-initialised at startup and shared by all
    --   Repa computations.
    --
    --   In a data parallel setting, it does not help to have multiple gangs
    --   running at the same time. This is because a single data parallel
    --   computation should already be able to keep all threads busy. If we had
    --   multiple gangs running at the same time, then the system as a whole would
    --   run slower as the gangs would contend for cache and thrash the scheduler.
    --
    --   If, due to laziness or otherwise, you try to start multiple parallel
    --   Repa computations at the same time, then you will get the following
    --   warning on stderr at runtime:
    --
    -- @Data.Array.Repa: Performing nested parallel computation sequentially.
    --    You've probably called the 'compute' or 'copy' function while another
    --    instance was already running. This can happen if the second version
    --    was suspended due to lazy evaluation. Use 'deepSeqArray' to ensure that
    --    each array is fully evaluated before you 'compute' the next one.
    -- @
    --
    theGang :: Gang
    {-# NOINLINE theGang #-}
    theGang
     = unsafePerformIO
     $ do   let caps        = numCapabilities
            forkGang caps

    -- Requests -------------------------------------------------------------------
    -- | The 'Req' type encapsulates work requests for individual members of a gang.
    data Req
            -- | Instruct the worker to run the given action.
            = ReqDo        (Int -> IO ())

            -- | Tell the worker that we're shutting the gang down.
            --   The worker should signal that it's receieved the request by
            --   writing to its result var before returning to the caller (forkGang).
            | ReqShutdown

    -- Gang -----------------------------------------------------------------------
    -- | A 'Gang' is a group of threads that execute arbitrary work requests.
    data Gang
            = Gang
            { -- | Number of threads in the gang.
              _gangThreads           :: !Int

              -- | Workers listen for requests on these vars.
            , _gangRequestVars       :: [MVar Req]

              -- | Workers put their results in these vars.
            , _gangResultVars        :: [MVar ()]

              -- | Indicates that the gang is busy.
            , _gangBusy              :: MVar Bool
            }

    instance Show Gang where
      showsPrec p (Gang n _ _ _)
            = showString "<<"
            . showsPrec p n
            . showString " threads>>"

    -- | O(1). Yield the number of threads in the 'Gang'.
    gangSize :: Gang -> Int
    gangSize (Gang n _ _ _)
            = n

    -- | Fork a 'Gang' with the given number of threads (at least 1).
    forkGang :: Int -> IO Gang
    forkGang n
     = assert (n > 0)
     $ do
            -- Create the vars we'll use to issue work requests.
            mvsRequest     <- sequence $ replicate n $ newEmptyMVar

            -- Create the vars we'll use to signal that threads are done.
            mvsDone        <- sequence $ replicate n $ newEmptyMVar

            -- Add finalisers so we can shut the workers down cleanly if they
            -- become unreachable.
            zipWithM_ (\varReq varDone
                            -> mkWeakMVar varReq (finaliseWorker varReq varDone))
                    mvsRequest
                    mvsDone

            -- Create all the worker threads
            zipWithM_ forkOn [0..]
                    $ zipWith3 gangWorker
                            [0 .. n-1] mvsRequest mvsDone

            -- The gang is currently idle.
            busy   <- newMVar False

            return $ Gang n mvsRequest mvsDone busy

    -- | The worker thread of a 'Gang'.
    --   The threads blocks on the MVar waiting for a work request.
    gangWorker :: Int -> MVar Req -> MVar () -> IO ()
    gangWorker threadId varRequest varDone
     = do
            -- Wait for a request
            req     <- takeMVar varRequest

            case req of
             ReqDo action
              -> do -- Run the action we were given.
                    action threadId

                    -- Signal that the action is complete.
                    putMVar varDone ()

                    -- Wait for more requests.
                    gangWorker threadId varRequest varDone

             ReqShutdown
              ->    putMVar varDone ()

    -- | Finaliser for worker threads.
    --   We want to shutdown the corresponding thread when it's MVar becomes
    --   unreachable.
    --   Without this Repa programs can complain about "Blocked indefinitely
    --   on an MVar" because worker threads are still blocked on the request
    --   MVars when the program ends. Whether the finalizer is called or not
    --   is very racey. It happens about 1 in 10 runs when for the
    --   repa-edgedetect benchmark, and less often with the others.
    --
    --   We're relying on the comment in System.Mem.Weak that says
    --    "If there are no other threads to run, the runtime system will
    --     check for runnablefinalizers before declaring the system to be
    --     deadlocked."
    --
    --   If we were creating and destroying the gang cleanly we wouldn't need
    --     this, but theGang is created with a top-level unsafePerformIO.
    --     Hacks beget hacks beget hacks...
    --
    finaliseWorker :: MVar Req -> MVar () -> IO ()
    finaliseWorker varReq varDone
     = do   putMVar varReq ReqShutdown
            takeMVar varDone
            return ()

    -- | Issue work requests for the 'Gang' and wait until they complete.
    --
    --   If the gang is already busy then print a warning to `stderr` and just
    --   run the actions sequentially in the requesting thread.
    gangIO  :: Gang
            -> (Int -> IO ())
            -> IO ()

    {-# NOINLINE gangIO #-}
    gangIO gang@(Gang _ _ _ busy) action
     = do   b <- swapMVar busy True
            if b
             then do
                    seqIO gang action

             else do
                    parIO gang action
                    _ <- swapMVar busy False
                    return ()

    -- | Run an action on the gang sequentially.
    seqIO   :: Gang -> (Int -> IO ()) -> IO ()
    seqIO (Gang n _ _ _) action
     = do   hPutStr stderr
             $ unlines
             [ "Data.Array.Repa: Performing nested parallel computation sequentially."
             , "  You've probably called the 'compute' or 'copy' function while another"
             , "  instance was already running. This can happen if the second version"
             , "  was suspended due to lazy evaluation. Use 'deepSeqArray' to ensure"
             , "  that each array is fully evaluated before you 'compute' the next one."
             , "" ]

            mapM_ action [0 .. n-1]

    -- | Run an action on the gang in parallel.
    parIO   :: Gang -> (Int -> IO ()) -> IO ()
    parIO (Gang _ mvsRequest mvsResult _) action
     = do
            -- Send requests to all the threads.
            mapM_ (\v -> putMVar v (ReqDo action)) mvsRequest

            -- Wait for all the requests to complete.
            mapM_ takeMVar mvsResult

    -- | Same as 'gangIO' but in the 'ST' monad.
    gangST :: Gang -> (Int -> ST s ()) -> ST s ()
    gangST g p = unsafeIOToST . gangIO g $ unsafeSTToIO . p