-- | Simple benchmarking of IO functions.
module Cryptol.Utils.Benchmark
  ( BenchmarkStats (..)
  , benchmark
  , secs
  ) where

import           Criterion.Measurement       (runBenchmark, secs, threshold)
import           Criterion.Measurement.Types
import           Data.Int
import qualified Data.Vector                 as V
import qualified Data.Vector.Unboxed         as U

-- | Statistics returned by 'benchmark'.
--
-- This is extremely crude compared to the full analysis that criterion can do,
-- but is enough for now.
data BenchmarkStats = BenchmarkStats
  { BenchmarkStats -> Double
benchAvgTime    :: !Double
  , BenchmarkStats -> Double
benchAvgCpuTime :: !Double
  , BenchmarkStats -> Int64
benchAvgCycles  :: !Int64
  } deriving Int -> BenchmarkStats -> ShowS
[BenchmarkStats] -> ShowS
BenchmarkStats -> String
(Int -> BenchmarkStats -> ShowS)
-> (BenchmarkStats -> String)
-> ([BenchmarkStats] -> ShowS)
-> Show BenchmarkStats
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> BenchmarkStats -> ShowS
showsPrec :: Int -> BenchmarkStats -> ShowS
$cshow :: BenchmarkStats -> String
show :: BenchmarkStats -> String
$cshowList :: [BenchmarkStats] -> ShowS
showList :: [BenchmarkStats] -> ShowS
Show

-- | Benchmark the application of the given function to the given input and the
-- execution of the resulting IO action to WHNF, spending at least the given
-- amount of time in seconds to collect measurements.
benchmark :: Double -> (a -> IO b) -> a -> IO BenchmarkStats
benchmark :: forall a b. Double -> (a -> IO b) -> a -> IO BenchmarkStats
benchmark Double
period a -> IO b
f a
x = do
  (Vector Measured
meas, Double
_) <- Benchmarkable -> Double -> IO (Vector Measured, Double)
runBenchmark ((a -> IO b) -> a -> Benchmarkable
forall a b. (a -> IO b) -> a -> Benchmarkable
whnfAppIO a -> IO b
f a
x) Double
period
  let meas' :: Vector Measured
meas' = Measured -> Measured
rescale (Measured -> Measured) -> Vector Measured -> Vector Measured
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Measured -> Bool) -> Vector Measured -> Vector Measured
forall a. (a -> Bool) -> Vector a -> Vector a
V.filter ((Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
threshold) (Double -> Bool) -> (Measured -> Double) -> Measured -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Measured -> Double
measTime) Vector Measured
meas
      len :: Int
len = Vector Measured -> Int
forall a. Vector a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector Measured
meas'
      sumMeasure :: (Measured -> a) -> a
sumMeasure Measured -> a
sel = Vector a -> a
forall a. (Unbox a, Num a) => Vector a -> a
U.sum (Vector a -> a) -> Vector a -> a
forall a b. (a -> b) -> a -> b
$ (Measured -> a) -> Vector Measured -> Vector a
forall a. Unbox a => (Measured -> a) -> Vector Measured -> Vector a
measure Measured -> a
sel Vector Measured
meas'
  BenchmarkStats -> IO BenchmarkStats
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure BenchmarkStats
    { benchAvgTime :: Double
benchAvgTime = (Measured -> Double) -> Double
forall {a}. (Unbox a, Num a) => (Measured -> a) -> a
sumMeasure Measured -> Double
measTime Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len
    , benchAvgCpuTime :: Double
benchAvgCpuTime = (Measured -> Double) -> Double
forall {a}. (Unbox a, Num a) => (Measured -> a) -> a
sumMeasure Measured -> Double
measCpuTime Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len
    , benchAvgCycles :: Int64
benchAvgCycles = (Measured -> Int64) -> Int64
forall {a}. (Unbox a, Num a) => (Measured -> a) -> a
sumMeasure Measured -> Int64
measCycles Int64 -> Int64 -> Int64
forall a. Integral a => a -> a -> a
`div` Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len }