Unsafe.hs

-- | A module with time measuring primitives that might not work in all monads
-- that building allows.
--
-- 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.Unsafe
  ( -- * Measuring
    unsafeMeasureM
  ) where

import Debug.TimeStats
         ( TimeStats(..)
         , TimeStatsRef
         , enabled
         , lookupTimeStatsRef
         , updateTimeStatsRef
         )
import GHC.Clock (getMonotonicTimeNSec)
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.
--
-- This function relies on a hack to perform IO in any monad, which does not
-- always work. In particular, we can expect it to fail in monads where
--
-- > (m >>= \_ -> undefined) == undefined -- for some computation m
--
-- An example of such a monad is the list monad
--
-- > ([()] >>= \_ -> undefined) == undefined
--
-- Another example is the @Control.Monad.Free.Free IO@.
--
-- > (Control.Monad.Free.Pure () >>= \_ -> undefined) == undefined
--
-- But it seems to work on @IO@ or @ReaderT IO@.
--
-- > seq (print () >>= \_ -> undefined) () == ()
--
-- Also, monads that run the continuation of bind multiple times might only
-- have accounted the time to run the first time only.
--
{-# INLINE unsafeMeasureM #-}
unsafeMeasureM :: Monad m => String -> m a -> m a
unsafeMeasureM label =
    -- See the documentation of 'enabled'
    if enabled then do
          -- @ref@ is the reference to the stats associated to the label.
          -- See note [Looking up stats with unsafePerformIO]
      let ref = unsafePerformIO $ lookupTimeStatsRef label
       in \action -> measureMWith ref action
    else
      id

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

---------------------
-- intersperseIOinM
---------------------

-- | Hack to intersperse IO actions into any monad
intersperseIOinM :: Monad m => IO a -> m a
intersperseIOinM m = do
    -- The fictitious state is only used to force @unsafePerformIO@
    -- to run @m@ every time @intersperseIOinM m@ is evaluated.
    s <- getStateM
    case unsafePerformIO $ (,) s <$> m of
      (_, r) -> pure r
  where
    -- We mark this function as NOINLINE to ensure the compiler cannot reason
    -- by unfolding that two calls of @getStateM@ yield the same value.
    {-# NOINLINE getStateM #-}
    getStateM = pure True