src/Parallel.hs


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

module Parallel where

import Control.Concurrent
import Control.Monad (replicateM_)
import Data.IORef


-- | Does not return results in-order.
parallelForM :: [a] -> (a -> IO b) -> IO [b]
parallelForM inputList action = do
  nthread <- getNumCapabilities

  listref <- newIORef inputList
  outref <- newIORef []
  donechan <- newChan
  replicateM_ nthread $ forkIO $
    let loop = do
          mitem <- atomicModifyIORef' listref (\case l@[] -> (l, Nothing)
                                                     item : l -> (l, Just item))
          case mitem of
            Just item -> do
              res <- action item
              atomicModifyIORef' outref (\l -> (res : l, ()))
              loop
            Nothing -> do
              writeChan donechan ()
    in loop

  replicateM_ nthread $ readChan donechan
  reverse <$> readIORef outref