module ParkBench.Driver
  ( benchmark1,
    Pull1 (..),
    benchmark,
    Pull,
    Pulls,
    pulls,
    pull,
  )
where

import Data.IORef
import Data.List.NonEmpty (NonEmpty (..))
import qualified Data.List.NonEmpty as NonEmpty
import ParkBench.Prelude
import ParkBench.RtsStats (RtsStats)
import ParkBench.Statistics

newtype Pull1 a
  = Pull1 (IO (Estimate a, Pull1 a))

-- | Like 'benchmark', but optimized for only running one benchmark.
benchmark1 :: forall a. Roll a => (Word64 -> IO (Timed a)) -> Pull1 a
benchmark1 :: (Word64 -> IO (Timed a)) -> Pull1 a
benchmark1 Word64 -> IO (Timed a)
run =
  IO (Estimate a, Pull1 a) -> Pull1 a
forall a. IO (Estimate a, Pull1 a) -> Pull1 a
Pull1 do
    Timed a
t <- Word64 -> IO (Timed a)
run Word64
1
    let another :: Estimate a -> Pull1 a
        another :: Estimate a -> Pull1 a
another Estimate a
e0 =
          IO (Estimate a, Pull1 a) -> Pull1 a
forall a. IO (Estimate a, Pull1 a) -> Pull1 a
Pull1 do
            Timed a
t2 <- Word64 -> IO (Timed a)
run Word64
n
            (Estimate a, Pull1 a) -> IO (Estimate a, Pull1 a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Estimate a -> (Estimate a, Pull1 a)
andAnother (Word64 -> Timed a -> Estimate a -> Estimate a
forall a. Roll a => Word64 -> Timed a -> Estimate a -> Estimate a
updateEstimate Word64
n Timed a
t2 Estimate a
e0))
          where
            n :: Word64
n = Estimate a -> Word64
forall a. Estimate a -> Word64
next Estimate a
e0
        andAnother :: Estimate a -> (Estimate a, Pull1 a)
        andAnother :: Estimate a -> (Estimate a, Pull1 a)
andAnother Estimate a
e =
          (Estimate a
e, Estimate a -> Pull1 a
another Estimate a
e)
    (Estimate a, Pull1 a) -> IO (Estimate a, Pull1 a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Estimate a -> (Estimate a, Pull1 a)
andAnother (Timed a -> Estimate a
forall a. Timed a -> Estimate a
initialEstimate Timed a
t))
{-# SPECIALIZE benchmark1 :: (Word64 -> IO (Timed RtsStats)) -> Pull1 RtsStats #-}

benchmark :: forall a. Roll a => (Word64 -> IO (Timed a)) -> IO (IO (Estimate a), Pull a)
benchmark :: (Word64 -> IO (Timed a)) -> IO (IO (Estimate a), Pull a)
benchmark Word64 -> IO (Timed a)
run = do
  Timed a
t <- Word64 -> IO (Timed a)
run Word64
1
  let e :: Estimate a
e = Timed a -> Estimate a
forall a. Timed a -> Estimate a
initialEstimate Timed a
t
  IORef (Estimate a)
ref <- Estimate a -> IO (IORef (Estimate a))
forall a. a -> IO (IORef a)
newIORef Estimate a
e
  let another :: Estimate a -> IO (Pull a)
      another :: Estimate a -> IO (Pull a)
another Estimate a
e0 = do
        Timed a
t2 <- Word64 -> IO (Timed a)
run Word64
n
        let !e1 :: Estimate a
e1 = Word64 -> Timed a -> Estimate a -> Estimate a
forall a. Roll a => Word64 -> Timed a -> Estimate a -> Estimate a
updateEstimate Word64
n Timed a
t2 Estimate a
e0
        IORef (Estimate a) -> Estimate a -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Estimate a)
ref Estimate a
e1
        Pull a -> IO (Pull a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Estimate a -> Pull a
andAnother Estimate a
e1)
        where
          n :: Word64
n = Estimate a -> Word64
forall a. Estimate a -> Word64
next Estimate a
e0
      andAnother :: Estimate a -> Pull a
      andAnother :: Estimate a -> Pull a
andAnother Estimate a
e0 =
        Rational -> IO (Pull a) -> Pull a
forall a. Rational -> IO (Pull a) -> Pull a
Pull (Word64 -> Rational
w2r (Estimate a -> Word64
forall a. Estimate a -> Word64
samples Estimate a
e0) Rational -> Rational -> Rational
forall a. Num a => a -> a -> a
* Timed a -> Rational
forall a. Timed a -> Rational
nanoseconds (Estimate a -> Timed a
forall a. Estimate a -> Timed a
mean Estimate a
e0)) (Estimate a -> IO (Pull a)
another Estimate a
e0)
  (IO (Estimate a), Pull a) -> IO (IO (Estimate a), Pull a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IORef (Estimate a) -> IO (Estimate a)
forall a. IORef a -> IO a
readIORef IORef (Estimate a)
ref, Estimate a -> Pull a
andAnother Estimate a
e)
{-# SPECIALIZE benchmark :: (Word64 -> IO (Timed RtsStats)) -> IO (IO (Estimate RtsStats), Pull RtsStats) #-}

-- target runs that take 0.1 seconds (e.g. 500_000_000 would be 0.5 seconds)
next :: Estimate a -> Word64
next :: Estimate a -> Word64
next Estimate {$sel:mean:Estimate :: forall a. Estimate a -> Timed a
mean = Timed Rational
nanoseconds a
_, Word64
samples :: Word64
$sel:samples:Estimate :: forall a. Estimate a -> Word64
samples} =
  Word64 -> Word64 -> Word64
forall a. Ord a => a -> a -> a
max Word64
1 (Word64 -> Word64 -> Word64
forall a. Ord a => a -> a -> a
min Word64
samples (Rational -> Word64
forall a b. (RealFrac a, Integral b) => a -> b
floor (Rational
100_000_000 Rational -> Rational -> Rational
forall a. Fractional a => a -> a -> a
/ Rational
nanoseconds)))

data Pull a
  = Pull
      -- amount of time this pull has gotten
      {-# UNPACK #-} !Rational
      !(IO (Pull a))

isMoreUrgentThan :: Pull a -> Pull a -> Bool
Pull Rational
t0 IO (Pull a)
_ isMoreUrgentThan :: Pull a -> Pull a -> Bool
`isMoreUrgentThan` Pull Rational
t1 IO (Pull a)
_ =
  Rational
t0 Rational -> Rational -> Bool
forall a. Ord a => a -> a -> Bool
< Rational
t1

-- | A @Pulls@ represents the suspended state of a collection of 1+ benchmarks.
data Pulls 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] -> Pulls a
pattern $mPn :: forall r a.
Pulls a -> (Pull a -> [Pull a] -> r) -> (Void# -> r) -> r
Pn p ps <- Pn_ (p : ps)

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

-- | Construct a 'Pulls' from a non-empty list of 'Pull'.
pulls :: NonEmpty (Pull a) -> Pulls a
pulls :: NonEmpty (Pull a) -> Pulls a
pulls =
  NonEmpty (Pull a) -> Pulls a
forall a. NonEmpty (Pull a) -> Pulls a
pulls' (NonEmpty (Pull a) -> Pulls a)
-> (NonEmpty (Pull a) -> NonEmpty (Pull a))
-> NonEmpty (Pull a)
-> Pulls a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Pull a -> Rational) -> NonEmpty (Pull a) -> NonEmpty (Pull a)
forall o a. Ord o => (a -> o) -> NonEmpty a -> NonEmpty a
NonEmpty.sortWith \(Pull Rational
t IO (Pull a)
_) -> Rational
t

pulls' :: NonEmpty (Pull a) -> Pulls a
pulls' :: NonEmpty (Pull a) -> Pulls a
pulls' = \case
  Pull a
a :| [] -> Pull a -> Pulls a
forall a. Pull a -> Pulls a
P1 Pull a
a
  Pull a
a :| [Pull a
b] -> Pull a -> Pull a -> Pulls a
forall a. Pull a -> Pull a -> Pulls a
P2 Pull a
a Pull a
b
  Pull a
a :| [Pull a
b, Pull a
c] -> Pull a -> Pull a -> Pull a -> Pulls a
forall a. Pull a -> Pull a -> Pull a -> Pulls a
P3 Pull a
a Pull a
b Pull a
c
  Pull a
a :| [Pull a]
as -> [Pull a] -> Pulls a
forall a. [Pull a] -> Pulls a
Pn_ (Pull a
a Pull a -> [Pull a] -> [Pull a]
forall a. a -> [a] -> [a]
: [Pull a]
as)

-- | Pull on a 'Pulls', which blocks until the benchmark that has heretofore accumulated the smallest amount of runtime
-- runs once more.
--
-- Returns the 'Pulls' to use next time, which reflects the latest benchmark run that just completed.
pull :: Pulls a -> IO (Pulls a)
pull :: Pulls a -> IO (Pulls a)
pull = \case
  P1 (Pull Rational
_ IO (Pull a)
p0) -> do
    Pull a
p <- IO (Pull a)
p0
    Pulls a -> IO (Pulls a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Pull a -> Pulls a
forall a. Pull a -> Pulls a
P1 Pull a
p)
  P2 (Pull Rational
_ IO (Pull a)
p0) Pull a
q -> do
    Pull a
p <- IO (Pull a)
p0
    Pulls a -> IO (Pulls a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure
      if Pull a
q Pull a -> Pull a -> Bool
forall a. Pull a -> Pull a -> Bool
`isMoreUrgentThan` Pull a
p
        then Pull a -> Pull a -> Pulls a
forall a. Pull a -> Pull a -> Pulls a
P2 Pull a
q Pull a
p
        else Pull a -> Pull a -> Pulls a
forall a. Pull a -> Pull a -> Pulls a
P2 Pull a
p Pull a
q
  P3 (Pull Rational
_ IO (Pull a)
p0) Pull a
q Pull a
r -> do
    Pull a
p <- IO (Pull a)
p0
    Pulls a -> IO (Pulls a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure
      if Pull a
q Pull a -> Pull a -> Bool
forall a. Pull a -> Pull a -> Bool
`isMoreUrgentThan` Pull a
p
        then
          if Pull a
r Pull a -> Pull a -> Bool
forall a. Pull a -> Pull a -> Bool
`isMoreUrgentThan` Pull a
p
            then Pull a -> Pull a -> Pull a -> Pulls a
forall a. Pull a -> Pull a -> Pull a -> Pulls a
P3 Pull a
q Pull a
r Pull a
p
            else Pull a -> Pull a -> Pull a -> Pulls a
forall a. Pull a -> Pull a -> Pull a -> Pulls a
P3 Pull a
q Pull a
p Pull a
r
        else Pull a -> Pull a -> Pull a -> Pulls a
forall a. Pull a -> Pull a -> Pull a -> Pulls a
P3 Pull a
p Pull a
q Pull a
r
  Pn (Pull Rational
_ IO (Pull a)
p0) [Pull a]
ps -> do
    Pull a
p <- IO (Pull a)
p0
    Pulls a -> IO (Pulls a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Pull a] -> Pulls a
forall a. [Pull a] -> Pulls a
Pn_ (Pull a -> [Pull a] -> [Pull a]
forall a. Pull a -> [Pull a] -> [Pull a]
insertPull Pull a
p [Pull a]
ps))

insertPull :: Pull a -> [Pull a] -> [Pull a]
insertPull :: Pull a -> [Pull a] -> [Pull a]
insertPull Pull a
p0 = \case
  [] -> [Pull a
p0]
  Pull a
p1 : [Pull a]
ps ->
    if Pull a
p0 Pull a -> Pull a -> Bool
forall a. Pull a -> Pull a -> Bool
`isMoreUrgentThan` Pull a
p1
      then Pull a
p0 Pull a -> [Pull a] -> [Pull a]
forall a. a -> [a] -> [a]
: Pull a
p1 Pull a -> [Pull a] -> [Pull a]
forall a. a -> [a] -> [a]
: [Pull a]
ps
      else Pull a
p1 Pull a -> [Pull a] -> [Pull a]
forall a. a -> [a] -> [a]
: Pull a -> [Pull a] -> [Pull a]
forall a. Pull a -> [Pull a] -> [Pull a]
insertPull Pull a
p0 [Pull a]
ps