{-# LANGUAGE RankNTypes, BangPatterns #-}

-- Compare low-level, fused, unfused, and partially fused

import Data.Conduit

import qualified Data.Conduit.List as CL

import Data.Conduit.Internal (Step (..), Stream (..), unstream, StreamConduit (..))

import Criterion.Main

import Data.Functor.Identity (runIdentity)

-- | unfused

enumFromToC :: (Eq a, Monad m, Enum a) => a -> a -> Producer m a

enumFromToC x0 y =

    loop x0

  where

    loop x

        | x == y = yield x

        | otherwise = yield x >> loop (succ x)

{-# INLINE enumFromToC #-}

-- | unfused

mapC :: Monad m => (a -> b) -> Conduit a m b

mapC f = awaitForever $ yield . f

{-# INLINE mapC #-}

-- | unfused

foldC :: Monad m => (b -> a -> b) -> b -> Consumer a m b

foldC f =

    loop

  where

    loop !b = await >>= maybe (return b) (loop . f b)

{-# INLINE foldC #-}

main :: IO ()

main = defaultMain

    [ bench "low level" $ flip whnf upper0 $ \upper ->

        let loop x t

                | x > upper = t

                | otherwise = loop (x + 1) (t + ((x * 2) + 1))

         in loop 1 0

    , bench "completely fused" $ flip whnf upper0 $ \upper ->

        runIdentity

            $ CL.enumFromTo 1 upper

           $$ CL.map (* 2)

           =$ CL.map (+ 1)

           =$ CL.fold (+) 0

    , bench "runConduit, completely fused" $ flip whnf upper0 $ \upper ->

        runIdentity

            $ runConduit

            $ CL.enumFromTo 1 upper

          =$= CL.map (* 2)

          =$= CL.map (+ 1)

          =$= CL.fold (+) 0

    , bench "completely unfused" $ flip whnf upper0 $ \upper ->

        runIdentity

            $ enumFromToC 1 upper

           $$ mapC (* 2)

           =$ mapC (+ 1)

           =$ foldC (+) 0

    , bench "beginning fusion" $ flip whnf upper0 $ \upper ->

        runIdentity

            $ (CL.enumFromTo 1 upper $= CL.map (* 2))

           $$ mapC (+ 1)

           =$ foldC (+) 0

    , bench "middle fusion" $ flip whnf upper0 $ \upper ->

        runIdentity

            $ enumFromToC 1 upper

           $$ (CL.map (* 2) =$= CL.map (+ 1))

           =$ foldC (+) 0

    , bench "ending fusion" $ flip whnf upper0 $ \upper ->

        runIdentity

            $ enumFromToC 1 upper

           $= mapC (* 2)

           $$ (CL.map (+ 1) =$ CL.fold (+) 0)

    , bench "performance of CL.enumFromTo without fusion" $ flip whnf upper0 $ \upper ->

        runIdentity

            $ CL.enumFromTo 1 upper

           $= mapC (* 2)

           $$ (CL.map (+ 1) =$ CL.fold (+) 0)

    ]

  where

    upper0 = 100000 :: Int