#!/usr/bin/env stack
-- stack --resolver lts-8.3 exec --package vector --package criterion -- ghc -O2
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Main where

import           Control.Monad.ST            (ST)
import           Criterion.Main
import qualified Data.Vector.Unboxed         as VU
import qualified Data.Vector.Unboxed.Mutable as MVU

-- | Delayed Array.
data DArray ix e = DArray { dSize        :: !ix
                          , dUnsafeIndex :: ix -> e }

-- | Manifest Array.
data MArray ix e = MArray { mSize        :: !ix
                          , mUnsafeIndex :: ix -> e }


-- | Efficient loop with an accumulator
loop :: Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> a) -> a
loop !init' condition increment !initAcc f = go init' initAcc where
  go !step !acc =
    case condition step of
      False -> acc
      True  -> go (increment step) (f step acc)
{-# INLINE loop #-}

-- | Efficient monadic loop
loopM_ :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ !init' condition increment f = go init' where
  go !step =
    case condition step of
      False -> return ()
      True  -> f step >> go (increment step)
{-# INLINE loopM_ #-}

-- | Load a delayed Array into memory. Unlike strem fusion, this can also be parallelized.
compute
  :: forall e . VU.Unbox e => DArray Int e -> MArray Int e
compute (DArray k f) = MArray k $ (VU.unsafeIndex v)
  where
    -- Make sure it's evaluated (deepseq == seq is for Unboxed Vector)
    !v = VU.create generateArray
    generateArray :: ST s (MVU.MVector s e)
    generateArray = do
      mv <- MVU.unsafeNew k
      loopM_ 0 (< k) (+ 1) $ \ !i -> MVU.unsafeWrite mv i (f i)
      return mv
    {-# INLINE generateArray #-}
{-# INLINE compute #-}


makeArray1D :: Int -> (Int -> e) -> DArray Int e
makeArray1D = DArray
{-# INLINE makeArray1D #-}

mapA :: (b -> e) -> DArray ix b -> DArray ix e
mapA f (DArray k g) = DArray k (f . g)
{-# INLINE mapA #-}

foldlA :: (t1 -> t -> t1) -> t1 -> DArray Int t -> t1
foldlA f !acc (DArray k g) = loop 0 (< k) (+ 1) acc $ \ !i !acc0 -> f acc0 (g i)
{-# INLINE foldlA #-}

sumA :: DArray Int Int -> Int
sumA = foldlA (+) 0
{-# INLINE sumA #-}

main :: IO ()
main = do
  let !sz = 640000
  let arr1D = (`makeArray1D` id)
  let vec = (`VU.generate` id)
      {-# INLINE vec #-} -- <-- Degraded performance if not inlined.
  defaultMain
    [ bgroup
        "Fold Only"
        [ bench "1D Array" $ whnf (sumA . arr1D) sz
        , bench "Vector Unboxed" $ whnf (VU.sum . vec) sz
        ]
    , bgroup
        "Fold Fused"
        [ bench "1D Array" $ whnf (sumA . mapA (+ 25) . arr1D) sz
        , bench "Vector Unboxed" $ whnf (VU.sum . VU.map (+ 25) . vec) sz
        ]
    , bgroup
        "Compute Only"
        [ bench "1D Array" $ whnf (compute . arr1D) sz
        , bench "Vector Unboxed" $ whnf (`VU.generate` id) sz
        , bench "Vector Unboxed (very slow - not inlined!?!?)" $ whnf vec sz
        ]
    , bgroup
        "Compute Fused"
        [ bench "1D Array" $ whnf (compute . mapA (+ 50) . arr1D) sz
        , bench "Vector Unboxed" $ whnf (VU.map (+ 50) . vec) sz
        ]
    ]