Infinite lists and parallelism in Haskell

parallelStreams.hs

module Main where

import qualified Data.List as L
import qualified System.Random as R

import qualified Data.Numbers.Primes as P

import Control.Parallel.Strategies (NFData, ($||))
import qualified Control.Parallel.Strategies as S

import qualified Criterion.Main as C

randoms :: IO [Int]
randoms = (L.nub . R.randoms) `fmap` R.getStdGen

primes :: [Int]
primes = P.primes

primeFactors :: Int -> [Int]
primeFactors = P.primeFactors

prefixReduce :: (a -> b -> b) -> b -> [a] -> b
prefixReduce = foldr

parMap :: (NFData b) => (a -> b) -> [a] -> [b]
-- parMap = S.parMap S.rdeepseq
parMap f = S.withStrategy (S.parBuffer bufferAmount S.rdeepseq) . map f

parFilter :: (NFData a) => (a -> Bool) -> [a] -> [a]
parFilter p = S.withStrategy (S.parBuffer bufferAmount S.rdeepseq) . filter p

parFoldrN :: (NFData b) => Int -> (a -> b -> b) -> b -> [a] -> b
parFoldrN n f z [] = S.withStrategy S.rdeepseq z
parFoldrN n f z xs =
    ((\acc -> parFoldrN n f acc rest) $|| S.rdeepseq) (foldr f z buffer)
    where (buffer, rest) = L.splitAt n xs

parFoldr :: (NFData b) => (a -> b -> b) -> b -> [a] -> b
parFoldr f z [] = S.withStrategy S.rdeepseq z
parFoldr f z (x : xs) = (f x $|| S.rdeepseq) (parFoldr f z xs)

parPrefixReduce :: (NFData b) => (a -> b -> b) -> b -> [a] -> b
parPrefixReduce = parFoldr

parPrefixReduceN :: (NFData b) => Int -> (a -> b -> b) -> b -> [a] -> b
parPrefixReduceN = parFoldrN

bufferAmount :: Int
bufferAmount = 10

takeAmount :: Int
takeAmount = 1000000

expensive :: Int -> Int
expensive = maximum . primeFactors . (+1)

main :: IO ()
main = C.defaultMain $
  C.bench "take randoms" (take takeAmount `fmap` randoms) :
  C.bench "take primes" (C.nf (take takeAmount) primes) :
  C.bench "take prime factors" (C.nf (take takeAmount . primeFactors) 9223372036854775) :
  C.bench "map randoms" ((take takeAmount . map expensive) `fmap` randoms) :
  C.bench "parMap randoms" ((take takeAmount . parMap expensive) `fmap` randoms) :
  C.bench "filter randoms" ((take takeAmount . filter (odd . expensive)) `fmap` randoms) :
  C.bench "parFilter randoms" ((take takeAmount . parFilter (odd . expensive)) `fmap` randoms) :
  C.bench "prefixReduce randoms" ((prefixReduce (\x y -> expensive (x + y)) 0 . take takeAmount) `fmap` randoms) :
  C.bench "parPrefixReduce randoms" ((parPrefixReduce (\x y -> expensive (x + y)) 0 . take takeAmount) `fmap` randoms) :
  C.bench "map primes" (C.nf (take takeAmount . map expensive) primes) :
  C.bench "parMap primes" (C.nf (take takeAmount . parMap expensive) primes) :
  C.bench "filter primes" (C.nf (take takeAmount . filter (odd . expensive)) primes) :
  C.bench "parFilter primes" (C.nf (take takeAmount . parFilter (odd . expensive)) primes) :
  C.bench "prefixReduce primes" (C.nf (prefixReduce (\x y -> expensive (x + y)) 0 . take takeAmount) primes) :
  C.bench "parPrefixReduce primes" (C.nf (parPrefixReduceN bufferAmount (\x y -> expensive (x + y)) 0 . take takeAmount) primes) :
  []