{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}

{-# OPTIONS_GHC -O2 #-}


import Control.Monad
import Data.Primitive
import Control.Monad.ST
import Criterion.Main
import GHC.Prim
import GHC.Int
import qualified Weigh as W

main :: IO ()
main = do
  let !thousand = upto 1000
      !tenThousand = upto 10000
  when (boxedFoldInt (+) 0 thousand /= (div (999 * 1000) 2)) $ do
    fail "boxedFoldInt is incorrect"
  when (unboxedFoldInt (+#) 0 thousand /= (div (999 * 1000) 2)) $ do
    fail "unboxedFoldInt is incorrect"
  let runCriterionBench = defaultMain
        [ bgroup "mapping ByteArray"
          [ bench "boxed" $ whnf (boxedFoldInt (+) 0) tenThousand
          , bench "unboxed" $ whnf (unboxedFoldInt (+#) 0) tenThousand
          ]
        ]
      runWeighBench = W.mainWith $ do
        W.func "boxed" (boxedFoldInt (+) 0) tenThousand
        W.func "unboxed" (unboxedFoldInt (+#) 0) tenThousand
  -- You can either run the criterion benchmark or the
  -- weigh benchmark but not both at the same time. Criterion
  -- will show you that the unboxed variant is faster.
  -- Weigh will show you why it's faster.
  runCriterionBench
  -- runWeighBench
  putStrLn "Finished Benchmarking"



upto :: Int -> ByteArray
upto n = runST $ unsafeFreezeByteArray =<< go 0 =<< newByteArray (sizeOf (undefined :: Int) * n)
  where
  go :: Int -> MutableByteArray s -> ST s (MutableByteArray s)
  go ix arr = if ix < n
    then do
      writeByteArray arr ix ix
      go (ix + 1) arr
    else return arr

-- If we don't put NOINLINE pragmas on these, GHC optimizer gets too
-- smart and inlines even the partial applications of these, making
-- the benchmark entirely worthless. Specifically, what this means is
-- that we end up a partial application of boxedFoldInt to (+) and 0.
-- GHC ends up proceeds to optimizes this new function, and since
-- it now knows the function being used as the fold, it can inline
-- it and unbox everything.

{-# NOINLINE boxedFoldInt #-}
boxedFoldInt :: (Int -> Int -> Int) -> Int -> ByteArray -> Int
boxedFoldInt f acc0 arr = go 0 acc0
  where
  !sz = div (sizeofByteArray arr) (sizeOf (undefined :: Int))
  go !ix !acc = if ix < sz
    then go (ix + 1) (f (indexByteArray arr ix) acc)
    else acc

{-# NOINLINE unboxedFoldInt #-}
unboxedFoldInt :: (Int# -> Int# -> Int#) -> Int -> ByteArray -> Int
unboxedFoldInt f (I# acc0) boxedArr@(ByteArray arr) = I# (go 0 acc0)
  where
  !sz = div (sizeofByteArray boxedArr) (sizeOf (undefined :: Int))
  go :: Int -> Int# -> Int#
  go !ix@(I# ix#) !acc = if ix < sz
    then go (ix + 1) (f (indexIntArray# arr ix#) acc)
    else acc