module ParkBench
  ( Benchmark,
    benchmark,
    function,
    action,
  )
where

import Control.Concurrent (threadDelay)
import qualified Data.ByteString as ByteString
import Data.List.NonEmpty (NonEmpty (..))
import qualified Data.List.NonEmpty as NonEmpty
import qualified Data.Text as Text (pack)
import qualified Data.Text.Encoding as Text (encodeUtf8)
import qualified ParkBench.Benchmark as Benchmark
import qualified ParkBench.Builder as Builder
import qualified ParkBench.Driver as Driver
import ParkBench.Named (Named (Named))
import qualified ParkBench.Named as Named
import ParkBench.Prelude
import ParkBench.Pretty (renderTable)
import ParkBench.Render (estimatesToTable)
import ParkBench.RtsStats (RtsStats)
import qualified ParkBench.Statistics as Statistics
import ParkBench.Terminal (clearFromCursor, cursorUp, withTerminal)

-- | A single benchmark.
newtype Benchmark
  = Benchmark (Named (Word64 -> IO ()))

-- | Run a collection of benchmarks.
benchmark ::
  -- |
  [Benchmark] ->
  IO void
benchmark :: [Benchmark] -> IO void
benchmark [Benchmark]
xs =
  case [Benchmark] -> Maybe (NonEmpty Benchmark)
forall a. [a] -> Maybe (NonEmpty a)
NonEmpty.nonEmpty [Benchmark]
xs of
    Maybe (NonEmpty Benchmark)
Nothing -> IO () -> IO void
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (Int -> IO ()
threadDelay Int
forall a. Bounded a => a
maxBound)
    Just NonEmpty Benchmark
ys -> NonEmpty (Named (Word64 -> IO ())) -> IO void
forall void. NonEmpty (Named (Word64 -> IO ())) -> IO void
benchmark' (NonEmpty Benchmark -> NonEmpty (Named (Word64 -> IO ()))
coerce NonEmpty Benchmark
ys)

benchmark' :: NonEmpty (Named (Word64 -> IO ())) -> IO void
benchmark' :: NonEmpty (Named (Word64 -> IO ())) -> IO void
benchmark' = \case
  Named Text
name Word64 -> IO ()
f :| [] ->
    IO void -> IO void
forall a. IO a -> IO a
withTerminal do
      let loop :: Driver.Pull1 RtsStats -> Int -> IO void
          loop :: Pull1 RtsStats -> Int -> IO void
loop (Driver.Pull1 IO (Estimate RtsStats, Pull1 RtsStats)
pull0) Int
newlines0 = do
            (Estimate RtsStats
estimate, Pull1 RtsStats
pull1) <- IO (Estimate RtsStats, Pull1 RtsStats)
pull0
            Int
newlines1 <- NonEmpty (Named (Estimate RtsStats)) -> Int -> IO Int
renderSummaries (Text -> Estimate RtsStats -> Named (Estimate RtsStats)
forall a. Text -> a -> Named a
Named Text
name Estimate RtsStats
estimate Named (Estimate RtsStats)
-> [Named (Estimate RtsStats)]
-> NonEmpty (Named (Estimate RtsStats))
forall a. a -> [a] -> NonEmpty a
:| []) Int
newlines0
            Pull1 RtsStats -> Int -> IO void
forall void. Pull1 RtsStats -> Int -> IO void
loop Pull1 RtsStats
pull1 Int
newlines1
      ByteString -> IO ()
ByteString.putStr (Word8 -> ByteString
ByteString.singleton Word8
newline)
      Pull1 RtsStats -> Int -> IO void
forall void. Pull1 RtsStats -> Int -> IO void
loop ((Word64 -> IO (Timed RtsStats)) -> Pull1 RtsStats
forall a. Roll a => (Word64 -> IO (Timed a)) -> Pull1 a
Driver.benchmark1 (IO () -> IO (Timed RtsStats)
Benchmark.measure (IO () -> IO (Timed RtsStats))
-> (Word64 -> IO ()) -> Word64 -> IO (Timed RtsStats)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word64 -> IO ()
f)) Int
0
  NonEmpty (Named (Word64 -> IO ()))
fs ->
    IO void -> IO void
forall a. IO a -> IO a
withTerminal do
      NonEmpty (Named (IO (Estimate RtsStats), Pull RtsStats))
summaries0 <- ((Named (Word64 -> IO ())
 -> IO (Named (IO (Estimate RtsStats), Pull RtsStats)))
-> NonEmpty (Named (Word64 -> IO ()))
-> IO (NonEmpty (Named (IO (Estimate RtsStats), Pull RtsStats)))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((Named (Word64 -> IO ())
  -> IO (Named (IO (Estimate RtsStats), Pull RtsStats)))
 -> NonEmpty (Named (Word64 -> IO ()))
 -> IO (NonEmpty (Named (IO (Estimate RtsStats), Pull RtsStats))))
-> (((Word64 -> IO ())
     -> IO (IO (Estimate RtsStats), Pull RtsStats))
    -> Named (Word64 -> IO ())
    -> IO (Named (IO (Estimate RtsStats), Pull RtsStats)))
-> ((Word64 -> IO ())
    -> IO (IO (Estimate RtsStats), Pull RtsStats))
-> NonEmpty (Named (Word64 -> IO ()))
-> IO (NonEmpty (Named (IO (Estimate RtsStats), Pull RtsStats)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Word64 -> IO ()) -> IO (IO (Estimate RtsStats), Pull RtsStats))
-> Named (Word64 -> IO ())
-> IO (Named (IO (Estimate RtsStats), Pull RtsStats))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse) (\Word64 -> IO ()
f -> (Word64 -> IO (Timed RtsStats))
-> IO (IO (Estimate RtsStats), Pull RtsStats)
forall a.
Roll a =>
(Word64 -> IO (Timed a)) -> IO (IO (Estimate a), Pull a)
Driver.benchmark (IO () -> IO (Timed RtsStats)
Benchmark.measure (IO () -> IO (Timed RtsStats))
-> (Word64 -> IO ()) -> Word64 -> IO (Timed RtsStats)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word64 -> IO ()
f)) NonEmpty (Named (Word64 -> IO ()))
fs
      let loop :: Driver.Pulls RtsStats -> Int -> IO void
          loop :: Pulls RtsStats -> Int -> IO void
loop Pulls RtsStats
pulls0 Int
newlines0 = do
            NonEmpty (Named (Estimate RtsStats))
summaries <- (Named (IO (Estimate RtsStats), Pull RtsStats)
 -> IO (Named (Estimate RtsStats)))
-> NonEmpty (Named (IO (Estimate RtsStats), Pull RtsStats))
-> IO (NonEmpty (Named (Estimate RtsStats)))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (((IO (Estimate RtsStats), Pull RtsStats) -> IO (Estimate RtsStats))
-> Named (IO (Estimate RtsStats), Pull RtsStats)
-> IO (Named (Estimate RtsStats))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (IO (Estimate RtsStats), Pull RtsStats) -> IO (Estimate RtsStats)
forall a b. (a, b) -> a
fst) NonEmpty (Named (IO (Estimate RtsStats), Pull RtsStats))
summaries0
            Int
newlines1 <- NonEmpty (Named (Estimate RtsStats)) -> Int -> IO Int
renderSummaries NonEmpty (Named (Estimate RtsStats))
summaries Int
newlines0
            Pulls RtsStats
pulls1 <- Pulls RtsStats -> IO (Pulls RtsStats)
forall a. Pulls a -> IO (Pulls a)
Driver.pull Pulls RtsStats
pulls0
            Pulls RtsStats -> Int -> IO void
forall void. Pulls RtsStats -> Int -> IO void
loop Pulls RtsStats
pulls1 Int
newlines1
      ByteString -> IO ()
ByteString.putStr (Word8 -> ByteString
ByteString.singleton Word8
newline)
      Pulls RtsStats -> Int -> IO void
forall void. Pulls RtsStats -> Int -> IO void
loop (NonEmpty (Pull RtsStats) -> Pulls RtsStats
forall a. NonEmpty (Pull a) -> Pulls a
Driver.pulls ((IO (Estimate RtsStats), Pull RtsStats) -> Pull RtsStats
forall a b. (a, b) -> b
snd ((IO (Estimate RtsStats), Pull RtsStats) -> Pull RtsStats)
-> (Named (IO (Estimate RtsStats), Pull RtsStats)
    -> (IO (Estimate RtsStats), Pull RtsStats))
-> Named (IO (Estimate RtsStats), Pull RtsStats)
-> Pull RtsStats
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Named (IO (Estimate RtsStats), Pull RtsStats)
-> (IO (Estimate RtsStats), Pull RtsStats)
forall a. Named a -> a
Named.thing (Named (IO (Estimate RtsStats), Pull RtsStats) -> Pull RtsStats)
-> NonEmpty (Named (IO (Estimate RtsStats), Pull RtsStats))
-> NonEmpty (Pull RtsStats)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NonEmpty (Named (IO (Estimate RtsStats), Pull RtsStats))
summaries0)) Int
0

renderSummaries :: NonEmpty (Named (Statistics.Estimate RtsStats)) -> Int -> IO Int
renderSummaries :: NonEmpty (Named (Estimate RtsStats)) -> Int -> IO Int
renderSummaries NonEmpty (Named (Estimate RtsStats))
summaries Int
newlines0 = do
  ByteString -> IO ()
ByteString.putStr ByteString
bytes
  Int -> IO Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Word8 -> ByteString -> Int
ByteString.count Word8
newline ByteString
bytes)
  where
    bytes :: ByteString
bytes = Text -> ByteString
Text.encodeUtf8 (Builder -> Text
Builder.build Builder
builder)
    builder :: Builder
builder = Int -> Builder
cursorUp Int
newlines0 Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Builder
clearFromCursor Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Table -> Builder
renderTable (NonEmpty (Named (Estimate RtsStats)) -> Table
estimatesToTable NonEmpty (Named (Estimate RtsStats))
summaries)

newline :: Word8
newline :: Word8
newline = Word8
10

-- | Benchmark a function. The result is evaluated to weak head normal form.
function ::
  -- |
  String ->
  -- |
  (a -> b) ->
  -- |
  a ->
  Benchmark
function :: String -> (a -> b) -> a -> Benchmark
function String
name a -> b
f a
x =
  Named (Word64 -> IO ()) -> Benchmark
Benchmark (Text -> (Word64 -> IO ()) -> Named (Word64 -> IO ())
forall a. Text -> a -> Named a
Named (String -> Text
Text.pack String
name) ((a -> b) -> a -> Word64 -> IO ()
forall a b. (a -> b) -> a -> Word64 -> IO ()
Benchmark.whnf a -> b
f a
x))

-- | Benchmark an IO action. The result is evaluated to weak head normal form.
action ::
  -- |
  String ->
  -- |
  IO a ->
  Benchmark
action :: String -> IO a -> Benchmark
action String
name IO a
x =
  Named (Word64 -> IO ()) -> Benchmark
Benchmark (Text -> (Word64 -> IO ()) -> Named (Word64 -> IO ())
forall a. Text -> a -> Named a
Named (String -> Text
Text.pack String
name) (IO a -> Word64 -> IO ()
forall a. IO a -> Word64 -> IO ()
Benchmark.whnfIO IO a
x))