{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE MonoLocalBinds #-}

module Test.Sandwich.TestTimer where

import Control.Concurrent
import Control.Exception.Safe
import Control.Monad.IO.Class
import Control.Monad.Reader
import Control.Monad.Trans.State
import qualified Data.Aeson as A
import qualified Data.ByteString.Lazy as BL
import qualified Data.List as L
import qualified Data.Sequence as S
import Data.String.Interpolate
import qualified Data.Text as T
import qualified Data.Text.IO as T
import Data.Time.Clock.POSIX
import Lens.Micro
import System.Directory
import System.FilePath
import System.IO
import Test.Sandwich.Types.RunTree
import Test.Sandwich.Types.Spec
import Test.Sandwich.Types.TestTimer
import Test.Sandwich.Util (whenJust)


type EventName = T.Text
type ProfileName = T.Text

-- * User functions

-- | Time a given action with a given event name. This name will be the "stack frame" of the given action in the profiling results. This function will use the current timing profile name.
timeAction :: (MonadMask m, MonadIO m, MonadReader context m, HasBaseContext context, HasTestTimer context) => EventName -> m a -> m a
timeAction :: forall (m :: * -> *) context a.
(MonadMask m, MonadIO m, MonadReader context m,
 HasBaseContext context, HasTestTimer context) =>
EventName -> m a -> m a
timeAction EventName
eventName m a
action = do
  TestTimer
tt <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall context. HasTestTimer context => context -> TestTimer
getTestTimer
  BaseContext {EventName
baseContextTestTimerProfile :: BaseContext -> EventName
baseContextTestTimerProfile :: EventName
baseContextTestTimerProfile} <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a. HasBaseContext a => a -> BaseContext
getBaseContext
  forall (m :: * -> *) a.
(MonadMask m, MonadIO m) =>
TestTimer -> EventName -> EventName -> m a -> m a
timeAction' TestTimer
tt EventName
baseContextTestTimerProfile EventName
eventName m a
action

-- | Time a given action with a given profile name and event name. Use when you want to manually specify the profile name.
timeActionByProfile :: (MonadMask m, MonadIO m, MonadReader context m, HasTestTimer context) => ProfileName -> EventName -> m a -> m a
timeActionByProfile :: forall (m :: * -> *) context a.
(MonadMask m, MonadIO m, MonadReader context m,
 HasTestTimer context) =>
EventName -> EventName -> m a -> m a
timeActionByProfile EventName
profileName EventName
eventName m a
action = do
  TestTimer
tt <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall context. HasTestTimer context => context -> TestTimer
getTestTimer
  forall (m :: * -> *) a.
(MonadMask m, MonadIO m) =>
TestTimer -> EventName -> EventName -> m a -> m a
timeAction' TestTimer
tt EventName
profileName EventName
eventName m a
action

-- | Introduce a new timing profile name.
withTimingProfile :: (Monad m) => ProfileName -> SpecFree (LabelValue "testTimerProfile" TestTimerProfile :> context) m () -> SpecFree context m ()
withTimingProfile :: forall (m :: * -> *) context.
Monad m =>
EventName
-> SpecFree
     (LabelValue "testTimerProfile" TestTimerProfile :> context) m ()
-> SpecFree context m ()
withTimingProfile EventName
name = forall intro (l :: Symbol) context (m :: * -> *).
(HasCallStack, Typeable intro) =>
NodeOptions
-> String
-> Label l intro
-> ExampleT context m intro
-> (intro -> ExampleT context m ())
-> SpecFree (LabelValue l intro :> context) m ()
-> SpecFree context m ()
introduce' NodeOptions
timingNodeOptions [i|Switch test timer profile to '#{name}'|] Label "testTimerProfile" TestTimerProfile
testTimerProfile (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ EventName -> TestTimerProfile
TestTimerProfile EventName
name) (\TestTimerProfile
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return ())

-- | Introduce a new timing profile name dynamically. The given 'ExampleT' should come up with the name and return it.
withTimingProfile' :: (Monad m) => ExampleT context m ProfileName -> SpecFree (LabelValue "testTimerProfile" TestTimerProfile :> context) m () -> SpecFree context m ()
withTimingProfile' :: forall (m :: * -> *) context.
Monad m =>
ExampleT context m EventName
-> SpecFree
     (LabelValue "testTimerProfile" TestTimerProfile :> context) m ()
-> SpecFree context m ()
withTimingProfile' ExampleT context m EventName
getName = forall intro (l :: Symbol) context (m :: * -> *).
(HasCallStack, Typeable intro) =>
NodeOptions
-> String
-> Label l intro
-> ExampleT context m intro
-> (intro -> ExampleT context m ())
-> SpecFree (LabelValue l intro :> context) m ()
-> SpecFree context m ()
introduce' NodeOptions
timingNodeOptions [i|Switch test timer profile to dynamic value|] Label "testTimerProfile" TestTimerProfile
testTimerProfile (EventName -> TestTimerProfile
TestTimerProfile forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExampleT context m EventName
getName) (\TestTimerProfile
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return ())

-- * Core

timingNodeOptions :: NodeOptions
timingNodeOptions :: NodeOptions
timingNodeOptions = NodeOptions
defaultNodeOptions { nodeOptionsRecordTime :: Bool
nodeOptionsRecordTime = Bool
False
                                       , nodeOptionsCreateFolder :: Bool
nodeOptionsCreateFolder = Bool
False
                                       , nodeOptionsVisibilityThreshold :: Int
nodeOptionsVisibilityThreshold = Int
systemVisibilityThreshold }

newSpeedScopeTestTimer :: FilePath -> Bool -> IO TestTimer
newSpeedScopeTestTimer :: String -> Bool -> IO TestTimer
newSpeedScopeTestTimer String
path Bool
writeRawTimings = do
  Bool -> String -> IO ()
createDirectoryIfMissing Bool
True String
path

  Maybe Handle
maybeHandle <- case Bool
writeRawTimings of
    Bool
False -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
    Bool
True -> do
      Handle
h <- String -> IOMode -> IO Handle
openFile (String
path String -> String -> String
</> String
"timings_raw.txt") IOMode
AppendMode
      Handle -> BufferMode -> IO ()
hSetBuffering Handle
h BufferMode
LineBuffering
      forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just Handle
h

  MVar SpeedScopeFile
speedScopeFile <- forall a. a -> IO (MVar a)
newMVar SpeedScopeFile
emptySpeedScopeFile
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ String -> Maybe Handle -> MVar SpeedScopeFile -> TestTimer
SpeedScopeTestTimer String
path Maybe Handle
maybeHandle MVar SpeedScopeFile
speedScopeFile

finalizeSpeedScopeTestTimer :: TestTimer -> IO ()
finalizeSpeedScopeTestTimer :: TestTimer -> IO ()
finalizeSpeedScopeTestTimer TestTimer
NullTestTimer = forall (m :: * -> *) a. Monad m => a -> m a
return ()
finalizeSpeedScopeTestTimer (SpeedScopeTestTimer {String
Maybe Handle
MVar SpeedScopeFile
$sel:testTimerSpeedScopeFile:SpeedScopeTestTimer :: TestTimer -> MVar SpeedScopeFile
$sel:testTimerHandle:SpeedScopeTestTimer :: TestTimer -> Maybe Handle
$sel:testTimerBasePath:SpeedScopeTestTimer :: TestTimer -> String
testTimerSpeedScopeFile :: MVar SpeedScopeFile
testTimerHandle :: Maybe Handle
testTimerBasePath :: String
..}) = do
  forall (m :: * -> *) a b. Monad m => Maybe a -> (a -> m b) -> m ()
whenJust Maybe Handle
testTimerHandle Handle -> IO ()
hClose
  forall a. MVar a -> IO a
readMVar MVar SpeedScopeFile
testTimerSpeedScopeFile forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= String -> ByteString -> IO ()
BL.writeFile (String
testTimerBasePath String -> String -> String
</> String
"speedscope.json") forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. ToJSON a => a -> ByteString
A.encode

timeAction' :: (MonadMask m, MonadIO m) => TestTimer -> T.Text -> T.Text -> m a -> m a
timeAction' :: forall (m :: * -> *) a.
(MonadMask m, MonadIO m) =>
TestTimer -> EventName -> EventName -> m a -> m a
timeAction' TestTimer
NullTestTimer EventName
_ EventName
_ = forall a. a -> a
id
timeAction' (SpeedScopeTestTimer {String
Maybe Handle
MVar SpeedScopeFile
testTimerSpeedScopeFile :: MVar SpeedScopeFile
testTimerHandle :: Maybe Handle
testTimerBasePath :: String
$sel:testTimerSpeedScopeFile:SpeedScopeTestTimer :: TestTimer -> MVar SpeedScopeFile
$sel:testTimerHandle:SpeedScopeTestTimer :: TestTimer -> Maybe Handle
$sel:testTimerBasePath:SpeedScopeTestTimer :: TestTimer -> String
..}) EventName
profileName EventName
eventName = forall (m :: * -> *) a b c. MonadMask m => m a -> m b -> m c -> m c
bracket_
  (forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVar_ MVar SpeedScopeFile
testTimerSpeedScopeFile forall a b. (a -> b) -> a -> b
$ \SpeedScopeFile
file -> do
    POSIXTime
now <- IO POSIXTime
getPOSIXTime
    SpeedScopeFile -> POSIXTime -> IO SpeedScopeFile
handleStartEvent SpeedScopeFile
file POSIXTime
now
  )
  (forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVar_ MVar SpeedScopeFile
testTimerSpeedScopeFile forall a b. (a -> b) -> a -> b
$ \SpeedScopeFile
file -> do
    POSIXTime
now <- IO POSIXTime
getPOSIXTime
    SpeedScopeFile -> POSIXTime -> IO SpeedScopeFile
handleEndEvent SpeedScopeFile
file POSIXTime
now
  )
  where
    handleStartEvent :: SpeedScopeFile -> POSIXTime -> IO SpeedScopeFile
handleStartEvent SpeedScopeFile
file POSIXTime
time = do
      forall (m :: * -> *) a b. Monad m => Maybe a -> (a -> m b) -> m ()
whenJust Maybe Handle
testTimerHandle forall a b. (a -> b) -> a -> b
$ \Handle
h -> Handle -> EventName -> IO ()
T.hPutStrLn Handle
h [i|#{time} START #{show profileName} #{eventName}|]
      forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ SpeedScopeFile
-> POSIXTime -> SpeedScopeEventType -> SpeedScopeFile
handleSpeedScopeEvent SpeedScopeFile
file POSIXTime
time SpeedScopeEventType
SpeedScopeEventTypeOpen

    handleEndEvent :: SpeedScopeFile -> POSIXTime -> IO SpeedScopeFile
handleEndEvent SpeedScopeFile
file POSIXTime
time = do
      forall (m :: * -> *) a b. Monad m => Maybe a -> (a -> m b) -> m ()
whenJust Maybe Handle
testTimerHandle forall a b. (a -> b) -> a -> b
$ \Handle
h -> Handle -> EventName -> IO ()
T.hPutStrLn Handle
h [i|#{time} END #{show profileName} #{eventName}|]
      forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ SpeedScopeFile
-> POSIXTime -> SpeedScopeEventType -> SpeedScopeFile
handleSpeedScopeEvent SpeedScopeFile
file POSIXTime
time SpeedScopeEventType
SpeedScopeEventTypeClose

    -- | TODO: maybe use an intermediate format so the frames (and possibly profiles) aren't stored as lists,
    -- so we don't have to do O(N) L.length and S.findIndexL
    handleSpeedScopeEvent :: SpeedScopeFile -> POSIXTime -> SpeedScopeEventType -> SpeedScopeFile
    handleSpeedScopeEvent :: SpeedScopeFile
-> POSIXTime -> SpeedScopeEventType -> SpeedScopeFile
handleSpeedScopeEvent SpeedScopeFile
initialFile POSIXTime
time SpeedScopeEventType
typ = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall s a. State s a -> s -> s
execState SpeedScopeFile
initialFile forall a b. (a -> b) -> a -> b
$ do
      Int
frameID <- forall (m :: * -> *) s. Monad m => StateT s m s
get forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \SpeedScopeFile
f -> case forall a. (a -> Bool) -> Seq a -> Maybe Int
S.findIndexL (forall a. Eq a => a -> a -> Bool
== EventName -> SpeedScopeFrame
SpeedScopeFrame EventName
eventName) (SpeedScopeFile
f forall s a. s -> Getting a s a -> a
^. forall s a. HasShared s a => Lens' s a
shared forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s a. HasFrames s a => Lens' s a
frames) of
        Just Int
j -> forall (m :: * -> *) a. Monad m => a -> m a
return Int
j
        Maybe Int
Nothing -> do
          forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
modify' forall a b. (a -> b) -> a -> b
$ forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over (forall s a. HasShared s a => Lens' s a
shared forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s a. HasFrames s a => Lens' s a
frames) (forall a. Seq a -> a -> Seq a
S.|> (EventName -> SpeedScopeFrame
SpeedScopeFrame EventName
eventName))
          forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Seq a -> Int
S.length forall a b. (a -> b) -> a -> b
$ SpeedScopeFile
f forall s a. s -> Getting a s a -> a
^. forall s a. HasShared s a => Lens' s a
shared forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s a. HasFrames s a => Lens' s a
frames

      Int
profileIndex <- forall (m :: * -> *) s. Monad m => StateT s m s
get forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \SpeedScopeFile
f -> case forall a. (a -> Bool) -> [a] -> Maybe Int
L.findIndex ((forall a. Eq a => a -> a -> Bool
== EventName
profileName) forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall s a. s -> Getting a s a -> a
^. forall s a. HasName s a => Lens' s a
name)) (SpeedScopeFile
f forall s a. s -> Getting a s a -> a
^. forall s a. HasProfiles s a => Lens' s a
profiles) of
        Just Int
j -> forall (m :: * -> *) a. Monad m => a -> m a
return Int
j
        Maybe Int
Nothing -> do
          forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
modify' forall a b. (a -> b) -> a -> b
$ forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over forall s a. HasProfiles s a => Lens' s a
profiles (\[SpeedScopeProfile]
x -> [SpeedScopeProfile]
x forall a. Semigroup a => a -> a -> a
<> [EventName -> POSIXTime -> SpeedScopeProfile
newProfile EventName
profileName POSIXTime
time])
          forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t a -> Int
L.length (SpeedScopeFile
f forall s a. s -> Getting a s a -> a
^. forall s a. HasProfiles s a => Lens' s a
profiles)

      forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
modify' forall a b. (a -> b) -> a -> b
$ forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over (forall s a. HasProfiles s a => Lens' s a
profiles forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall m. Ixed m => Index m -> Traversal' m (IxValue m)
ix Int
profileIndex forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s a. HasEvents s a => Lens' s a
events) (forall a. Seq a -> a -> Seq a
S.|> (SpeedScopeEventType -> Int -> POSIXTime -> SpeedScopeEvent
SpeedScopeEvent SpeedScopeEventType
typ Int
frameID POSIXTime
time))
              forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over (forall s a. HasProfiles s a => Lens' s a
profiles forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall m. Ixed m => Index m -> Traversal' m (IxValue m)
ix Int
profileIndex forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s a. HasEndValue s a => Lens' s a
endValue) (forall a. Ord a => a -> a -> a
max POSIXTime
time)