{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Main where

import Control.DeepSeq (NFData(..))
import Criterion.Main as Criterion
import Control.Monad (replicateM, replicateM_, void)
import Data.Int
import Data.IORef
import Data.Singletons (SomeSing(..))
import Data.Text (Text)
import qualified Foreign.Concurrent as Concurrent
import qualified Foreign.ForeignPtr as ForeignPtr
import Foreign.JNI
import Foreign.Marshal.Alloc (mallocBytes, finalizerFree, free)
import Foreign.Marshal.Array (callocArray, mallocArray)
import Foreign.Ptr (castPtr)
import Language.Java

newtype Box a = Box { unBox :: a }

-- Not much sense in deepseq'ing foreign pointers. But needed for call to 'env'
-- below.
instance NFData (Box a) where rnf (Box a) = seq a ()

jabs :: Int32 -> IO Int32
jabs x = callStatic "java.lang.Math" "abs" x

jniAbs :: JClass -> JMethodID -> Int32 -> IO Int32
jniAbs klass method x = callStaticIntMethod klass method [coerce x]

intValue :: Int32 -> IO Int32
intValue x = do
    jx <- reflect x
    call jx "intValue"

compareTo :: Int32 -> Int32 -> IO Int32
compareTo x y = do
    jx <- reflect x
    jy <- reflect y
    call jx "compareTo" jy

incrHaskell :: Int32 -> IO Int32
incrHaskell x = return (x + 1)

foreign import ccall unsafe getpid :: IO Int

benchCalls :: Benchmark
benchCalls =
    env ini $ \ ~(Box klass, method) ->
      bgroup "Calls"
      [ bgroup "Java calls"
        [ bench "static method call: unboxed single arg / unboxed return" $ nfIO $ jabs 1
        , bench "static method call: boxed single arg / boxed return" $
          perBatchEnvWithCleanup
            (\batchSize -> do
               pushLocalFrame (2 * fromIntegral batchSize)
               Box <$> reflect ("123" :: Text)
            )
            (\_ _ -> void (popLocalFrame jnull)) $
            \(Box jStringInteger) -> do
              _ <- callStatic "java.lang.Integer" "valueOf" jStringInteger
                 :: IO (J ('Class "java.lang.Integer"))
              return ()
        , bench "jni static method call: unboxed single arg / unboxed return" $ nfIO $ jniAbs klass method 1
        , bench "method call: no args / unboxed return" $ nfIO $ intValue 1
        , bench "method call: boxed single arg / unboxed return" $ nfIO $ compareTo 1 1
        , bench "getClass" $
          perBatchEnvWithCleanup
            (pushLocalFrame . (2*) . fromIntegral)
            (\_ _ -> void (popLocalFrame jnull)) $
            \() -> do
              _ <- getClass (SClass "java/lang/Math")
              return ()
        , bench "getStaticMethodID" $
          perBatchEnvWithCleanup
            (\_ -> Box <$> getClass (SClass "java/lang/Math"))
            (\_ (Box c) -> deleteLocalRef c) $
            \ ~(Box c) -> do
              _ <- getStaticMethodID c "abs" absSignature
              return ()
        , bench "getMethodID" $
          perBatchEnvWithCleanup
            (\_ -> Box <$> getClass (SClass "java/lang/Integer"))
            (\_ (Box c) -> deleteLocalRef c) $
            \ ~(Box c) -> do
              _ <- getMethodID c "intValue" (methodSignature [] (SPrim "int"))
              return ()
        ]
      , bgroup "Haskell calls"
        [ bench "incr haskell" $ nfIO $ incrHaskell 1
        , bench "ffi haskell" $ nfIO $ getpid
        ]
      ]
  where
    absSignature = methodSignature [SomeSing (sing :: Sing ('Prim "int"))] (SPrim "int")
    ini = do
      klass <- findClass (referenceTypeName (SClass "java/lang/Math"))
      method <- getStaticMethodID klass "abs" absSignature
      return (Box klass, method)

benchRefs :: Benchmark
benchRefs =
    env (Box <$> (new >>= newGlobalRefNonFinalized)) $ \ ~(Box (jobj :: JObject)) ->
    bgroup "References"
    [ bench "local reference" $ nfIO $ do
        _ <- newLocalRef jobj
        return ()
    , bench "global reference" $ nfIO $ do
        _ <- newGlobalRef jobj
        return ()
    ,  bench "global reference (no finalizer)" $ nfIO $ do
        _ <- newGlobalRefNonFinalized jobj
        return ()
    -- The next three benchmarks are to be compared with one another:
    -- The goal is to evaluate the cost of attaching a thread to the JVM
    -- when deleting a non-finalized global ref, versus the cost
    -- of having a dedicated thread to do the deleting
    , bench "delete global reference in attached thread" $ nfIO $
        newGlobalRefNonFinalized jobj >>= deleteGlobalRefNonFinalized
    , envWithCleanup
      detachCurrentThread
      (const attachCurrentThreadAsDaemon) $
      \_ -> bench "delete global reference in non-attached thread" $ nfIO $ runInAttachedThread $
        newGlobalRefNonFinalized jobj >>= deleteGlobalRefNonFinalized
    , bench "pass global references to another thread for deletion" $ nfIO $
        newGlobalRefNonFinalized jobj >>=
          submitToFinalizerThread . deleteGlobalRefNonFinalized
    , bench "Foreign.Concurrent.newForeignPtr" $ nfIO $ do
        _ <- Concurrent.newForeignPtr (unsafeObjectToPtr jobj) (return ())
        return ()
    , bench "Foreign.ForeignPtr.newForeignPtr" $ nfIO $ do
        -- Approximate cost of malloc: 50ns. Ideally would move out of benchmark
        -- but finalizer not idempotent.
        ptr <- mallocBytes 4
        _ <- ForeignPtr.newForeignPtr finalizerFree ptr
        return ()
    , bench "local frame / 1 reference" $ nfIO $ do
        pushLocalFrame 30
        _ <- newLocalRef jobj
        _ <- popLocalFrame jnull
        return ()
    , bench "delete 1 local ref" $ nfIO $
        newLocalRef jobj >>= deleteLocalRef
    , bench "local frame / 30 references" $ nfIO $ do
        pushLocalFrame 30
        replicateM_ 30 $ newLocalRef jobj
        _ <- popLocalFrame jnull
        return ()
    , bench "delete 30 local refs" $ nfIO $
        replicateM_ 30 $ newLocalRef jobj >>= deleteLocalRef
    ]

benchNew :: Benchmark
benchNew =
    bgroup "new"
    [ bench "Integer" $
      perBatchEnvWithCleanup
        (pushLocalFrame . (2*) . fromIntegral)
        (\_ _ -> void (popLocalFrame jnull)) $
        \() -> do
          _ <- new (2 :: Int32) :: IO (J ('Class "java.lang.Integer"))
          return ()
    , bench "Integer.valueOf" $
      perBatchEnvWithCleanup
        (pushLocalFrame . (2*) . fromIntegral)
        (\_ _ -> void (popLocalFrame jnull)) $
        \() -> do
          _ <- callStatic "java.lang.Integer" "valueOf" (2 :: Int32)
                 :: IO (J ('Class "java.lang.Integer"))
          return ()
    , envWithCleanup allocTextPtr freeTextPtr $ \ ~(Box (ptr, len)) ->
      bench "newString" $
      perBatchEnvWithCleanup
        (pushLocalFrame . (2*) . fromIntegral)
        (\_ _ -> void (popLocalFrame jnull)) $
        \() ->
          void $ newString ptr (fromIntegral len)
    , envWithCleanup allocTextPtr freeTextPtr $ \ ~(Box (_ptr, len)) ->
      bench "newArray" $
      perBatchEnvWithCleanup
        (pushLocalFrame . (2*) . fromIntegral)
        (\_ _ -> void (popLocalFrame jnull)) $
        \() ->
          void $ newByteArray (fromIntegral len)
    , envWithCleanup allocTextPtr freeTextPtr $ \ ~(Box (ptr, len)) ->
      bench "newDirectByteBuffer" $
      perBatchEnvWithCleanup
        (pushLocalFrame . (2*) . fromIntegral)
        (\_ _ -> void (popLocalFrame jnull)) $
        \() ->
          void $ newDirectByteBuffer (castPtr ptr) (2 * len)
    ]
  where
    allocTextPtr = do
      let len = 128
      dst <- callocArray len
      return $ Box (dst, fromIntegral len)
    freeTextPtr (Box (p, _)) = free p

benchArrays :: Benchmark
benchArrays =
    bgroup "Arrays" $ (++ otherBenchmarks) $ (`map` [128, 256, 512]) $ \arraySize ->
    let n :: Num b => b
        n = fromIntegral (arraySize :: Int) in
    env (Box <$> mallocArray n) $ \ ~(Box bytes) ->
    env (Box <$> newArray n) $ \ ~(Box jbytes) ->
    bgroup (show (n :: Int))
    [ bench "getByteArrayElements" $
      perBatchEnvWithCleanup (\_ -> newIORef []) (const cleanArrays) $
      \ref -> do
        p <- getByteArrayElements jbytes
        modifyIORef ref ((jbytes, p) :)
    , bench "releaseByteArrayElements" $
      perBatchEnv
        (\batchSize ->
           replicateM (fromIntegral batchSize) (getByteArrayElements jbytes)
           >>= newIORef
        ) $
      \ref -> do
         arrays <- readIORef ref
         case arrays of
           x : xs -> do
             releaseByteArrayElements jbytes x
             writeIORef ref xs
           _ -> error "not enough arrays"
    , bench "getByteArrayRegion" $ nfIO $
         getByteArrayRegion jbytes 0 n bytes
    , bench "setByteArrayRegion" $ nfIO $
         setByteArrayRegion jbytes 0 n bytes
    ]
  where
    otherBenchmarks =
     [ bench "getObjectArrayElement" $
       perBatchEnvWithCleanup
         (\batchSize -> do
            pushLocalFrame (2 * fromIntegral batchSize)
            Box <$> newArray 100
         )
         (\_ _ -> void (popLocalFrame jnull)) $
         \(Box jObjectArray) -> do
           _ <- getObjectArrayElement (jObjectArray :: JObjectArray) 40 :: IO JObject
           return ()
     ]

    cleanArrays ref =
      readIORef ref >>= mapM_ (uncurry releaseByteArrayElements)

benchDirectBuffers :: Benchmark
benchDirectBuffers =
    bgroup "DirectBuffers" $ (`map` [128, 256, 512]) $ \bufferSize ->
    let n :: Num b => b
        n = fromIntegral (bufferSize :: Int) in
    env (Box <$> mallocArray n) $ \ ~(Box bytes) ->
    bgroup (show bufferSize)
    [ bench "getDirectBufferAddress" $
      perBatchEnvWithCleanup
        (\_ -> Box <$> newDirectByteBuffer bytes n)
        (\_ -> deleteLocalRef . unBox) $
      void . getDirectBufferAddress . unBox
    , bench "getDirectBufferCapacity" $
      perBatchEnvWithCleanup
        (\_ -> Box <$> newDirectByteBuffer bytes n)
        (\_ -> deleteLocalRef . unBox) $
      void . getDirectBufferCapacity . unBox
    ]

main :: IO ()
main = withJVM [] $ do
    Criterion.defaultMain
      [ benchCalls
      , benchRefs
      , benchNew
      , benchArrays
      , benchDirectBuffers
      ]