{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UnboxedTuples       #-}

module Main where

-- criterion
import qualified Criterion.Main as Criterion
  ( bench, bgroup, defaultMain, nf )
import Data.Foldable
import Data.STRef
import Control.Monad.ST
import Data.STRef.Unboxed

testList :: [Int]
testList = [0..10^6]

sum_recurs :: [Int] -> Int
sum_recurs = go 0
  where
    go :: Int -> [Int] -> Int
    go acc [] = acc
    go acc (n : ns) = go (acc + n) ns


sum_boxed_ref :: [Int] -> Int
sum_boxed_ref = go
  where
    go :: [Int] -> Int
    go ns = runST $ do
      acc <- newSTRef 0
      for_ ns $ \n ->
        modifySTRef' acc (+ n)
      readSTRef acc

sum_unboxed_ref :: [Int] -> Int
sum_unboxed_ref = go
  where
    go :: [Int] -> Int
    go ns = runST $ do
      acc <- newSTRefU 0
      for_ ns $ \n ->
        modifySTRefU acc (+ n)
      readSTRefU acc

main :: IO ()
main = Criterion.defaultMain
     [ Criterion.bgroup "STRefs"
        [ Criterion.bench "foldr (+) 0"     $ Criterion.nf (foldr (+) 0)   testList
        , Criterion.bench "sum"             $ Criterion.nf sum             testList
        , Criterion.bench "sum_recurs"      $ Criterion.nf sum_recurs      testList
        , Criterion.bench "sum_boxed_ref"   $ Criterion.nf sum_boxed_ref   testList
        , Criterion.bench "sum_unboxed_ref" $ Criterion.nf sum_unboxed_ref testList
        ]
     ]