{-# LANGUAGE Trustworthy #-}
-- |
-- Module      : Criterion.Monad
-- Copyright   : (c) 2009 Neil Brown
--
-- License     : BSD-style
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : GHC
--
-- The environment in which most criterion code executes.
module Criterion.Monad
    (
      Criterion
    , withConfig
    , getGen
    , getOverhead
    ) where

import Control.Monad.Reader (asks, runReaderT)
import Control.Monad.Trans (liftIO)
import Control.Monad (when)
import Criterion.Measurement (measure, runBenchmark, secs)
import Criterion.Monad.Internal (Criterion(..), Crit(..))
import Criterion.Types hiding (measure)
import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import Statistics.Regression (olsRegress)
import System.Random.MWC (GenIO, createSystemRandom)
import qualified Data.Vector.Generic as G

-- | Run a 'Criterion' action with the given 'Config'.
withConfig :: Config -> Criterion a -> IO a
withConfig cfg (Criterion act) = do
  g <- newIORef Nothing
  o <- newIORef Nothing
  runReaderT act (Crit cfg g o)

-- | Return a random number generator, creating one if necessary.
--
-- This is not currently thread-safe, but in a harmless way (we might
-- call 'createSystemRandom' more than once if multiple threads race).
getGen :: Criterion GenIO
getGen = memoise gen createSystemRandom

-- | Return an estimate of the measurement overhead.
getOverhead :: Criterion Double
getOverhead = do
  verbose <- asks ((== Verbose) . verbosity)
  memoise overhead $ do
    (meas,_) <- runBenchmark (whnfIO $ measure (whnfIO $ return ()) 1) 1
    let metric get = G.convert . G.map get $ meas
    let o = G.head . fst $
            olsRegress [metric (fromIntegral . measIters)] (metric measTime)
    when verbose . liftIO $
      putStrLn $ "measurement overhead " ++ secs o
    return o

-- | Memoise the result of an 'IO' action.
--
-- This is not currently thread-safe, but hopefully in a harmless way.
-- We might call the given action more than once if multiple threads
-- race, so our caller's job is to write actions that can be run
-- multiple times safely.
memoise :: (Crit -> IORef (Maybe a)) -> IO a -> Criterion a
memoise ref generate = do
  r <- Criterion $ asks ref
  liftIO $ do
    mv <- readIORef r
    case mv of
      Just rv -> return rv
      Nothing -> do
        rv <- generate
        writeIORef r (Just rv)
        return rv
