content: {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE DeriveGeneric #-}
  {-# LANGUAGE DeriveAnyClass #-}
  {-# LANGUAGE FlexibleContexts #-}
  module Main ( main ) where
  
  import Criterion.Main
  import qualified Language.Souffle.Compiled as S
  import qualified Data.Text as T
  import qualified Data.Vector as V
  import GHC.Generics
  import Data.Word
  import Data.Int
  import Control.Monad
  import Control.Monad.IO.Class
  import Control.DeepSeq
  
  
  data Benchmarks = Benchmarks
  
  data NumbersFact
    = NumbersFact Word32 Int32 Float
    deriving (Generic, NFData)
  
  data StringsFact
    = StringsFact Word32 T.Text Int32 Float
    deriving (Generic, NFData)
  
  newtype FromDatalogFact
    = FromDatalogFact Int32
    deriving (Generic, NFData)
  
  data FromDatalogStringFact
    = FromDatalogStringFact Int32 T.Text
    deriving (Generic, NFData)
  
  instance S.Program Benchmarks where
    type ProgramFacts Benchmarks =
        '[NumbersFact, StringsFact, FromDatalogFact, FromDatalogStringFact]
    programName = const "bench"
  
  instance S.Fact NumbersFact where
    type FactDirection NumbersFact = 'S.InputOutput
    factName = const "numbers_fact"
  
  instance S.Fact StringsFact where
    type FactDirection StringsFact = 'S.InputOutput
    factName = const "strings_fact"
  
  instance S.Fact FromDatalogFact where
    type FactDirection FromDatalogFact = 'S.InputOutput
    factName = const "from_datalog_fact"
  
  instance S.Fact FromDatalogStringFact where
    type FactDirection FromDatalogStringFact = 'S.InputOutput
    factName = const "from_datalog_string_fact"
  
  instance S.Marshal NumbersFact
  instance S.Marshal StringsFact
  instance S.Marshal FromDatalogFact
  instance S.Marshal FromDatalogStringFact
  
  
  -- TODO: fix cases with larger numbers (crashes due to large memory allocations?)
  main :: IO ()
  main = defaultMain
       $ roundTripBenchmarks
      ++ serializationBenchmarks
      ++ deserializationBenchmarks
  
  roundTripBenchmarks :: [Benchmark]
  roundTripBenchmarks =
    [ bgroup "round trip facts (without strings)"
      [ bench "1"      $ nfIO $ roundTrip $ mkVec 1
      , bench "10"     $ nfIO $ roundTrip $ mkVec 10
      , bench "100"    $ nfIO $ roundTrip $ mkVec 100
      , bench "1000"   $ nfIO $ roundTrip $ mkVec 1000
      , bench "10000"  $ nfIO $ roundTrip $ mkVec 10000
      --, bench "100000" $ nfIO $ roundTrip $ mkVec 100000
      ]
    , bgroup "round trip facts (with strings)"
      [ bench "1"      $ nfIO $ roundTrip $ mkVecStr 1
      , bench "10"     $ nfIO $ roundTrip $ mkVecStr 10
      , bench "100"    $ nfIO $ roundTrip $ mkVecStr 100
      , bench "1000"   $ nfIO $ roundTrip $ mkVecStr 1000
      --, bench "10000"  $ nfIO $ roundTrip $ mkVecStr 10000
      --, bench "100000" $ nfIO $ roundTrip $ mkVecStr 100000
      ]
    , bgroup "round trip facts (with long strings)"
      [ bench "1"      $ nfIO $ roundTrip $ mkVecLongStr 1
      , bench "10"     $ nfIO $ roundTrip $ mkVecLongStr 10
      , bench "100"    $ nfIO $ roundTrip $ mkVecLongStr 100
      --, bench "1000"   $ nfIO $ roundTrip $ mkVecLongStr 1000
      --, bench "10000"  $ nfIO $ roundTrip $ mkVecLongStr 10000
      --, bench "100000" $ nfIO $ roundTrip $ mkVecLongStr 100000
      ]
    ]
    where mkVec count = V.generate count $ \i -> NumbersFact (fromIntegral i) (-42) 3.14
          mkVecStr count = V.generate count $ \i -> StringsFact (fromIntegral i) "abcdef" (-42) 3.14
          mkVecLongStr count = V.generate count $ \i -> StringsFact (fromIntegral i) (T.replicate 10 "abcdef") (-42) 3.14
  
  roundTrip :: (S.ContainsInputFact Benchmarks a, S.ContainsOutputFact Benchmarks a, S.Fact a, S.Submit a)
            => V.Vector a -> IO (V.Vector a)
  roundTrip vec = S.runSouffle Benchmarks $ \case
    Nothing -> do
      liftIO $ print "Failed to load roundtrip benchmarks!"
      pure V.empty
    Just prog -> do
      S.addFacts prog vec
      -- No run needed
      S.getFacts prog
  
  
  serializeNumbers :: Int -> IO ()
  serializeNumbers iterationCount = S.runSouffle Benchmarks $ \case
    Nothing -> liftIO $ print "Failed to load serialize benchmarks!"
    Just prog ->
      replicateM_ iterationCount $ S.addFacts prog vec
      -- No run needed
    where vec = V.generate 100 $ \i -> NumbersFact (fromIntegral i) (-42) 3.14
  
  deserializeNumbers :: Int -> IO ()
  deserializeNumbers iterationCount = S.runSouffle Benchmarks $ \case
    Nothing -> liftIO $ print "Failed to load deserialize benchmarks!"
    Just prog -> do
      S.run prog
      replicateM_ iterationCount $ do
        fs <- S.getFacts prog
        pure (fs :: V.Vector FromDatalogFact)
  
  serializeWithStrings :: Int -> IO ()
  serializeWithStrings iterationCount = S.runSouffle Benchmarks $ \case
    Nothing -> liftIO $ print "Failed to load serialize benchmarks!"
    Just prog ->
      replicateM_ iterationCount $ S.addFacts prog vec
      -- No run needed
    where vec = V.generate 100 $ \i -> StringsFact (fromIntegral i) "abcdef" (-42) 3.14
  
  deserializeWithStrings :: Int -> IO ()
  deserializeWithStrings iterationCount = S.runSouffle Benchmarks $ \case
    Nothing -> liftIO $ print "Failed to load deserialize benchmarks!"
    Just prog -> do
      S.run prog
      replicateM_ iterationCount $ do
        fs <- S.getFacts prog
        pure (fs :: V.Vector FromDatalogStringFact)
  
  serializationBenchmarks :: [Benchmark]
  serializationBenchmarks =
    [ bgroup "serializing facts (without strings)"
      [ bench "1"      $ nfIO $ serializeNumbers 1
      , bench "10"     $ nfIO $ serializeNumbers 10
      , bench "100"    $ nfIO $ serializeNumbers 100
      , bench "1000"   $ nfIO $ serializeNumbers 1000
      , bench "10000"  $ nfIO $ serializeNumbers 10000
      ]
    , bgroup "serializing facts (with strings)"
      [ bench "1"      $ nfIO $ serializeWithStrings 1
      , bench "10"     $ nfIO $ serializeWithStrings 10
      , bench "100"    $ nfIO $ serializeWithStrings 100
      , bench "1000"   $ nfIO $ serializeWithStrings 1000
      , bench "10000"  $ nfIO $ serializeWithStrings 10000
      ]
    ]
  
  deserializationBenchmarks :: [Benchmark]
  deserializationBenchmarks =
    [ bgroup "deserializing facts (without strings)"
      [ bench "1"      $ nfIO $ deserializeNumbers 1
      , bench "10"     $ nfIO $ deserializeNumbers 10
      , bench "100"    $ nfIO $ deserializeNumbers 100
      , bench "1000"   $ nfIO $ deserializeNumbers 1000
      , bench "10000"  $ nfIO $ deserializeNumbers 10000
      ]
    , bgroup "deserializing facts (with strings)"
      [ bench "1"      $ nfIO $ deserializeWithStrings 1
      , bench "10"     $ nfIO $ deserializeWithStrings 10
      , bench "100"    $ nfIO $ deserializeWithStrings 100
      , bench "1000"   $ nfIO $ deserializeWithStrings 1000
      , bench "10000"  $ nfIO $ deserializeWithStrings 10000
      ]
    ]
  