{-# OPTIONS_GHC -fno-warn-name-shadowing #-}

module Parallel
  ( streamsOf,
    parMapM,
  )
where

import Control.Concurrent.ParallelIO.Local (parallel, withPool)
import GHC.Conc (numCapabilities)

streamsOf :: Int -> [a] -> [[a]]
streamsOf 1 xs = [xs]
streamsOf n xs
  | n > 0 = [everyN k xs | k <- [0 .. n - 1]]
  | otherwise = []
  where
    everyN k xs = map snd $ filter ((== k) . (`mod` n) . fst) $ zip [0 ..] xs

parMapM :: (a -> IO b) -> [a] -> IO [b]
parMapM f xs = do
  withPool numCapabilities $ \pool -> parallel pool (map f xs)