-- |
-- Module      : Foundation.Timing
-- License     : BSD-style
-- Maintainer  : Foundation maintainers
--
-- An implementation of a timing framework
--
{-# LANGUAGE CPP #-}
module Foundation.Timing
    ( Timing(..)
    , Measure(..)
    , stopWatch
    , measure
    ) where

import           Basement.Imports hiding (from)
import           Basement.From (from)
#if __GLASGOW_HASKELL__ < 802
import           Basement.Cast (cast)
#endif
import           Basement.Monad
-- import           Basement.UArray hiding (unsafeFreeze)
import           Basement.UArray.Mutable (MUArray)
import           Foundation.Collection
import           Foundation.Time.Types
import           Foundation.Numerical
import           Foundation.Time.Bindings
import           Control.Exception (evaluate)
import           System.Mem (performGC)
import           Data.Function (on)
import qualified GHC.Stats as GHC


data Timing = Timing
    { Timing -> NanoSeconds
timeDiff           :: !NanoSeconds
    , Timing -> Maybe Word64
timeBytesAllocated :: !(Maybe Word64)
    }

data Measure = Measure
    { Measure -> UArray NanoSeconds
measurements :: UArray NanoSeconds
    , Measure -> Word
iters        :: Word
    }

#if __GLASGOW_HASKELL__ >= 802
type GCStats = GHC.RTSStats

getGCStats :: IO (Maybe GCStats)
getGCStats :: IO (Maybe GCStats)
getGCStats = do
    Bool
r <- IO Bool
GHC.getRTSStatsEnabled
    if Bool
r then Maybe GCStats -> IO (Maybe GCStats)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe GCStats
forall a. Maybe a
Nothing else GCStats -> Maybe GCStats
forall a. a -> Maybe a
Just (GCStats -> Maybe GCStats) -> IO GCStats -> IO (Maybe GCStats)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO GCStats
GHC.getRTSStats

diffGC :: Maybe GHC.RTSStats -> Maybe GHC.RTSStats -> Maybe Word64
diffGC :: Maybe GCStats -> Maybe GCStats -> Maybe Word64
diffGC Maybe GCStats
gc2 Maybe GCStats
gc1 = ((-) (Word64 -> Word64 -> Word64)
-> (GCStats -> Word64) -> GCStats -> GCStats -> Word64
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` GCStats -> Word64
GHC.allocated_bytes) (GCStats -> GCStats -> Word64)
-> Maybe GCStats -> Maybe (GCStats -> Word64)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe GCStats
gc2 Maybe (GCStats -> Word64) -> Maybe GCStats -> Maybe Word64
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe GCStats
gc1
#else
type GCStats = GHC.GCStats

getGCStats :: IO (Maybe GCStats)
getGCStats = do
    r <- GHC.getGCStatsEnabled
    if r then pure Nothing else Just <$> GHC.getGCStats

diffGC :: Maybe GHC.GCStats -> Maybe GHC.GCStats -> Maybe Word64
diffGC gc2 gc1 = cast <$> (((-) `on` GHC.bytesAllocated) <$> gc2 <*> gc1)
#endif

-- | Simple one-time measurement of time & other metrics spent in a function
stopWatch :: (a -> b) -> a -> IO Timing
stopWatch :: (a -> b) -> a -> IO Timing
stopWatch a -> b
f !a
a = do
    IO ()
performGC
    Maybe GCStats
gc1 <- IO (Maybe GCStats)
getGCStats
    (b
_, NanoSeconds
ns) <- IO b -> IO (b, NanoSeconds)
forall a. IO a -> IO (a, NanoSeconds)
measuringNanoSeconds (b -> IO b
forall a. a -> IO a
evaluate (b -> IO b) -> b -> IO b
forall a b. (a -> b) -> a -> b
$ a -> b
f a
a)
    Maybe GCStats
gc2 <- IO (Maybe GCStats)
getGCStats
    Timing -> IO Timing
forall (m :: * -> *) a. Monad m => a -> m a
return (Timing -> IO Timing) -> Timing -> IO Timing
forall a b. (a -> b) -> a -> b
$ NanoSeconds -> Maybe Word64 -> Timing
Timing NanoSeconds
ns (Maybe GCStats -> Maybe GCStats -> Maybe Word64
diffGC Maybe GCStats
gc2 Maybe GCStats
gc1)

-- | In depth timing & other metrics analysis of a function
measure :: Word -> (a -> b) -> a -> IO Measure
measure :: Word -> (a -> b) -> a -> IO Measure
measure Word
nbIters a -> b
f a
a = do
    MUArray NanoSeconds RealWorld
d <- CountOf (MutableValue (MUArray NanoSeconds))
-> IO (MUArray NanoSeconds (PrimState IO))
forall (c :: * -> *) (prim :: * -> *).
(MutableCollection c, PrimMonad prim) =>
CountOf (MutableValue c) -> prim (c (PrimState prim))
mutNew (Word -> CountOf NanoSeconds
forall a b. From a b => a -> b
from Word
nbIters) :: IO (MUArray NanoSeconds (PrimState IO))
    MUArray NanoSeconds RealWorld -> Word -> IO ()
loop MUArray NanoSeconds RealWorld
d Word
0
    UArray NanoSeconds -> Word -> Measure
Measure (UArray NanoSeconds -> Word -> Measure)
-> IO (UArray NanoSeconds) -> IO (Word -> Measure)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MUArray NanoSeconds (PrimState IO)
-> IO (MutableFreezed (MUArray NanoSeconds))
forall (c :: * -> *) (prim :: * -> *).
(MutableCollection c, PrimMonad prim) =>
c (PrimState prim) -> prim (MutableFreezed c)
unsafeFreeze MUArray NanoSeconds RealWorld
MUArray NanoSeconds (PrimState IO)
d
            IO (Word -> Measure) -> IO Word -> IO Measure
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Word -> IO Word
forall (f :: * -> *) a. Applicative f => a -> f a
pure Word
nbIters
  where
    loop :: MUArray NanoSeconds RealWorld -> Word -> IO ()
loop MUArray NanoSeconds RealWorld
d !Word
i
        | Word
i Word -> Word -> Bool
forall a. Eq a => a -> a -> Bool
== Word
nbIters = () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        | Bool
otherwise    = do
            (b
_, NanoSeconds
r) <- IO b -> IO (b, NanoSeconds)
forall a. IO a -> IO (a, NanoSeconds)
measuringNanoSeconds (b -> IO b
forall a. a -> IO a
evaluate (b -> IO b) -> b -> IO b
forall a b. (a -> b) -> a -> b
$ a -> b
f a
a)
            MUArray NanoSeconds (PrimState IO)
-> MutableKey (MUArray NanoSeconds)
-> MutableValue (MUArray NanoSeconds)
-> IO ()
forall (c :: * -> *) (prim :: * -> *).
(MutableCollection c, PrimMonad prim) =>
c (PrimState prim) -> MutableKey c -> MutableValue c -> prim ()
mutUnsafeWrite MUArray NanoSeconds RealWorld
MUArray NanoSeconds (PrimState IO)
d (Word -> Offset NanoSeconds
forall a b. From a b => a -> b
from Word
i) MutableValue (MUArray NanoSeconds)
NanoSeconds
r
            MUArray NanoSeconds RealWorld -> Word -> IO ()
loop MUArray NanoSeconds RealWorld
d (Word
iWord -> Word -> Word
forall a. Additive a => a -> a -> a
+Word
1)