-- | A module to collect aggregates on how much time is spent in a computation
--
-- Aggregates can be identified with a label that determines where the time of
-- each computation is accounted for.
--
-- Measures are collected only if the environment variable
-- @DEBUG_TIMESTATS_ENABLE@ is set to any value ahead of invoking any function
-- in this module.
--
module Debug.TimeStats
  ( -- * Measuring
    measureM
  , measurePure
    -- * Time stats manipulation
  , printTimeStats
  , hPrintTimeStats
  , reset
  , TimeStats(..)
  , collect
  , asText
  , scope
    -- * Not intended for direct use
    --
    -- | These definitions are not intended for instrumenting applications,
    -- but they can be handy to implement other measuring primitives.
    --
  , TimeStatsRef
  , lookupTimeStatsRef
  , updateTimeStatsRef
  ) where

import Control.Exception (evaluate)
import Control.Monad (forM, forM_, unless)
import Control.Monad.IO.Class (MonadIO(liftIO))
import Data.IORef
import Data.Map (Map)
import Data.Maybe (isJust)
import qualified Data.Map as Map
import Data.Text (Text)
import qualified Data.Text as Text
import qualified Data.Text.IO as Text
import Data.Word (Word64)
import Debug.TimeStats.Internal (formatIntWithSeparator)
import GHC.Clock (getMonotonicTimeNSec)
import Text.Printf (printf)
import System.Environment (lookupEnv)
import System.IO (Handle, stderr)
import System.IO.Unsafe (unsafePerformIO)


-- | Measure the time it takes to run the action.
--
-- Add the time to the stats of the given label and increase its count by one.
--
-- 'measureM' keeps the stats in a globally available store in order to minimize
-- the changes necessary when instrumenting a program. Otherwise a reference to
-- the store would need to be passed to every function that might invoke
-- functions that need this reference.
--
-- A time measure isn't collected if the given action fails with an exception.
-- This is a deliberate choice to demand less of the monad in which measures are
-- taken.
--
-- Time measures aren't collected either if the environment variable
-- @DEBUG_TIMESTATS_ENABLE@ isn't set the first time this function is
-- evaluated.
--
{-# INLINE measureM #-}
measureM :: MonadIO m => String -> m a -> m a
measureM :: forall (m :: * -> *) a. MonadIO m => String -> m a -> m a
measureM String
label =
    -- See the documentation of 'enabled'
    if Bool
enabled then do
          -- @ref@ is the reference to the stats associated to the label.
          -- See note [Looking up stats with unsafePerformIO]
      let ref :: TimeStatsRef
ref = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ String -> IO TimeStatsRef
lookupTimeStatsRef String
label
       in \m a
action -> forall (m :: * -> *) a. MonadIO m => TimeStatsRef -> m a -> m a
measureMWith TimeStatsRef
ref m a
action
    else
      forall a. a -> a
id

-- | Pure version of 'measureM'. Measures the time taken to reduce the given
-- value to head normal form.
--
-- 'measurePure' is a bit dangerous to use in contexts where there are monadic
-- computations. If 'measurePure' is applied to a monadic computation it
-- will measure the time of constructing the computation rather than the time
-- of executing it, and the typechecker won't catch the mistake. We try to
-- fence against it with a longer name.
{-# INLINE measurePure #-}
measurePure :: String -> a -> a
measurePure :: forall a. String -> a -> a
measurePure String
label =
    if Bool
enabled then
      forall a. IO a -> a
unsafePerformIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. MonadIO m => String -> m a -> m a
measureM String
label forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. a -> IO a
evaluate
    else
      forall a. a -> a
id

-- Note [Looking up stats with unsafePerformIO]
--
-- When calling 'measureM' we would like to save the trouble of looking the
-- stats to update on every invocation. Hence, we use unsafePerformIO, and
-- we ask to inline 'measureM'.
--
-- Most of the time 'measureM' should be called with a statically known label.
-- When inlining, GHC should notice this fact and move the lookup closure to
-- the top-level, thus performing it only once per invocation, and perhaps
-- only once per label for all 'measureM' calls in the same module.


-- | @True@ iff the environment variable @DEBUG_TIMESTATS_ENABLE@ is set to any
-- value
--
-- We assume the value of the environment variable doesn't change during the
-- lifetime of the program.
--
-- The purpose of making this a top-level value is to have all calls to
-- 'measureM' checking it only the first time. Thus we save the trouble of
-- looking up the environment variable repeteadly.
{-# NOINLINE enabled #-}
enabled :: Bool
enabled :: Bool
enabled = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ forall a. Maybe a -> Bool
isJust forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> IO (Maybe String)
lookupEnv String
"DEBUG_TIMESTATS_ENABLE"

-- | A unique global reference to the map associating labels to their
-- stats.
{-# NOINLINE labelStatsMapRef #-}
labelStatsMapRef :: IORef (Map String TimeStatsRef)
labelStatsMapRef :: IORef (Map String TimeStatsRef)
labelStatsMapRef = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ forall a. a -> IO (IORef a)
newIORef forall k a. Map k a
Map.empty

-- | Set all statistics to initial values.
reset :: MonadIO m => m ()
reset :: forall (m :: * -> *). MonadIO m => m ()
reset = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$
    if Bool
enabled then do
      Map String TimeStatsRef
m <- forall a. IORef a -> IO a
readIORef IORef (Map String TimeStatsRef)
labelStatsMapRef
      forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (forall k a. Map k a -> [a]
Map.elems Map String TimeStatsRef
m) forall a b. (a -> b) -> a -> b
$ \(TimeStatsRef IORef TimeStats
ref) ->
        forall a. IORef a -> a -> IO ()
writeIORef IORef TimeStats
ref TimeStats
initialTimeStats
     else
      forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | Run an action by previously reseting all stats to initial values
-- and printing them afterwards.
scope :: MonadIO m => m a -> m a
scope :: forall (m :: * -> *) a. MonadIO m => m a -> m a
scope =
    if Bool
enabled then
      \m a
m -> do
        forall (m :: * -> *). MonadIO m => m ()
reset
        a
a <- m a
m
        forall (m :: * -> *). MonadIO m => Handle -> m ()
hPrintTimeStats Handle
stderr
        forall (m :: * -> *) a. Monad m => a -> m a
return a
a
     else
      forall a. a -> a
id

-- | Looks up the stats of a label. If no stats are found for the label,
-- a new TimeStatsRef is created with initial values.
--
lookupTimeStatsRef :: String -> IO TimeStatsRef
lookupTimeStatsRef :: String -> IO TimeStatsRef
lookupTimeStatsRef String
label = do
    TimeStatsRef
r0 <- forall (m :: * -> *). MonadIO m => m TimeStatsRef
newTimeStatsRef
    forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef (Map String TimeStatsRef)
labelStatsMapRef forall a b. (a -> b) -> a -> b
$ \Map String TimeStatsRef
m ->
      case forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup String
label Map String TimeStatsRef
m of
        Maybe TimeStatsRef
Nothing -> (forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert String
label TimeStatsRef
r0 Map String TimeStatsRef
m, TimeStatsRef
r0)
        Just TimeStatsRef
r -> (Map String TimeStatsRef
m, TimeStatsRef
r)

-- | Yields the labels and the stats collected thus far.
collect :: MonadIO m => m [(String, TimeStats)]
collect :: forall (m :: * -> *). MonadIO m => m [(String, TimeStats)]
collect = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    Map String TimeStatsRef
m <- forall a. IORef a -> IO a
readIORef IORef (Map String TimeStatsRef)
labelStatsMapRef
    forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (forall k a. Map k a -> [(k, a)]
Map.toList Map String TimeStatsRef
m) forall a b. (a -> b) -> a -> b
$ \(String
label, TimeStatsRef IORef TimeStats
ref) ->
      (,) String
label forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. IORef a -> IO a
readIORef IORef TimeStats
ref

-- | Prints the time stats to the given handle.
hPrintTimeStats :: MonadIO m => Handle -> m ()
hPrintTimeStats :: forall (m :: * -> *). MonadIO m => Handle -> m ()
hPrintTimeStats Handle
h = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    [(String, TimeStats)]
xs <- forall (m :: * -> *). MonadIO m => m [(String, TimeStats)]
collect
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(String, TimeStats)]
xs) forall a b. (a -> b) -> a -> b
$
      Handle -> Text -> IO ()
Text.hPutStrLn Handle
h ([(String, TimeStats)] -> Text
asText [(String, TimeStats)]
xs)

-- | Prints the time stats to stderr.
printTimeStats :: MonadIO m => m ()
printTimeStats :: forall (m :: * -> *). MonadIO m => m ()
printTimeStats = forall (m :: * -> *). MonadIO m => Handle -> m ()
hPrintTimeStats Handle
stderr

-- | Renders the given time stats in a tabular format
asText :: [(String, TimeStats)] -> Text
asText :: [(String, TimeStats)] -> Text
asText [(String, TimeStats)]
stats =
    let ([String]
lbls, [TimeStats]
timestats) = forall a b. [(a, b)] -> ([a], [b])
unzip [(String, TimeStats)]
stats
        ([String]
times, [String]
counts) = forall a b. [(a, b)] -> ([a], [b])
unzip forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map TimeStats -> (String, String)
formatTimeStats [TimeStats]
timestats
        widthLbls :: Int
widthLbls = forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall (t :: * -> *) a. Foldable t => t a -> Int
length [String]
lbls
        widthTimes :: Int
widthTimes = forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall (t :: * -> *) a. Foldable t => t a -> Int
length [String]
times
        widthCounts :: Int
widthCounts = forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall (t :: * -> *) a. Foldable t => t a -> Int
length [String]
counts
     in [Text] -> Text
Text.unlines forall a b. (a -> b) -> a -> b
$
        forall a b. (a -> b) -> [a] -> [b]
map (String -> Text
Text.pack forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Int -> (String, String, String) -> String
printStat Int
widthLbls Int
widthTimes Int
widthCounts) forall a b. (a -> b) -> a -> b
$
        forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [String]
lbls [String]
times [String]
counts
  where
    formatTimeStats :: TimeStats -> (String, String)
    formatTimeStats :: TimeStats -> (String, String)
formatTimeStats TimeStats
t =
      ( forall r. PrintfType r => String -> r
printf String
"%.3f" (forall a b. (Integral a, Num b) => a -> b
fromIntegral (TimeStats -> Word64
timeStat TimeStats
t) forall a. Fractional a => a -> a -> a
/ Double
1e9 :: Double)
      , Char -> Int -> String -> String
formatIntWithSeparator Char
'_' (TimeStats -> Int
countStat TimeStats
t) String
""
      )

    -- At the time of this writing printf can't render to 'Text'.
    printStat :: Int -> Int -> Int -> (String, String, String) -> String
    printStat :: Int -> Int -> Int -> (String, String, String) -> String
printStat Int
widthLbls Int
widthTimes Int
widthCounts (String
label, String
time, String
count) =
      let fmt :: String
fmt = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
            [ String
"%", forall a. Show a => a -> String
show Int
widthLbls
            , String
"s: %", forall a. Show a => a -> String
show Int
widthTimes
            , String
"ss  count: %", forall a. Show a => a -> String
show Int
widthCounts, String
"s"
            ]
       in forall r. PrintfType r => String -> r
printf String
fmt (String -> Text
Text.pack String
label) String
time String
count

---------------------
-- TimeStats
---------------------

-- | A reference to a 'TimeStats' value
newtype TimeStatsRef = TimeStatsRef (IORef TimeStats)

-- | Reports how much time (in nanoseconds) the invocations to 'measureM' took
-- for a given label and how many times it was invoked on a given label.
data TimeStats = TimeStats
    { TimeStats -> Word64
timeStat :: {-# UNPACK #-} !Word64
    , TimeStats -> Int
countStat :: {-# UNPACK #-} !Int
    }
  deriving Int -> TimeStats -> String -> String
[TimeStats] -> String -> String
TimeStats -> String
forall a.
(Int -> a -> String -> String)
-> (a -> String) -> ([a] -> String -> String) -> Show a
showList :: [TimeStats] -> String -> String
$cshowList :: [TimeStats] -> String -> String
show :: TimeStats -> String
$cshow :: TimeStats -> String
showsPrec :: Int -> TimeStats -> String -> String
$cshowsPrec :: Int -> TimeStats -> String -> String
Show

-- | Measured time is 0 and call count is 0.
initialTimeStats :: TimeStats
initialTimeStats :: TimeStats
initialTimeStats = Word64 -> Int -> TimeStats
TimeStats Word64
0 Int
0

-- | Creates a reference to time stats with intial values
newTimeStatsRef :: MonadIO m => m TimeStatsRef
newTimeStatsRef :: forall (m :: * -> *). MonadIO m => m TimeStatsRef
newTimeStatsRef = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ IORef TimeStats -> TimeStatsRef
TimeStatsRef forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. a -> IO (IORef a)
newIORef TimeStats
initialTimeStats

-- | Measure the time it takes to run the given action and update with it
-- the given reference to time stats.
measureMWith :: MonadIO m => TimeStatsRef -> m a -> m a
measureMWith :: forall (m :: * -> *) a. MonadIO m => TimeStatsRef -> m a -> m a
measureMWith TimeStatsRef
tref m a
m = do
    Word64
t0 <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO Word64
getMonotonicTimeNSec
    a
a <- m a
m
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
      Word64
tf <- IO Word64
getMonotonicTimeNSec
      TimeStatsRef -> (TimeStats -> TimeStats) -> IO ()
updateTimeStatsRef TimeStatsRef
tref forall a b. (a -> b) -> a -> b
$ \TimeStats
st ->
        TimeStats
st
          { timeStat :: Word64
timeStat = (Word64
tf forall a. Num a => a -> a -> a
- Word64
t0) forall a. Num a => a -> a -> a
+ TimeStats -> Word64
timeStat TimeStats
st
          , countStat :: Int
countStat = Int
1 forall a. Num a => a -> a -> a
+ TimeStats -> Int
countStat TimeStats
st
          }
    forall (m :: * -> *) a. Monad m => a -> m a
return a
a

-- | Updates the TimeStats in a TimeStatsRef
updateTimeStatsRef :: TimeStatsRef -> (TimeStats -> TimeStats) -> IO ()
updateTimeStatsRef :: TimeStatsRef -> (TimeStats -> TimeStats) -> IO ()
updateTimeStatsRef (TimeStatsRef IORef TimeStats
ref) TimeStats -> TimeStats
f =
    forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef TimeStats
ref forall a b. (a -> b) -> a -> b
$ \TimeStats
st -> (TimeStats -> TimeStats
f TimeStats
st, ())