{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RankNTypes #-}
-- | 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
  , enabled
  , lookupTimeStatsRef
  , measureMWithLiftIO
  , 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.
--
-- Monads that run the continuation of bind multiple times might only
-- have accounted the time to run the first time only.
--
{-# INLINE measureM #-}
measureM :: MonadIO m => String -> m a -> m a
measureM :: forall (m :: * -> *) a. MonadIO m => String -> m a -> m a
measureM String
label = String -> (forall b. IO b -> m b) -> m a -> m a
forall (m :: * -> *) a.
Monad m =>
String -> (forall b. IO b -> m b) -> m a -> m a
measureMWithLiftIO String
label IO b -> m b
forall b. IO b -> m b
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO

-- | Like 'measureM' but allows to change the function to lift IO into the
-- monad.
{-# INLINE measureMWithLiftIO #-}
measureMWithLiftIO :: Monad m => String -> (forall b. IO b -> m b) -> m a -> m a
measureMWithLiftIO :: forall (m :: * -> *) a.
Monad m =>
String -> (forall b. IO b -> m b) -> m a -> m a
measureMWithLiftIO String
label forall b. IO b -> m b
lift =
    -- 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 = IO TimeStatsRef -> TimeStatsRef
forall a. IO a -> a
unsafePerformIO (IO TimeStatsRef -> TimeStatsRef)
-> IO TimeStatsRef -> TimeStatsRef
forall a b. (a -> b) -> a -> b
$ String -> IO TimeStatsRef
lookupTimeStatsRef String
label
       in \m a
action -> (forall b. IO b -> m b) -> TimeStatsRef -> m a -> m a
forall (m :: * -> *) a.
Monad m =>
(forall b. IO b -> m b) -> TimeStatsRef -> m a -> m a
measureMWithRef IO b -> m b
forall b. IO b -> m b
lift TimeStatsRef
ref m a
action
    else
      m a -> m a
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
      IO a -> a
forall a. IO a -> a
unsafePerformIO (IO a -> a) -> (a -> IO a) -> a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> IO a -> IO a
forall (m :: * -> *) a. MonadIO m => String -> m a -> m a
measureM String
label (IO a -> IO a) -> (a -> IO a) -> a -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> IO a
forall a. a -> IO a
evaluate
    else
      a -> a
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 = IO Bool -> Bool
forall a. IO a -> a
unsafePerformIO (IO Bool -> Bool) -> IO Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Maybe String -> Bool
forall a. Maybe a -> Bool
isJust (Maybe String -> Bool) -> IO (Maybe String) -> IO Bool
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 = IO (IORef (Map String TimeStatsRef))
-> IORef (Map String TimeStatsRef)
forall a. IO a -> a
unsafePerformIO (IO (IORef (Map String TimeStatsRef))
 -> IORef (Map String TimeStatsRef))
-> IO (IORef (Map String TimeStatsRef))
-> IORef (Map String TimeStatsRef)
forall a b. (a -> b) -> a -> b
$ Map String TimeStatsRef -> IO (IORef (Map String TimeStatsRef))
forall a. a -> IO (IORef a)
newIORef Map String TimeStatsRef
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 = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$
    if Bool
enabled then do
      m <- IORef (Map String TimeStatsRef) -> IO (Map String TimeStatsRef)
forall a. IORef a -> IO a
readIORef IORef (Map String TimeStatsRef)
labelStatsMapRef
      forM_ (Map.elems m) $ \(TimeStatsRef IORef TimeStats
ref) ->
        IORef TimeStats -> TimeStats -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef TimeStats
ref TimeStats
initialTimeStats
     else
      () -> IO ()
forall a. a -> IO a
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
        m ()
forall (m :: * -> *). MonadIO m => m ()
reset
        a <- m a
m
        hPrintTimeStats stderr
        return a
     else
      m a -> m a
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
    r0 <- IO TimeStatsRef
forall (m :: * -> *). MonadIO m => m TimeStatsRef
newTimeStatsRef
    atomicModifyIORef labelStatsMapRef $ \Map String TimeStatsRef
m ->
      case String -> Map String TimeStatsRef -> Maybe TimeStatsRef
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup String
label Map String TimeStatsRef
m of
        Maybe TimeStatsRef
Nothing -> (String
-> TimeStatsRef
-> Map String TimeStatsRef
-> Map String TimeStatsRef
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 = IO [(String, TimeStats)] -> m [(String, TimeStats)]
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO [(String, TimeStats)] -> m [(String, TimeStats)])
-> IO [(String, TimeStats)] -> m [(String, TimeStats)]
forall a b. (a -> b) -> a -> b
$ do
    m <- IORef (Map String TimeStatsRef) -> IO (Map String TimeStatsRef)
forall a. IORef a -> IO a
readIORef IORef (Map String TimeStatsRef)
labelStatsMapRef
    forM (Map.toList m) $ \(String
label, TimeStatsRef IORef TimeStats
ref) ->
      (,) String
label (TimeStats -> (String, TimeStats))
-> IO TimeStats -> IO (String, TimeStats)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef TimeStats -> IO TimeStats
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 = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    xs <- IO [(String, TimeStats)]
forall (m :: * -> *). MonadIO m => m [(String, TimeStats)]
collect
    unless (null xs) $
      Text.hPutStrLn h (asText xs)

-- | Prints the time stats to stderr.
printTimeStats :: MonadIO m => m ()
printTimeStats :: forall (m :: * -> *). MonadIO m => m ()
printTimeStats = Handle -> m ()
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) = [(String, TimeStats)] -> ([String], [TimeStats])
forall a b. [(a, b)] -> ([a], [b])
unzip [(String, TimeStats)]
stats
        ([String]
times, [String]
counts) = [(String, String)] -> ([String], [String])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(String, String)] -> ([String], [String]))
-> [(String, String)] -> ([String], [String])
forall a b. (a -> b) -> a -> b
$ (TimeStats -> (String, String))
-> [TimeStats] -> [(String, String)]
forall a b. (a -> b) -> [a] -> [b]
map TimeStats -> (String, String)
formatTimeStats [TimeStats]
timestats
        widthLbls :: Int
widthLbls = [Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ (String -> Int) -> [String] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [String]
lbls
        widthTimes :: Int
widthTimes = [Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ (String -> Int) -> [String] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [String]
times
        widthCounts :: Int
widthCounts = [Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ (String -> Int) -> [String] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [String]
counts
     in [Text] -> Text
Text.unlines ([Text] -> Text) -> [Text] -> Text
forall a b. (a -> b) -> a -> b
$
        ((String, String, String) -> Text)
-> [(String, String, String)] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (String -> Text
Text.pack (String -> Text)
-> ((String, String, String) -> String)
-> (String, String, String)
-> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Int -> (String, String, String) -> String
printStat Int
widthLbls Int
widthTimes Int
widthCounts) ([(String, String, String)] -> [Text])
-> [(String, String, String)] -> [Text]
forall a b. (a -> b) -> a -> b
$
        [String] -> [String] -> [String] -> [(String, String, String)]
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 =
      ( String -> Double -> String
forall r. PrintfType r => String -> r
printf String
"%.3f" (Word64 -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (TimeStats -> Word64
timeStat TimeStats
t) Double -> Double -> Double
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 = [String] -> String
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
            [ String
"%", Int -> String
forall a. Show a => a -> String
show Int
widthLbls
            , String
"s: %", Int -> String
forall a. Show a => a -> String
show Int
widthTimes
            , String
"ss  count: %", Int -> String
forall a. Show a => a -> String
show Int
widthCounts, String
"s"
            ]
       in String -> Text -> String -> String -> String
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
(Int -> TimeStats -> String -> String)
-> (TimeStats -> String)
-> ([TimeStats] -> String -> String)
-> Show TimeStats
forall a.
(Int -> a -> String -> String)
-> (a -> String) -> ([a] -> String -> String) -> Show a
$cshowsPrec :: Int -> TimeStats -> String -> String
showsPrec :: Int -> TimeStats -> String -> String
$cshow :: TimeStats -> String
show :: TimeStats -> String
$cshowList :: [TimeStats] -> String -> String
showList :: [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 = IO TimeStatsRef -> m TimeStatsRef
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO TimeStatsRef -> m TimeStatsRef)
-> IO TimeStatsRef -> m TimeStatsRef
forall a b. (a -> b) -> a -> b
$ IORef TimeStats -> TimeStatsRef
TimeStatsRef (IORef TimeStats -> TimeStatsRef)
-> IO (IORef TimeStats) -> IO TimeStatsRef
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TimeStats -> IO (IORef TimeStats)
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, using the given IO lifting function.
measureMWithRef :: Monad m => (forall b. IO b -> m b) -> TimeStatsRef -> m a -> m a
measureMWithRef :: forall (m :: * -> *) a.
Monad m =>
(forall b. IO b -> m b) -> TimeStatsRef -> m a -> m a
measureMWithRef forall b. IO b -> m b
lift TimeStatsRef
tref m a
m = do
    !t0 <- IO Word64 -> m Word64
forall b. IO b -> m b
lift IO Word64
getMonotonicTimeNSec
    a <- m
    lift $ do
      !tf <- getMonotonicTimeNSec
      updateTimeStatsRef tref $ \TimeStats
st ->
        TimeStats
st
          { timeStat = (tf - t0) + timeStat st
          , countStat = 1 + countStat st
          }
    return 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 =
    IORef TimeStats -> (TimeStats -> (TimeStats, ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef TimeStats
ref ((TimeStats -> (TimeStats, ())) -> IO ())
-> (TimeStats -> (TimeStats, ())) -> IO ()
forall a b. (a -> b) -> a -> b
$ \TimeStats
st -> (TimeStats -> TimeStats
f TimeStats
st, ())