{-# LANGUAGE MagicHash, ScopedTypeVariables, UnboxedTuples, BangPatterns, OverlappingInstances, TypeFamilies, OverloadedStrings, NoMonomorphismRestriction, FlexibleInstances, FlexibleContexts, FunctionalDependencies, MultiParamTypeClasses, RankNTypes, DataKinds, GADTs, GeneralizedNewtypeDeriving, TemplateHaskell #-}
module Main
  ( main
  , foo
  , bar
  , baz
  , fast
  , fastST ) where

import           Criterion
import           Criterion.Main
import           Criterion.Types
import           Data.IORef
import           Data.STRef
import qualified Data.List.Stream as S
import           Control.Monad.ST
import           Control.Monad.Stream
import           Control.Monad.Trans
import           Control.Monad.Trans.RWS.Strict
import           Control.Monad.Trans.State.Strict
import           Control.Monad.Trans.Maybe
import           Data.Text (Text)
import qualified Data.Text as T

newtype Foo a = Foo { runFoo :: IO a }
newtype Bar a = Bar { unBar :: StateT Int IO a }
newtype Baz a = Baz { unBaz :: RWST Int [Text] Int (StateT Int IO) a }

runBar x = (flip evalStateT) 0 . unBar $ x
runBaz x = (flip evalStateT) 0 . liftM fst . (\y -> evalRWST y 0 0) . unBaz $ x

n = 10000 :: Int

foo :: Int -> Foo Int
foo inp = Foo $ do
  x <- newIORef inp
  replicateM_ n $ do
    modifyIORef x (+1)
  readIORef x

bar :: Int -> Bar Int
bar inp = Bar $ do
  x <- liftIO $ newIORef inp
  replicateM_ n $ do
    liftIO $ modifyIORef x (\i -> i + 1)
  liftIO $ readIORef x

baz :: Int -> Baz Int
baz inp = Baz $ do
  x <- liftIO $ newIORef inp
  replicateM_ n $ do
    liftIO $ modifyIORef x (+1)
  liftIO $ readIORef x

fast :: Int -> Int -> Int
fast 0 acc = acc
fast k acc = fast (k - 1) (acc + 1)

fastST :: Int -> Int
fastST inp = runST $ do
  x <- newSTRef inp
  replicateM_ n $ do
    modifySTRef x (\i -> i +1)
  readSTRef x

main = defaultMain $
  [ bench "IO: "                       $ runFoo $ foo 0
  , bench "IO with StateT: "           $ runBar $ bar 0
  , bench "IO with StateT and RWST: "  $ runBaz $ baz 0 
  , bench "Pure: "                     $ whnf (\x -> fast x 0) n
  , bench "Pure with ST: "             $ whnf (\x -> fastST x) n ]