module System.Directory.Parallel ( parTraverse
, parTraverseAll
) where
import Control.Concurrent (getNumCapabilities)
import Control.Concurrent.ParallelIO.Local (Pool, parallel_, withPool)
import Control.Monad (filterM)
import System.Directory (doesDirectoryExist, listDirectory)
import System.FilePath ((</>))
partitionM :: Monad m => (a -> m Bool) -> [a] -> m ([a], [a])
partitionM _ [] = pure ([], [])
partitionM f (x:xs) = do
res <- f x
(as,bs) <- partitionM f xs
if res
then pure (x:as, bs)
else pure (as, x:bs)
parTraverseAll :: (FilePath -> IO ())
-> (FilePath -> IO Bool)
-> (FilePath -> IO Bool)
-> [FilePath]
-> IO ()
parTraverseAll act fileP dirP fps = do
(dirs, files) <- partitionM doesDirectoryExist fps
parTraverseFiles act fileP dirP dirs files
parTraverseFiles :: (FilePath -> IO ())
-> (FilePath -> IO Bool)
-> (FilePath -> IO Bool)
-> [FilePath]
-> [FilePath]
-> IO ()
parTraverseFiles act fileP dirP dirs files = do
ncpu <- getNumCapabilities
withPool ncpu $ \pool ->
parallel_ pool $ fmap act files ++ fmap (loopPool pool) dirs
where loopPool :: Pool -> FilePath -> IO ()
loopPool pool fp = do
all' <- fmap (fp </>) <$> listDirectory fp
(dirs', files') <- partitionM doesDirectoryExist all'
dirs'' <- filterM dirP dirs'
files'' <- filterM fileP files'
parallel_ pool (fmap act files'' ++ fmap (loopPool pool) dirs'')
parTraverse :: (FilePath -> IO ())
-> (FilePath -> IO Bool)
-> (FilePath -> IO Bool)
-> [FilePath]
-> IO ()
parTraverse act fileP dirP dirs = parTraverseFiles act fileP dirP dirs []