{-# LANGUAGE CPP           #-}
{-# LANGUAGE MagicHash     #-}
{-# LANGUAGE UnboxedTuples #-}
{-# OPTIONS_HADDOCK hide #-}
-- |
-- Module      : Data.Array.Accelerate.Async
-- Copyright   : [2009..2020] The Accelerate Team
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--

module Data.Array.Accelerate.Async (

  -- * Asynchronous actions
  Async,

  -- ** Spawning
  async, asyncOn, asyncBound,

  -- ** Querying 'Async's
  wait, poll, cancel,

) where

import Control.Exception
import Control.Concurrent

import GHC.Exts
import GHC.Conc
import GHC.IO


-- We need to execute the main thread asynchronously to give finalisers a chance
-- to run. Make sure to catch exceptions to avoid "blocked indefinitely on MVar"
-- errors.
--
data Async a = Async {-# UNPACK #-} !ThreadId
                     {-# UNPACK #-} !(MVar (Either SomeException a))

-- | Spawn an asynchronous action in a separate thread.
--
async :: IO a -> IO (Async a)
async :: IO a -> IO (Async a)
async = ((IO () -> IO ThreadId) -> IO a -> IO (Async a))
-> (IO () -> IO ThreadId) -> IO a -> IO (Async a)
forall a. a -> a
inline (IO () -> IO ThreadId) -> IO a -> IO (Async a)
forall a. (IO () -> IO ThreadId) -> IO a -> IO (Async a)
asyncUsing IO () -> IO ThreadId
rawForkIO

-- | Like 'async', but using 'forkOn' internally.
--
asyncOn :: Int -> IO a -> IO (Async a)
asyncOn :: Int -> IO a -> IO (Async a)
asyncOn Int
cpu = ((IO () -> IO ThreadId) -> IO a -> IO (Async a))
-> (IO () -> IO ThreadId) -> IO a -> IO (Async a)
forall a. a -> a
inline (IO () -> IO ThreadId) -> IO a -> IO (Async a)
forall a. (IO () -> IO ThreadId) -> IO a -> IO (Async a)
asyncUsing (Int -> IO () -> IO ThreadId
rawForkOn Int
cpu)

-- | Like 'async', but using 'forkOS' internally.
--
asyncBound :: IO a -> IO (Async a)
asyncBound :: IO a -> IO (Async a)
asyncBound = ((IO () -> IO ThreadId) -> IO a -> IO (Async a))
-> (IO () -> IO ThreadId) -> IO a -> IO (Async a)
forall a. a -> a
inline (IO () -> IO ThreadId) -> IO a -> IO (Async a)
forall a. (IO () -> IO ThreadId) -> IO a -> IO (Async a)
asyncUsing IO () -> IO ThreadId
forkOS

asyncUsing :: (IO () -> IO ThreadId) -> IO a -> IO (Async a)
asyncUsing :: (IO () -> IO ThreadId) -> IO a -> IO (Async a)
asyncUsing IO () -> IO ThreadId
fork IO a
action = do
   MVar (Either SomeException a)
var <- IO (MVar (Either SomeException a))
forall a. IO (MVar a)
newEmptyMVar
   ThreadId
tid <- ((forall a. IO a -> IO a) -> IO ThreadId) -> IO ThreadId
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask (((forall a. IO a -> IO a) -> IO ThreadId) -> IO ThreadId)
-> ((forall a. IO a -> IO a) -> IO ThreadId) -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
restore ->
            IO () -> IO ThreadId
fork (IO () -> IO ThreadId) -> IO () -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ IO a -> IO (Either SomeException a)
forall e a. Exception e => IO a -> IO (Either e a)
try (IO a -> IO a
forall a. IO a -> IO a
restore IO a
action) IO (Either SomeException a)
-> (Either SomeException a -> IO ()) -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MVar (Either SomeException a) -> Either SomeException a -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar (Either SomeException a)
var
   Async a -> IO (Async a)
forall (m :: * -> *) a. Monad m => a -> m a
return (ThreadId -> MVar (Either SomeException a) -> Async a
forall a. ThreadId -> MVar (Either SomeException a) -> Async a
Async ThreadId
tid MVar (Either SomeException a)
var)

-- | Block the calling thread until the computation completes, then return the
-- result.
--
{-# INLINE wait #-}
wait :: Async a -> IO a
wait :: Async a -> IO a
wait (Async ThreadId
_ MVar (Either SomeException a)
var) = (SomeException -> IO a)
-> (a -> IO a) -> Either SomeException a -> IO a
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> IO a
forall e a. Exception e => e -> IO a
throwIO a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either SomeException a -> IO a)
-> IO (Either SomeException a) -> IO a
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MVar (Either SomeException a) -> IO (Either SomeException a)
forall a. MVar a -> IO a
readMVar MVar (Either SomeException a)
var

-- | Test whether the asynchronous computation has already completed. If so,
-- return the result, else 'Nothing'.
--
{-# INLINE poll #-}
poll :: Async a -> IO (Maybe a)
poll :: Async a -> IO (Maybe a)
poll (Async ThreadId
_ MVar (Either SomeException a)
var) =
  IO (Maybe a)
-> (Either SomeException a -> IO (Maybe a))
-> Maybe (Either SomeException a)
-> IO (Maybe a)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Maybe a -> IO (Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing) ((SomeException -> IO (Maybe a))
-> (a -> IO (Maybe a)) -> Either SomeException a -> IO (Maybe a)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> IO (Maybe a)
forall e a. Exception e => e -> IO a
throwIO (Maybe a -> IO (Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe a -> IO (Maybe a)) -> (a -> Maybe a) -> a -> IO (Maybe a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Maybe a
forall a. a -> Maybe a
Just)) (Maybe (Either SomeException a) -> IO (Maybe a))
-> IO (Maybe (Either SomeException a)) -> IO (Maybe a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MVar (Either SomeException a)
-> IO (Maybe (Either SomeException a))
forall a. MVar a -> IO (Maybe a)
tryReadMVar MVar (Either SomeException a)
var

-- | Cancel a running asynchronous computation.
--
{-# INLINE cancel #-}
cancel :: Async a -> IO ()
cancel :: Async a -> IO ()
cancel (Async ThreadId
tid MVar (Either SomeException a)
_) = ThreadId -> AsyncException -> IO ()
forall e. Exception e => ThreadId -> e -> IO ()
throwTo ThreadId
tid AsyncException
ThreadKilled


-- A version of 'forkIO' that does not include the outer exception handler. This
-- saves a bit of time when we will be installing our own exception handler.
--
-- Stolen from Simon Marlow's 'async' package (BSD).
--
{-# INLINE rawForkIO #-}
rawForkIO :: IO () -> IO ThreadId
rawForkIO :: IO () -> IO ThreadId
rawForkIO IO ()
action = (State# RealWorld -> (# State# RealWorld, ThreadId #))
-> IO ThreadId
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, ThreadId #))
 -> IO ThreadId)
-> (State# RealWorld -> (# State# RealWorld, ThreadId #))
-> IO ThreadId
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s ->
  case IO () -> State# RealWorld -> (# State# RealWorld, ThreadId# #)
forall a.
a -> State# RealWorld -> (# State# RealWorld, ThreadId# #)
fork# IO ()
action State# RealWorld
s of
    (# State# RealWorld
s', ThreadId#
tid #) -> (# State# RealWorld
s', ThreadId# -> ThreadId
ThreadId ThreadId#
tid #)

{-# INLINE rawForkOn #-}
rawForkOn :: Int -> IO () -> IO ThreadId
rawForkOn :: Int -> IO () -> IO ThreadId
rawForkOn (I# Int#
cpu) IO ()
action = (State# RealWorld -> (# State# RealWorld, ThreadId #))
-> IO ThreadId
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, ThreadId #))
 -> IO ThreadId)
-> (State# RealWorld -> (# State# RealWorld, ThreadId #))
-> IO ThreadId
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s ->
  case Int#
-> IO () -> State# RealWorld -> (# State# RealWorld, ThreadId# #)
forall a.
Int# -> a -> State# RealWorld -> (# State# RealWorld, ThreadId# #)
forkOn# Int#
cpu IO ()
action State# RealWorld
s of
    (# State# RealWorld
s', ThreadId#
tid #) -> (# State# RealWorld
s', ThreadId# -> ThreadId
ThreadId ThreadId#
tid #)