module Main where

import Control.Concurrent.Async
import Control.Monad
import Criterion
import Criterion.Main
import Data.ByteString.SuperBuffer
import qualified Data.ByteString.SuperBuffer.Pure as P
import Data.Int
import qualified Data.BufferBuilder as BB
import qualified Data.ByteString as BS
import qualified Data.ByteString.Builder as BSB
import qualified Data.ByteString.Lazy as BSL

main :: IO ()
main =
    defaultMain
    [ bgroup "main"
      [ mkGroup "small" 5000 4000
      , mkGroup "med" 500 40000
      , mkGroup "large" 50 400000
      ]
    ]

mkGroup :: String -> Int -> Int -> Benchmark
mkGroup name steps chunkSize =
    bgroup name $
    mkSizedGroup steps chunkSize bufName buildBuf
    ++ mkSizedGroup steps chunkSize bufNameT buildBufT
    ++ mkSizedGroup steps chunkSize bufNameP buildBufP
    ++ mkSizedGroup steps chunkSize bufNamePT buildBufPT
    ++ mkSizedGroup steps chunkSize bufBBName buildBufBB
    ++
    [ bench "bytestring builder" $ nfIO $ BS.reverse <$> buildBufBuilder steps chunkSize
    , bench "bytestring fromChunks" $ nfIO $ BS.reverse <$> buildBufChunks steps chunkSize
    , bench "bytestring concat" $ nfIO $ BS.reverse <$> buildBufConcat steps chunkSize
    ]
    where
      bufBBName is = "buffer-builder (init=" ++ show is ++ " bytes, trim=yes)"
      bufName is = "superbuffer (init=" ++ show is ++ " bytes)"
      bufNameT is = "superbuffer (init=" ++ show is ++ " bytes, threadsafe, 2 concurrent writes)"
      bufNameP is = "superbuffer (pure haskell, init=" ++ show is ++ " bytes)"
      bufNamePT is = "superbuffer (pure haskell, init=" ++ show is ++ " bytes, threadsafe, 2 concurrent writes)"

mkSizedGroup ::
    Int -> Int -> (Int64 -> String) -> (Int64 -> Int -> Int -> IO BS.ByteString) -> [Benchmark]
mkSizedGroup steps chunkSize bufName builder =
    [ bench (bufName iBufSize128) $ nfIO $ BS.reverse <$> builder iBufSize128 steps chunkSize
    , bench (bufName iBufSize) $ nfIO $ BS.reverse <$> builder iBufSize steps chunkSize
    , bench (bufName iBufSize2) $ nfIO $ BS.reverse <$> builder iBufSize2 steps chunkSize
    , bench (bufName iBufSize4) $ nfIO $ BS.reverse <$> builder iBufSize4 steps chunkSize
    , bench (bufName iBufSizeAll) $ nfIO $ BS.reverse <$> builder iBufSizeAll steps chunkSize
    ]
    where
      iBufSize128 = 128
      iBufSize = fromIntegral chunkSize
      iBufSize2 = 2 * fromIntegral chunkSize
      iBufSize4 = 4 * fromIntegral chunkSize
      iBufSizeAll = fromIntegral $ steps * chunkSize

mkChunk :: Int -> Int -> BS.ByteString
mkChunk step chunkSize =
    BS.replicate chunkSize (fromIntegral $ (step `mod` 100) + 50)
{-# INLINE mkChunk #-}

buildBuf :: Int64 -> Int -> Int -> IO BS.ByteString
buildBuf bufSize steps chunkSize =
    withBuffer bufSize $ \buf ->
    forM_ [0..steps] $ \step ->
    appendBuffer buf (mkChunk step chunkSize)

buildBufT :: Int64 -> Int -> Int -> IO BS.ByteString
buildBufT bufSize steps chunkSize =
    withBuffer bufSize $ \buf ->
    forM_ [0..(ceiling halfSteps)] $ \step ->
    concurrently_
       (appendBufferT buf (mkChunk step chunkSize))
       (appendBufferT buf (mkChunk step chunkSize))
    where
      halfSteps :: Double
      halfSteps = fromIntegral steps / 2.0

buildBufP :: Int64 -> Int -> Int -> IO BS.ByteString
buildBufP bufSize steps chunkSize =
    P.withBuffer (fromIntegral bufSize) $ \buf ->
    forM_ [0..steps] $ \step ->
    P.appendBuffer buf (mkChunk step chunkSize)

buildBufPT :: Int64 -> Int -> Int -> IO BS.ByteString
buildBufPT bufSize steps chunkSize =
    P.withBuffer (fromIntegral bufSize) $ \buf ->
    forM_ [0..(ceiling halfSteps)] $ \step ->
    concurrently_
       (P.appendBufferT buf (mkChunk step chunkSize))
       (P.appendBufferT buf (mkChunk step chunkSize))
    where
      halfSteps :: Double
      halfSteps = fromIntegral steps / 2.0

buildBufBB :: Int64 -> Int -> Int -> IO BS.ByteString
buildBufBB bufSize steps chunkSize =
    pure $ BB.runBufferBuilderWithOptions opts $
    forM_ [0..steps] $ \step ->
    BB.appendBS (mkChunk step chunkSize)
    where
      opts =
          BB.Options
          { BB.initialCapacity = fromIntegral bufSize
          , BB.trimFinalBuffer = True
          }

buildBufBuilder :: Int -> Int -> IO BS.ByteString
buildBufBuilder steps chunkSize =
    BSL.toStrict . BSB.toLazyByteString <$>
    foldM (\b a ->  pure $ b `mappend` BSB.byteString (mkChunk a chunkSize)) mempty [0..steps]

buildBufChunks :: Int -> Int -> IO BS.ByteString
buildBufChunks steps chunkSize =
    BSL.toStrict . BSL.fromChunks <$> (
    forM [0..steps] $ \step ->
            pure (mkChunk step chunkSize))


buildBufConcat :: Int -> Int -> IO BS.ByteString
buildBufConcat steps chunkSize =
    BS.concat <$> (
    forM [0..steps] $ \step ->
            pure (mkChunk step chunkSize))