{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}

-- | Utilities for tracing to the eventlog.
--
-- These functions should be preferred over "Debug.Trace".
module Mikan.Utils.Trace
  ( -- * Eventlog tracing
    -- $eventlogTracing
    traceEventIO
    -- * Execution phase markers
    -- $eventlogMarkers
  , traceMarkerIO
    -- * Era profiling
    -- $eraProfiling
  , getUserEra
  , setUserEra
  , incrementUserEra
  , incrementUserEra_
  ) where

import Control.Monad.IO.Class

import Data.ByteString qualified as B
import Data.Functor ((<&>))
import Data.Text (Text)
import Data.Text.Encoding qualified as T

import GHC.Exts (Ptr(..), traceEvent#, traceMarker#)
import GHC.IO (IO(..))
import GHC.Profiling.Eras qualified as Eras
import GHC.RTS.Flags qualified as RTS

import Mikan.Utils.Monad

import System.IO.Unsafe

--------------------------------------------------------------------------------
-- Eventlog tracing

{-# NOINLINE userTracingEnabled #-}
-- | Are user events enabled?
userTracingEnabled :: Bool
userTracingEnabled :: Bool
userTracingEnabled =
  -- This @unsafeDupablePerformIO@ is safe, as we will not
  -- change the eventlog flags during runtime. Moreover, we
  -- do not care if this gets run multiple times across different
  -- threads, as this just a read and nothing else is writing to
  -- the flags.
  IO Bool -> Bool
forall a. IO a -> a
unsafeDupablePerformIO (TraceFlags -> Bool
RTS.user (TraceFlags -> Bool) -> IO TraceFlags -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO TraceFlags
RTS.getTraceFlags)

-- $eventlogTracing
--
-- Eventlog tracing is a performance profiling system. These functions emit
-- extra events into the eventlog. In combination with eventlog profiling
-- tools these functions can be used for monitoring execution and
-- investigating performance problems.
--
-- The easiest way to enable the eventlog is to pass @+RTS -l-au -RTS@ to @mikan@.
-- The [GHC Manual](https://ghc.gitlab.haskell.org/ghc/doc/users_guide/runtime_control.html#rts-eventlog)
-- provides documentation on further options.

{-# INLINE traceEventIO #-}
-- | \(\mathcal{O}(n)\). Log a @UserMessage@ event to the eventlog.
-- The input 'Text' is marshalled to a null-terminated 'Foreign.C.String.CString'.
--
-- This function does not force the 'Text' if the eventlog is disabled.
--
-- Precondition: the string should be shorter than \(2^{16}\) bytes.
traceEventIO :: (MonadIO m) => Text -> m ()
traceEventIO :: forall (m :: * -> *). MonadIO m => Text -> m ()
traceEventIO Text
txt = 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
$
  Bool -> IO () -> IO ()
forall b (m :: * -> *). (IsBool b, Monad m) => b -> m () -> m ()
when Bool
userTracingEnabled (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
  ByteString -> (CString -> IO ()) -> IO ()
forall a. ByteString -> (CString -> IO a) -> IO a
B.useAsCString (Text -> ByteString
T.encodeUtf8 Text
txt) \(Ptr Addr#
ptr) -> (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO \State# RealWorld
s ->
    case Addr# -> State# RealWorld -> State# RealWorld
forall d. Addr# -> State# d -> State# d
traceEvent# Addr#
ptr State# RealWorld
s of
      State# RealWorld
s' -> (# State# RealWorld
s', () #)

-- $eventlogMarkers
--
-- When looking at a profile for the execution of a program we often want to
-- be able to mark certain points or phases in the execution and see that
-- visually in the profile.
--
-- For example, a program might have several distinct phases with different
-- performance or resource behaviour in each phase. To properly interpret the
-- profile graph we really want to see when each phase starts and ends.
--
-- Markers let us do this: we can annotate the program to emit a marker at
-- an appropriate point during execution and then see that in a profile.

{-# INLINE traceMarkerIO #-}
-- | \(\mathcal{O}(n)\). Emit a marker to the eventlog.
-- The input 'Text' is marshalled to a null-terminated 'Foreign.C.String.CString'.
--
-- This function does not force the 'Text' if the eventlog is disabled.
--
-- Precondition: the string should be shorter than \(2^{16}\) bytes.
traceMarkerIO :: (MonadIO m) => Text -> m ()
traceMarkerIO :: forall (m :: * -> *). MonadIO m => Text -> m ()
traceMarkerIO Text
txt = 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
$
  Bool -> IO () -> IO ()
forall b (m :: * -> *). (IsBool b, Monad m) => b -> m () -> m ()
when Bool
userTracingEnabled (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
  ByteString -> (CString -> IO ()) -> IO ()
forall a. ByteString -> (CString -> IO a) -> IO a
B.useAsCString (Text -> ByteString
T.encodeUtf8 Text
txt) \(Ptr Addr#
ptr) -> (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO \State# RealWorld
s ->
    case Addr# -> State# RealWorld -> State# RealWorld
forall d. Addr# -> State# d -> State# d
traceMarker# Addr#
ptr State# RealWorld
s of
      State# RealWorld
s' -> (# State# RealWorld
s', () #)

--------------------------------------------------------------------------------
-- Era profiling

-- $eraProfiling
--
-- Era profiling lets us mark each closure with an /era/ that tracks when
-- the closure was allocated. It can be enabled by passing @+RTS -he -RTS@ to
-- a profiled build of @mikan@.
--
-- There are two ways to control the current era:
--
-- 1. Automatically via @+RTS --automatic-era-increment -RTS@, which increments
--    the era on every major GC.
-- 2. Manually via 'setUserEra', 'incrementUserEra_', and 'incrementUserEra'.

{-# NOINLINE heapProfilingEnabled #-}
-- | Is heap profiling enabled?
heapProfilingEnabled :: Bool
heapProfilingEnabled :: Bool
heapProfilingEnabled =
  -- This @unsafeDupablePerformIO@ is safe, as we will not
  -- enable heap profiling during runtime. Moreover, we
  -- do not care if this gets run multiple times across different
  -- threads, as this just a read and nothing else is writing to
  -- the flags.
  IO Bool -> Bool
forall a. IO a -> a
unsafeDupablePerformIO (IO Bool -> Bool) -> IO Bool -> Bool
forall a b. (a -> b) -> a -> b
$ IO ProfFlags
RTS.getProfFlags IO ProfFlags -> (ProfFlags -> Bool) -> IO Bool
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \case
    RTS.ProfFlags { doHeapProfile :: ProfFlags -> DoHeapProfile
RTS.doHeapProfile = DoHeapProfile
RTS.NoHeapProfiling } -> Bool
False
    ProfFlags
_ -> Bool
True

{-# INLINE getUserEra #-}
-- | Get the user era count if heap profiling is enabled.
--
-- If heap profiling is disabled, this function will always return @0@.
getUserEra :: (MonadIO m) => m Word
getUserEra :: forall (m :: * -> *). MonadIO m => m Word
getUserEra = IO Word -> m Word
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Word -> m Word) -> IO Word -> m Word
forall a b. (a -> b) -> a -> b
$
  if Bool
heapProfilingEnabled then
    IO Word
Eras.getUserEra
  else
    Word -> IO Word
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Word
0

{-# INLINE setUserEra #-}
-- | Set the user era count if heap profiling is enabled.
--
-- If heap profiling is disabled, this function is a no-op.
setUserEra :: (MonadIO m) => Word -> m ()
setUserEra :: forall (m :: * -> *). MonadIO m => Word -> m ()
setUserEra Word
n = 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
$
  Bool -> IO () -> IO ()
forall b (m :: * -> *). (IsBool b, Monad m) => b -> m () -> m ()
when Bool
heapProfilingEnabled (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
    Word -> IO ()
Eras.setUserEra Word
n

{-# INLINE incrementUserEra_ #-}
-- | Increment the user era count if heap profiling is enabled.
--
-- If heap profiling is disabled, this function will always return @0@.
incrementUserEra_ :: (MonadIO m) => Word -> m ()
incrementUserEra_ :: forall (m :: * -> *). MonadIO m => Word -> m ()
incrementUserEra_ Word
n = 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
$
  Bool -> IO () -> IO ()
forall b (m :: * -> *). (IsBool b, Monad m) => b -> m () -> m ()
when Bool
heapProfilingEnabled (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
    IO Word -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Word -> IO ()) -> IO Word -> IO ()
forall a b. (a -> b) -> a -> b
$ Word -> IO Word
Eras.incrementUserEra Word
n

{-# INLINE incrementUserEra #-}
-- | Increment the user era count if heap profiling is enabled, and
-- return the new era.
--
-- If heap profiling is disabled, this function will always return @0@.
incrementUserEra :: (MonadIO m) => Word -> m Word
incrementUserEra :: forall (m :: * -> *). MonadIO m => Word -> m Word
incrementUserEra Word
n = IO Word -> m Word
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Word -> m Word) -> IO Word -> m Word
forall a b. (a -> b) -> a -> b
$
  if Bool
heapProfilingEnabled then
    Word -> IO Word
Eras.incrementUserEra Word
n
  else
    Word -> IO Word
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Word
0