-- |Intended for internal use: Parallel evaluation of @IO@ values
module System.IO.Parallel
    ( twoParallel
    , threeParallel
    , fourParallel
    , manyParallel
    ) where

import Control.Concurrent (forkIO, yield)
import Control.Concurrent.MVar (MVar, newEmptyMVar, takeMVar, putMVar)


-------------------

-- |Run an @IO@ computation in parallel. The result will appear in the @MVar@.
async :: IO a -> IO (MVar a)
async :: IO a -> IO (MVar a)
async IO a
m = do
    MVar a
v <- IO (MVar a)
forall a. IO (MVar a)
newEmptyMVar
    ThreadId
_ <- IO () -> IO ThreadId
forkIO (IO () -> IO ThreadId) -> IO () -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ do
        a
x <- IO a
m
        IO ()
yield
        MVar a -> a -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar a
v a
x
    MVar a -> IO (MVar a)
forall (m :: * -> *) a. Monad m => a -> m a
return MVar a
v

-- |Run two @IO@ computations in parallel and wait for the results.
twoParallel :: IO a -> IO b -> IO (a, b)
twoParallel :: IO a -> IO b -> IO (a, b)
twoParallel IO a
a IO b
b = do
    MVar a
a' <- IO a -> IO (MVar a)
forall a. IO a -> IO (MVar a)
async IO a
a
    MVar b
b' <- IO b -> IO (MVar b)
forall a. IO a -> IO (MVar a)
async IO b
b
    a
a'' <- MVar a -> IO a
forall a. MVar a -> IO a
takeMVar MVar a
a'
    b
b'' <- MVar b -> IO b
forall a. MVar a -> IO a
takeMVar MVar b
b'
    (a, b) -> IO (a, b)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
a'', b
b'')

-- |Run three @IO@ computations in parallel and wait for the results.
threeParallel :: IO a -> IO b -> IO c -> IO (a, b, c)
threeParallel :: IO a -> IO b -> IO c -> IO (a, b, c)
threeParallel IO a
a IO b
b IO c
c = do
    MVar a
a' <- IO a -> IO (MVar a)
forall a. IO a -> IO (MVar a)
async IO a
a
    MVar b
b' <- IO b -> IO (MVar b)
forall a. IO a -> IO (MVar a)
async IO b
b
    MVar c
c' <- IO c -> IO (MVar c)
forall a. IO a -> IO (MVar a)
async IO c
c
    a
a'' <- MVar a -> IO a
forall a. MVar a -> IO a
takeMVar MVar a
a'
    b
b'' <- MVar b -> IO b
forall a. MVar a -> IO a
takeMVar MVar b
b'
    c
c'' <- MVar c -> IO c
forall a. MVar a -> IO a
takeMVar MVar c
c'
    (a, b, c) -> IO (a, b, c)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
a'', b
b'', c
c'')

-- |Run four @IO@ computations in parallel and wait for the results.
fourParallel :: IO a -> IO b -> IO c -> IO d -> IO (a, b, c, d)
fourParallel :: IO a -> IO b -> IO c -> IO d -> IO (a, b, c, d)
fourParallel IO a
a IO b
b IO c
c IO d
d = do
    MVar a
a' <- IO a -> IO (MVar a)
forall a. IO a -> IO (MVar a)
async IO a
a
    MVar b
b' <- IO b -> IO (MVar b)
forall a. IO a -> IO (MVar a)
async IO b
b
    MVar c
c' <- IO c -> IO (MVar c)
forall a. IO a -> IO (MVar a)
async IO c
c
    MVar d
d' <- IO d -> IO (MVar d)
forall a. IO a -> IO (MVar a)
async IO d
d
    a
a'' <- MVar a -> IO a
forall a. MVar a -> IO a
takeMVar MVar a
a'
    b
b'' <- MVar b -> IO b
forall a. MVar a -> IO a
takeMVar MVar b
b'
    c
c'' <- MVar c -> IO c
forall a. MVar a -> IO a
takeMVar MVar c
c'
    d
d'' <- MVar d -> IO d
forall a. MVar a -> IO a
takeMVar MVar d
d'
    (a, b, c, d) -> IO (a, b, c, d)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
a'', b
b'', c
c'', d
d'')

-- |Run computations in parallel and wait for the results.
manyParallel :: [IO a] -> IO [a]
manyParallel :: [IO a] -> IO [a]
manyParallel [IO a]
m 
    = (IO a -> IO (MVar a)) -> [IO a] -> IO [MVar a]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM IO a -> IO (MVar a)
forall a. IO a -> IO (MVar a)
async [IO a]
m IO [MVar a] -> ([MVar a] -> IO [a]) -> IO [a]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (MVar a -> IO a) -> [MVar a] -> IO [a]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM MVar a -> IO a
forall a. MVar a -> IO a
takeMVar