content: {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE QuasiQuotes #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TemplateHaskell #-}
  {-# LANGUAGE TypeApplications #-}
  
  module Main where
  
  import Control.DeepSeq (NFData(..))
  import Criterion.Main as Criterion
  import Control.Monad (void)
  import qualified Data.Coerce as Coerce
  import Data.Int
  import Data.Singletons (SomeSing(..))
  import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
  import qualified Foreign.JNI as JNI
  import qualified Foreign.JNI.Types as JNI
  import Foreign.Ptr
  import GHC.Stable
  import Language.Java
  import Language.Java.Inline
  
  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 ()
  
  benchCallbacks :: Benchmark
  benchCallbacks =
      env ini $ \ ~(Box (fun, obj, jstr, jstr2, jstr3, jstr4)) ->
        bgroup "Callbacks"
        [ bgroup "inline-java"
          [ bench "return-x100" $ withLocalFrame 1 $
              void @_ @JObject [java| {
                  String s = null;
                  for(int i=0;i<100;++i) {
                    s = ((java.util.function.BiFunction<String,String,String>) $fun).apply($jstr, $jstr2);          };
                  return s.concat($jstr);
                } |]
          , bench "no-callback" $ withLocalFrame 1 $
              void @_ @JObject [java| { if ($jstr3 == null){}; return $jstr.concat($jstr2); } |]
          , bench "no-callback-4-args" $ withLocalFrame 1 $
              void @_ @JObject [java| { if ($jstr3 == $jstr4){}; return $jstr.concat($jstr2); } |]
          , bench "loadJavaWrappers" $ nfIO $ loadJavaWrappers
          ]
        , bgroup "jvm"
          [ bench "return" $ withLocalFrame 1 $
              void @_ @JObject $ call fun "apply" obj (JNI.upcast jstr)
          , bench "no-callback" $ withLocalFrame 1 $
              void @_ @JString $ call jstr "concat" jstr2
          ]
        ]
    where
      ini = do
        loadJavaWrappers
        Box <$> ((,,,,,)
                  <$> newCallbackFunction (const . return)
                  <*> ([java| new Object() |] :: IO JObject)
                  <*> ([java| "Hello, " |] :: IO JString)
                  <*> ([java| "World!" |] :: IO JString)
                  <*> ([java| "Hello again, " |] :: IO JString)
                  <*> ([java| "The World." |] :: IO JString)
                )
      withLocalFrame refsPerRun action =
        perBatchEnvWithCleanup
          (JNI.pushLocalFrame . (refsPerRun*) . fromIntegral)
          (\_ _ -> void (JNI.popLocalFrame jnull)) $
          \() -> action
  
  type JNIFun = JNIEnv -> Ptr JObject -> Ptr JObject -> Ptr JObject -> IO (Ptr JObject)
  
  foreign import ccall "wrapper" wrapObjectFun
    :: JNIFun -> IO (FunPtr JNIFun)
  
  -- Export only to get a FunPtr.
  foreign export ccall "inline_java_callbacks_benchmark_freeIterator" freeIterator
    :: JNIEnv -> Ptr JObject -> Int64 -> IO ()
  foreign import ccall "&inline_java_callbacks_benchmark_freeIterator" freeIteratorPtr
    :: FunPtr (JNIEnv -> Ptr JObject -> Int64 -> IO ())
  
  data FunPtrTable = FunPtrTable
    { callbackPtr :: FunPtr JNIFun
    }
  
  freeIterator :: JNIEnv -> Ptr JObject -> Int64 -> IO ()
  freeIterator _ _ ptr = do
      let sptr = castPtrToStablePtr $ intPtrToPtr $ fromIntegral ptr
      FunPtrTable{..} <- deRefStablePtr sptr
      freeHaskellFunPtr callbackPtr
      freeStablePtr sptr
  
  -- | Reflects a stream with no batching.
  newCallbackFunction
    :: (JObject -> JObject -> IO JObject)
    -> IO (J ('Iface "java.util.function.BiFunction"))
  newCallbackFunction f = do
      callbackPtr <- wrapObjectFun $ \_ _jthis jarg1 jarg2 ->
        unsafeForeignPtrToPtr <$> Coerce.coerce <$> do
          a1 <- JNI.objectFromPtr jarg1
          a2 <- JNI.objectFromPtr jarg2
          f a1 a2
      -- Keep FunPtr's in a table that can be referenced from the Java side, so
      -- that they can be freed.
      tblPtr :: Int64 <- fromIntegral . ptrToIntPtr . castStablePtrToPtr <$> newStablePtr FunPtrTable{..}
      fun <-
        [java| new java.util.function.BiFunction() {
            @Override
            public native Object apply(Object o1, Object o2);
  
            private native void hsFinalize(long tblPtr);
  
            @Override
            public void finalize() { hsFinalize($tblPtr); }
          } |]
      klass <- JNI.getObjectClass fun >>= JNI.newGlobalRef
      JNI.registerNatives klass
        [ JNI.JNINativeMethod
            "apply"
            (methodSignature
              [ SomeSing (sing :: Sing ('Class "java.lang.Object"))
              , SomeSing (sing :: Sing ('Class "java.lang.Object"))
              ]
              (sing :: Sing ('Class "java.lang.Object")))
            callbackPtr
        , JNI.JNINativeMethod
            "hsFinalize"
            (methodSignature [SomeSing (sing :: Sing ('Prim "long"))] (sing :: Sing 'Void))
            freeIteratorPtr
        ]
      return fun
  
  
  main :: IO ()
  main = withJVM [] $ do
      Criterion.defaultMain
        [ benchCallbacks ]
  