{-# LANGUAGE RankNTypes #-}
module Main (main) where

import Criterion.Main
import Common (commonMain)
import Control.Monad.Identity (Identity, runIdentity)
import Pipes
import qualified Pipes.Prelude as P
import Prelude hiding (enumFromTo)

defaultMax :: Int
defaultMax = 10000

main :: IO ()
main = commonMain defaultMax preludeBenchmarks

enumFromTo :: (Int -> a) -> Int -> Int -> Producer a Identity ()
enumFromTo f n1 n2 = loop n1
    where
        loop n =
            if n <= n2
            then do
                yield $! f n
                loop $! n + 1
            else return ()
{-# INLINABLE enumFromTo #-}

drain :: Producer b Identity r -> r
drain p = runIdentity $ runEffect $ for p discard

msum :: (Monad m) => Producer Int m () -> m Int
msum = P.foldM (\a b -> return $ a + b) (return 0) return

scanMSum :: (Monad m) => Pipe Int Int m r
scanMSum = P.scanM (\x y -> return (x + y)) (return 0) return

-- Using runIdentity seems to reduce outlier counts.
preludeBenchmarks :: Int -> [Benchmark]
preludeBenchmarks vmax =
    let applyBench b = b benchEnum_p
        benchEnum_p  = enumFromTo id 1 vmax
    in
    [
      bgroup "Folds" $ map applyBench
        [
          bench "all"       . whnf (runIdentity . P.all (<= vmax))
        , bench "any"       . whnf (runIdentity . P.any (> vmax))
        , bench "find"      . whnf (runIdentity . P.find (== vmax))
        , bench "findIndex" . whnf (runIdentity . P.findIndex (== vmax))
        , bench "fold"      . whnf (runIdentity . P.fold (+) 0 id)
        , bench "foldM"     . whnf (runIdentity . msum)
        , bench "head"      . nf (runIdentity . P.head)
        , bench "index"     . nf (runIdentity . P.index (vmax-1))
        , bench "last"      . nf (runIdentity . P.last)
        , bench "length"    . whnf (runIdentity . P.length)
        , bench "null"      . whnf (runIdentity  . P.null)
        , bench "toList"    . nf P.toList
        ]
    , bgroup "Pipes" $ map applyBench
        [
          bench "chain"       . whnf (drain . (>-> P.chain (\_ -> return ())))
        , bench "drop"        . whnf (drain . (>-> P.drop vmax))
        , bench "dropWhile"   . whnf (drain . (>-> P.dropWhile (<= vmax)))
        , bench "filter"      . whnf (drain . (>-> P.filter even))
        , bench "findIndices" . whnf (drain . (>-> P.findIndices (<= vmax)))
        , bench "map"         . whnf (drain . (>-> P.map id))
        , bench "mapM"        . whnf (drain . (>-> P.mapM return))
        , bench "take"        . whnf (drain . (>-> P.take vmax))
        , bench "takeWhile"   . whnf (drain . (>-> P.takeWhile (<= vmax)))
        , bench "scan"        . whnf (drain . (>-> P.scan (+) 0 id))
        , bench "scanM"       . whnf (drain . (>-> scanMSum))
        ] ++ [
          bench "concat" $ whnf (drain . (>-> P.concat)) $ enumFromTo Just 1 vmax
        ]
    , bgroup "Zips" $ map applyBench
        [
          bench "zip"     . whnf (drain . P.zip benchEnum_p)
        , bench "zipWith" . whnf (drain . P.zipWith (+) benchEnum_p)
        ]
    , bgroup "enumFromTo.vs.each"
        [
          bench "enumFromTo" $ whnf (drain . enumFromTo id 1) vmax
        , bench "each"       $ whnf (drain . each) [1..vmax]
        ]
    ]