{-# LANGUAGE BangPatterns   #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveAnyClass #-}
module Main where

import           Prelude as P

import           Criterion.Main
import           Control.DeepSeq
import           Data.Monoids
import           Foreign.Storable (Storable, alignment, peek,
                                   peekByteOff, poke, pokeByteOff,
                                   sizeOf)
import GHC.Generics (Generic)
import System.Random

import qualified Data.Vector.Storable         as SVec
import           Data.Vector.Storable         (Vector)
import qualified Data.Vector.Storable.Mutable as Vector
import           Data.Vector.Storable.Mutable (IOVector)


data Foo = Foo Int Int deriving (Show, Eq, Generic, NFData)

chunkSize :: Int
chunkSize = sizeOf (undefined :: Int)
{-# INLINE chunkSize #-}

instance Storable Foo where
    sizeOf    _ = 2 * chunkSize ; {-# INLINE sizeOf    #-}
    alignment _ = chunkSize     ; {-# INLINE alignment #-}
    peek ptr = Foo
        <$> peekByteOff ptr 0
        <*> peekByteOff ptr chunkSize
    {-# INLINE peek #-}
    poke ptr (Foo a b) = do
        pokeByteOff ptr 0 a
        pokeByteOff ptr chunkSize b
    {-# INLINE poke #-}

-- create an empty vector
mkFooVec :: Int -> IO (Vector Foo)
mkFooVec !i = SVec.unsafeFreeze =<< Vector.new (i + 1)

-- fill the vector with sample structs
populateFooVec :: Int -> Vector Foo -> IO (Vector Foo)
populateFooVec !i !v = do
    v' <- SVec.unsafeThaw v
    let go 0 = return ()
        go j = Vector.unsafeWrite v' j (Foo j $ j + 1) >> go (j - 1)
    go i
    SVec.unsafeFreeze v'

-- trying to force the evaluation by accessing random elements
populateAndAccess :: Int -> Vector Foo -> IO Foo
populateAndAccess !i !v = do
    vec <- populateFooVec i v
    idx <- getStdRandom (randomR (0, i - 1))
    return $ SVec.unsafeIndex vec idx

-- trying once again, this time printing the elements to the console
accessAndPrint :: Int -> Vector Foo -> IO ()
accessAndPrint !i !v = do
    foo <- populateAndAccess i v
    print foo

main :: IO ()
main = do
    defaultMain [
          bgroup "Storable vector (mutable)"
            $ (\(i :: Int) -> env (mkFooVec (10 ^ i))
            $ \v -> bench ("10e" <> show i)
            $ nfIO (populateFooVec (10 ^ i) v))  <$> [6..8]
        , bgroup "Storable vector with index (mutable)"
            $ (\(i :: Int) -> env (mkFooVec (10 ^ i))
            $ \v -> bench ("10e" <> show i)
            $ nfIO (populateAndAccess (10 ^ i) v))  <$> [6..8]
        , bgroup "Storable vector with index and print (mutable)"
            $ (\(i :: Int) -> env (mkFooVec (10 ^ i))
            $ \v -> bench ("10e" <> show i)
            $ nfIO (accessAndPrint (10 ^ i) v))  <$> [6..8]
        ]