module ParkBench.Internal.Driver
  ( BenchmarkConfig (..),
    benchmark,
    LiveBenchmark,
    sampleLiveBenchmark,
    LiveBenchmarks,
    makeLiveBenchmarks,
    stepLiveBenchmarks,
  )
where

import Data.Foldable (toList)
import Data.IORef
import ParkBench.Internal.Array1 (Array1)
import qualified ParkBench.Internal.Array1 as Array1
import ParkBench.Internal.Benchable (Benchable)
import qualified ParkBench.Internal.Benchable as Benchable
import ParkBench.Internal.Prelude
import ParkBench.Internal.RtsStats (RtsStats)
import ParkBench.Internal.Statistics

data BenchmarkConfig = BenchmarkConfig
  { BenchmarkConfig -> Rational
runlen :: !Rational
  }

data LiveBenchmark a = LiveBenchmark
  { forall a. LiveBenchmark a -> IO (Estimate a)
_liveBenchmarkSample :: !(IO (Estimate a)),
    forall a. LiveBenchmark a -> Pull a
_liveBenchmarkPull :: !(Pull a)
  }

sampleLiveBenchmark :: LiveBenchmark a -> IO (Estimate a)
sampleLiveBenchmark :: forall a. LiveBenchmark a -> IO (Estimate a)
sampleLiveBenchmark =
  forall a. LiveBenchmark a -> IO (Estimate a)
_liveBenchmarkSample
{-# INLINE sampleLiveBenchmark #-}

benchmark :: forall a. Roll a => BenchmarkConfig -> Benchable (Timed a) -> IO (LiveBenchmark a)
benchmark :: forall a.
Roll a =>
BenchmarkConfig -> Benchable (Timed a) -> IO (LiveBenchmark a)
benchmark BenchmarkConfig {Rational
runlen :: Rational
$sel:runlen:BenchmarkConfig :: BenchmarkConfig -> Rational
runlen} Benchable (Timed a)
benchable = do
  Estimate a
e0 <- forall a. Timed a -> Estimate a
initialEstimate forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Benchable a -> Word64 -> IO a
Benchable.run Benchable (Timed a)
benchable Word64
1
  IORef (Estimate a)
estimateRef <- forall a. a -> IO (IORef a)
newIORef Estimate a
e0
  let go :: Estimate a -> Pull a
      go :: Estimate a -> Pull a
go Estimate a
oldEstimate =
        forall a. Rational -> IO (Pull a) -> Pull a
Pull (forall a. Estimate a -> Rational
elapsed Estimate a
oldEstimate) do
          let newIters :: Word64
newIters = forall a. Estimate a -> Rational -> Word64
itersInNanoseconds Estimate a
oldEstimate Rational
runlen
          Timed a
newTime <- forall a. Benchable a -> Word64 -> IO a
Benchable.run Benchable (Timed a)
benchable Word64
newIters
          let !newEstimate :: Estimate a
newEstimate = forall a. Roll a => Word64 -> Timed a -> Estimate a -> Estimate a
updateEstimate Word64
newIters Timed a
newTime Estimate a
oldEstimate
          forall a. IORef a -> a -> IO ()
writeIORef IORef (Estimate a)
estimateRef Estimate a
newEstimate
          forall (f :: * -> *) a. Applicative f => a -> f a
pure (Estimate a -> Pull a
go Estimate a
newEstimate)
  forall (f :: * -> *) a. Applicative f => a -> f a
pure
    LiveBenchmark
      { $sel:_liveBenchmarkSample:LiveBenchmark :: IO (Estimate a)
_liveBenchmarkSample = forall a. IORef a -> IO a
readIORef IORef (Estimate a)
estimateRef,
        $sel:_liveBenchmarkPull:LiveBenchmark :: Pull a
_liveBenchmarkPull = Estimate a -> Pull a
go Estimate a
e0
      }
{-# SPECIALIZE benchmark :: BenchmarkConfig -> Benchable (Timed RtsStats) -> IO (LiveBenchmark RtsStats) #-}

-- Given this latest estimate, how many iters could we run in the given number of nanoseconds?
itersInNanoseconds :: Estimate a -> Rational -> Word64
itersInNanoseconds :: forall a. Estimate a -> Rational -> Word64
itersInNanoseconds Estimate {Timed a
$sel:mean:Estimate :: forall a. Estimate a -> Timed a
mean :: Timed a
mean, Word64
$sel:samples:Estimate :: forall a. Estimate a -> Word64
samples :: Word64
samples} Rational
nanos =
  forall a. Ord a => a -> a -> a
max Word64
1 (forall a. Ord a => a -> a -> a
min Word64
samples (forall a b. (RealFrac a, Integral b) => a -> b
floor (Rational
nanos forall a. Fractional a => a -> a -> a
/ forall a. Timed a -> Rational
nanoseconds Timed a
mean)))

data Pull a = Pull
  { -- amount of time this pull has gotten
    forall a. Pull a -> Rational
pullElapsed :: {-# UNPACK #-} !Rational,
    forall a. Pull a -> IO (Pull a)
pullNext :: !(IO (Pull a))
  }

isMoreUrgentThan :: Pull a -> Pull a -> Bool
isMoreUrgentThan :: forall a. Pull a -> Pull a -> Bool
isMoreUrgentThan Pull a
p0 Pull a
p1 =
  forall a. Pull a -> Rational
pullElapsed Pull a
p0 forall a. Ord a => a -> a -> Bool
< forall a. Pull a -> Rational
pullElapsed Pull a
p1

-- | A @LiveBenchmarks@ represents the suspended state of a collection of 1+ benchmarks.
data LiveBenchmarks a
  = -- Most benchmark runs are probably only comparing 1-3 things, so we optimize those cases.
    P1 !(Pull a)
  | P2 !(Pull a) !(Pull a)
  | P3 !(Pull a) !(Pull a) !(Pull a)
  | -- invariant: 4+ elements
    Pn_ ![Pull a]

pattern Pn :: Pull a -> [Pull a] -> LiveBenchmarks a
pattern $mPn :: forall {r} {a}.
LiveBenchmarks a -> (Pull a -> [Pull a] -> r) -> ((# #) -> r) -> r
Pn p ps <- Pn_ (p : ps)

{-# COMPLETE P1, P2, P3, Pn #-}

-- | Construct a 'LiveBenchmarks' from a non-empty array of 'LiveBenchmark'.
makeLiveBenchmarks :: Array1 (LiveBenchmark a) -> LiveBenchmarks a
makeLiveBenchmarks :: forall a. Array1 (LiveBenchmark a) -> LiveBenchmarks a
makeLiveBenchmarks (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. LiveBenchmark a -> Pull a
_liveBenchmarkPull -> Array1 (Pull a)
xs)
  | Int
n forall a. Eq a => a -> a -> Bool
== Int
1 = forall a. Pull a -> LiveBenchmarks a
P1 (forall a. Int -> Array1 a -> a
Array1.get Int
0 Array1 (Pull a)
xs)
  | Int
n forall a. Eq a => a -> a -> Bool
== Int
2 = forall a. Pull a -> Pull a -> LiveBenchmarks a
P2 (forall a. Int -> Array1 a -> a
Array1.get Int
0 Array1 (Pull a)
xs) (forall a. Int -> Array1 a -> a
Array1.get Int
1 Array1 (Pull a)
xs)
  | Int
n forall a. Eq a => a -> a -> Bool
== Int
3 = forall a. Pull a -> Pull a -> Pull a -> LiveBenchmarks a
P3 (forall a. Int -> Array1 a -> a
Array1.get Int
0 Array1 (Pull a)
xs) (forall a. Int -> Array1 a -> a
Array1.get Int
1 Array1 (Pull a)
xs) (forall a. Int -> Array1 a -> a
Array1.get Int
2 Array1 (Pull a)
xs)
  | Bool
otherwise = forall a. [Pull a] -> LiveBenchmarks a
Pn_ (forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Array1 (Pull a)
xs)
  where
    n :: Int
n = forall (t :: * -> *) a. Foldable t => t a -> Int
length Array1 (Pull a)
xs

-- | Step forward a 'LiveBenchmarks', which blocks until the benchmark that has heretofore accumulated the smallest
-- amount of runtime finishes one more run.
--
-- Returns the 'LiveBenchmarks' to use next time, which reflects the latest benchmark run that just completed.
stepLiveBenchmarks :: LiveBenchmarks a -> IO (LiveBenchmarks a)
stepLiveBenchmarks :: forall a. LiveBenchmarks a -> IO (LiveBenchmarks a)
stepLiveBenchmarks = \case
  P1 Pull a
p0 -> do
    Pull a
x0 <- forall a. Pull a -> IO (Pull a)
pullNext Pull a
p0
    forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Pull a -> LiveBenchmarks a
P1 Pull a
x0)
  P2 Pull a
p0 Pull a
x1 -> do
    Pull a
x0 <- forall a. Pull a -> IO (Pull a)
pullNext Pull a
p0
    forall (f :: * -> *) a. Applicative f => a -> f a
pure
      if Pull a
x1 forall a. Pull a -> Pull a -> Bool
`isMoreUrgentThan` Pull a
x0
        then forall a. Pull a -> Pull a -> LiveBenchmarks a
P2 Pull a
x1 Pull a
x0
        else forall a. Pull a -> Pull a -> LiveBenchmarks a
P2 Pull a
x0 Pull a
x1
  P3 Pull a
p0 Pull a
x1 Pull a
x2 -> do
    Pull a
x0 <- forall a. Pull a -> IO (Pull a)
pullNext Pull a
p0
    forall (f :: * -> *) a. Applicative f => a -> f a
pure
      if Pull a
x1 forall a. Pull a -> Pull a -> Bool
`isMoreUrgentThan` Pull a
x0
        then
          if Pull a
x2 forall a. Pull a -> Pull a -> Bool
`isMoreUrgentThan` Pull a
x0
            then forall a. Pull a -> Pull a -> Pull a -> LiveBenchmarks a
P3 Pull a
x1 Pull a
x2 Pull a
x0
            else forall a. Pull a -> Pull a -> Pull a -> LiveBenchmarks a
P3 Pull a
x1 Pull a
x0 Pull a
x2
        else forall a. Pull a -> Pull a -> Pull a -> LiveBenchmarks a
P3 Pull a
x0 Pull a
x1 Pull a
x2
  Pn Pull a
p0 [Pull a]
xs -> do
    Pull a
x0 <- forall a. Pull a -> IO (Pull a)
pullNext Pull a
p0
    forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. [Pull a] -> LiveBenchmarks a
Pn_ (forall a. Pull a -> [Pull a] -> [Pull a]
insertPull Pull a
x0 [Pull a]
xs))

insertPull :: Pull a -> [Pull a] -> [Pull a]
insertPull :: forall a. Pull a -> [Pull a] -> [Pull a]
insertPull Pull a
x0 = \case
  [] -> [Pull a
x0]
  Pull a
x1 : [Pull a]
xs ->
    if Pull a
x0 forall a. Pull a -> Pull a -> Bool
`isMoreUrgentThan` Pull a
x1
      then Pull a
x0 forall a. a -> [a] -> [a]
: Pull a
x1 forall a. a -> [a] -> [a]
: [Pull a]
xs
      else Pull a
x1 forall a. a -> [a] -> [a]
: forall a. Pull a -> [Pull a] -> [Pull a]
insertPull Pull a
x0 [Pull a]
xs