{-# LANGUAGE BangPatterns               #-}
{-# LANGUAGE CPP                        #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}

module Main (main) where

import           Control.DeepSeq
import           Control.Monad
import           Control.Monad.ST
import           Control.Monad.Trans
import           Data.Benchmarks.UnorderedCollections.Distributions
import           Data.Benchmarks.UnorderedCollections.Types
import           Data.Bits
import           Data.ByteString                                    (ByteString)
import qualified Data.ByteString                                    as B
import qualified Data.ByteString.Base16                             as B16
import           Data.Hashable
import qualified Data.HashMap.Strict                                as UC
import qualified Data.HashTable.IO                                  as IOH
import           Data.IORef
import qualified Data.Map                                           as Map
import           System.Environment
import           System.FilePath
import           System.Random.MWC

import           Criterion.Collection.Main
import           Criterion.Collection.Sample
import           Criterion.Collection.Types


------------------------------------------------------------------------------
#if !MIN_VERSION_bytestring(0,10,0)
instance NFData ByteString
#endif

------------------------------------------------------------------------------
dataMap :: (Ord k, Eq k) => DataStructure (Operation k)
dataMap = setupData Map.empty f
  where
    f !m !op = case op of
                 (Insert k v) -> let !m' = Map.insert k v m in m'
                 (Lookup k)   -> let !_  = Map.lookup k m in m
                 (Delete k)   -> let !m' = Map.delete k m in m'
{-# INLINE dataMap #-}


------------------------------------------------------------------------------
hashMap :: (Hashable k, Eq k) => DataStructure (Operation k)
hashMap = setupData UC.empty f
  where
    f !m !op = case op of
                 (Insert k v) -> let !m' = UC.insert k v m in m'
                 (Lookup k)   -> let !_  = UC.lookup k m in m
                 (Delete k)   -> let !m' = UC.delete k m in m'
{-# INLINE hashMap #-}


------------------------------------------------------------------------------
basicHashTable :: (Hashable k, Eq k) => DataStructure (Operation k)
basicHashTable = setupDataIO (IOH.newSized :: Int -> IO (IOH.BasicHashTable k v))
                             f
  where
    f !m !op = case op of
                 (Insert k v) -> IOH.insert m k v >> return m
                 (Lookup k)   -> do
                           !_ <- IOH.lookup m k
                           return m
                 (Delete k)   -> IOH.delete m k >> return m
{-# INLINE basicHashTable #-}


------------------------------------------------------------------------------
cuckooHashTable :: (Hashable k, Eq k) => DataStructure (Operation k)
cuckooHashTable = setupDataIO (IOH.newSized :: Int -> IO (IOH.CuckooHashTable k v))
                             f
  where
    f !m !op = case op of
                 (Insert k v) -> IOH.insert m k v >> return m
                 (Lookup k)   -> do
                           !_ <- IOH.lookup m k
                           return m
                 (Delete k)   -> IOH.delete m k >> return m
{-# INLINE cuckooHashTable #-}


------------------------------------------------------------------------------
linearHashTable :: (Hashable k, Eq k) => DataStructure (Operation k)
linearHashTable = setupDataIO
                    (IOH.newSized :: Int -> IO (IOH.LinearHashTable k v))
                    f
  where
    f !m !op = case op of
                 (Insert k v) -> IOH.insert m k v >> return m
                 (Lookup k)   -> do
                           !_ <- IOH.lookup m k
                           return m
                 (Delete k)   -> IOH.delete m k >> return m
{-# INLINE linearHashTable #-}


------------------------------------------------------------------------------
mkByteString :: GenIO -> IO ByteString
mkByteString rng = do
    n <- uniformR (4,16) rng
    xs <- replicateM n (uniform rng)
    let !s = B.pack xs
    return $! B16.encode s


------------------------------------------------------------------------------
mkConsecutiveIntegers :: IORef Int -> GenIO -> IO Int
mkConsecutiveIntegers ref _ = do
    !i <- atomicModifyIORef ref f
    return $! i
  where
    f !i = let !j = i+1 in (j,j)


------------------------------------------------------------------------------
newtype IntMix = IntMix Int
  deriving (Num, Read, Show, Ord, Eq, NFData)


------------------------------------------------------------------------------
instance Hashable IntMix where
    hash (IntMix a) = hashWithSalt 1102951999 a
    hashWithSalt salt (IntMix a) = hashWithSalt salt a


------------------------------------------------------------------------------
loadConsecutiveIntegersWorkload :: WorkloadGenerator (Operation Int)
loadConsecutiveIntegersWorkload size = do
    ref <- liftIO $ newIORef 0
    loadOnly (mkConsecutiveIntegers ref) size


------------------------------------------------------------------------------
loadConsecutiveIntegersWorkload' :: WorkloadGenerator (Operation IntMix)
loadConsecutiveIntegersWorkload' size = do
    ref <- liftIO $ newIORef 0
    loadOnly (\rng -> IntMix `fmap` mkConsecutiveIntegers ref rng) size


------------------------------------------------------------------------------
testStructures = [ ("Data.Map"             , dataMap        )
                 , ("Data.HashMap"         , hashMap        )
                 , ("Data.BasicHashTable"  , basicHashTable )
                 , ("Data.LinearHashTable" , linearHashTable)
                 , ("Data.CuckooHashTable" , cuckooHashTable)
                 ]

intStructures = [ ("Data.Map"            , dataMap        )
                , ("Data.HashMap"        , hashMap        )
                , ("Data.BasicHashTable" , basicHashTable )
                , ("Data.CuckooHashTable", cuckooHashTable)
                ]

intStructures' = [ ("Data.Map"            , dataMap        )
                 , ("Data.HashMap"        , hashMap        )
                 , ("Data.BasicHashTable" , basicHashTable )
                 , ("Data.CuckooHashTable", cuckooHashTable)
                 ]


------------------------------------------------------------------------------
testSizes :: [Int]
testSizes = [ 250
            , 500
            , 1000
            , 2000
            , 4000
            , 8000
            , 16000
            , 32000
            , 64000
            , 128000
            , 256000
            , 512000
            , 1024000
            , 2048000 ]


------------------------------------------------------------------------------
lookupBenchmark :: Benchmark (Operation ByteString)
lookupBenchmark = Benchmark "Lookup Performance"
                      testStructures
                      testSizes
                      (loadAndUniformLookup mkByteString)


------------------------------------------------------------------------------
insertBenchmark :: Benchmark (Operation ByteString)
insertBenchmark = Benchmark "Insert Performance"
                      testStructures
                      testSizes
                      (loadOnly mkByteString)


------------------------------------------------------------------------------
consecutiveIntBenchmark :: Benchmark (Operation Int)
consecutiveIntBenchmark = Benchmark "Insert consecutive ints"
                              intStructures
                              testSizes
                              loadConsecutiveIntegersWorkload


------------------------------------------------------------------------------
consecutiveIntWithMixBenchmark :: Benchmark (Operation IntMix)
consecutiveIntWithMixBenchmark = Benchmark "Insert consecutive ints (mixed)"
                                     intStructures'
                                     testSizes
                                     loadConsecutiveIntegersWorkload'

------------------------------------------------------------------------------
main :: IO ()
main = do
    args <- getArgs
    let fn = case args of []    -> Nothing
                          (x:_) -> Just (dropExtensions x)

    let cfg = defaultCriterionCollectionConfig
    runBenchmark PerBatch Mutating insertBenchmark cfg (fmap (++".insert") fn)
    runBenchmark PerBatch Pure lookupBenchmark cfg (fmap (++".lookup") fn)
    runBenchmark PerBatch Mutating consecutiveIntBenchmark cfg
                 (fmap (++".int") fn)
    runBenchmark PerBatch Mutating consecutiveIntWithMixBenchmark cfg
                 (fmap (++".intmix") fn)