    {-# LANGUAGE RecordWildCards #-}

    module General.Timing(Timing, withTiming, timed, timedOverwrite) where

    import Data.List.Extra
    import System.Time.Extra
    import Data.IORef
    import Control.Monad.Extra
    import System.IO
    import General.Util
    import Control.Monad.IO.Class

    -- | A mutable object to keep timing information
    data Timing = Timing
        { -- | Get time since the initialization of this 'Timing'.
          Timing -> IO Seconds
    timingOffset :: IO Seconds
          -- | Record timings for writing to a file
        , Timing -> IORef [(String, Seconds)]
    timingStore :: IORef [(String, Seconds)]
          -- | If you are below T you may overwrite N characters
          -- at the end of the current terminal output.
          -- Only used iff @timingTerminal == True@.
        , Timing -> IORef (Maybe (Seconds, Int))
    timingOverwrite :: IORef (Maybe (Seconds, Int))
          -- | whether is this a terminal
        , Timing -> Bool
    timingTerminal :: Bool
        }

    -- | Time an action, printing timing information to the terminal
    withTiming ::
      -- | A file to optionally write all timings to, after the action is finished
      Maybe FilePath ->
      -- | An action that can write timings into 'Timing'
      (Timing -> IO a) ->
      IO a
    withTiming :: forall a. Maybe String -> (Timing -> IO a) -> IO a
    withTiming Maybe String
    writeTimingsTo Timing -> IO a
    act = do
        IO Seconds
    timingOffset <- IO (IO Seconds)
    offsetTime
        IORef [(String, Seconds)]
    timingStore <- [(String, Seconds)] -> IO (IORef [(String, Seconds)])
    forall a. a -> IO (IORef a)
    newIORef []
        IORef (Maybe (Seconds, Int))
    timingOverwrite <- Maybe (Seconds, Int) -> IO (IORef (Maybe (Seconds, Int)))
    forall a. a -> IO (IORef a)
    newIORef Maybe (Seconds, Int)
    forall a. Maybe a
    Nothing
        Bool
    timingTerminal <- Handle -> IO Bool
    hIsTerminalDevice Handle
    stdout

        a
    res <- Timing -> IO a
    act Timing{Bool
    IO Seconds
    IORef [(String, Seconds)]
    IORef (Maybe (Seconds, Int))
    timingOffset :: IO Seconds
    timingStore :: IORef [(String, Seconds)]
    timingOverwrite :: IORef (Maybe (Seconds, Int))
    timingTerminal :: Bool
    timingOffset :: IO Seconds
    timingStore :: IORef [(String, Seconds)]
    timingOverwrite :: IORef (Maybe (Seconds, Int))
    timingTerminal :: Bool
    ..}
        Seconds
    total <- IO Seconds
    timingOffset
        Maybe String -> (String -> IO ()) -> IO ()
    forall (m :: * -> *) a.
    Applicative m =>
    Maybe a -> (a -> m ()) -> m ()
    whenJust Maybe String
    writeTimingsTo ((String -> IO ()) -> IO ()) -> (String -> IO ()) -> IO ()
    forall a b. (a -> b) -> a -> b
    $ \String
    file -> do
            [(String, Seconds)]
    xs <- IORef [(String, Seconds)] -> IO [(String, Seconds)]
    forall a. IORef a -> IO a
    readIORef IORef [(String, Seconds)]
    timingStore
            -- Expecting unrecorded of ~2s
            -- Most of that comes from the pipeline - we get occasional 0.01 between items as one flushes
            -- Then at the end there is ~0.5 while the final item flushes
            [(String, Seconds)]
    xs <- [(String, Seconds)] -> IO [(String, Seconds)]
    forall a. a -> IO a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ([(String, Seconds)] -> IO [(String, Seconds)])
    -> [(String, Seconds)] -> IO [(String, Seconds)]
    forall a b. (a -> b) -> a -> b
    $ ((String, Seconds) -> Seconds)
    -> [(String, Seconds)] -> [(String, Seconds)]
    forall b a. Ord b => (a -> b) -> [a] -> [a]
    sortOn (Seconds -> Seconds
    forall a. Num a => a -> a
    negate (Seconds -> Seconds)
    -> ((String, Seconds) -> Seconds) -> (String, Seconds) -> Seconds
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (String, Seconds) -> Seconds
    forall a b. (a, b) -> b
    snd) ([(String, Seconds)] -> [(String, Seconds)])
    -> [(String, Seconds)] -> [(String, Seconds)]
    forall a b. (a -> b) -> a -> b
    $ (String
    "Unrecorded", Seconds
    total Seconds -> Seconds -> Seconds
    forall a. Num a => a -> a -> a
    - [Seconds] -> Seconds
    forall a. Num a => [a] -> a
    forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
    sum (((String, Seconds) -> Seconds) -> [(String, Seconds)] -> [Seconds]
    forall a b. (a -> b) -> [a] -> [b]
    map (String, Seconds) -> Seconds
    forall a b. (a, b) -> b
    snd [(String, Seconds)]
    xs)) (String, Seconds) -> [(String, Seconds)] -> [(String, Seconds)]
    forall a. a -> [a] -> [a]
    : [(String, Seconds)]
    xs
            String -> String -> IO ()
    writeFile String
    file (String -> IO ()) -> String -> IO ()
    forall a b. (a -> b) -> a -> b
    $ [String] -> String
    unlines ([String] -> String) -> [String] -> String
    forall a b. (a -> b) -> a -> b
    $ Int -> String -> [(String, Seconds)] -> [String]
    prettyTable Int
    2 String
    "Secs" [(String, Seconds)]
    xs
        String -> IO ()
    putStrLn (String -> IO ()) -> String -> IO ()
    forall a b. (a -> b) -> a -> b
    $ String
    "Took " String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ Seconds -> String
    showDuration Seconds
    total
        a -> IO a
    forall a. a -> IO a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure a
    res

    -- skip it if have written out in the last 1s and takes < 0.1

    -- | Time & write the given message to stdout
    timed :: MonadIO m => Timing -> String -> m a -> m a
    timed :: forall (m :: * -> *) a. MonadIO m => Timing -> String -> m a -> m a
    timed = Bool -> Timing -> String -> m a -> m a
    forall (m :: * -> *) a.
    MonadIO m =>
    Bool -> Timing -> String -> m a -> m a
    timedEx Bool
    False

    -- | Time & write the given message to stdout
    -- overwriting a previous message if it was marked as overwritable
    timedOverwrite :: MonadIO m => Timing -> String -> m a -> m a
    timedOverwrite :: forall (m :: * -> *) a. MonadIO m => Timing -> String -> m a -> m a
    timedOverwrite = Bool -> Timing -> String -> m a -> m a
    forall (m :: * -> *) a.
    MonadIO m =>
    Bool -> Timing -> String -> m a -> m a
    timedEx Bool
    True

    timedEx :: MonadIO m => Bool -> Timing -> String -> m a -> m a
    timedEx :: forall (m :: * -> *) a.
    MonadIO m =>
    Bool -> Timing -> String -> m a -> m a
    timedEx Bool
    overwrite Timing{Bool
    IO Seconds
    IORef [(String, Seconds)]
    IORef (Maybe (Seconds, Int))
    timingOffset :: Timing -> IO Seconds
    timingStore :: Timing -> IORef [(String, Seconds)]
    timingOverwrite :: Timing -> IORef (Maybe (Seconds, Int))
    timingTerminal :: Timing -> Bool
    timingOffset :: IO Seconds
    timingStore :: IORef [(String, Seconds)]
    timingOverwrite :: IORef (Maybe (Seconds, Int))
    timingTerminal :: Bool
    ..} String
    msg m a
    act = do
        Seconds
    start <- IO Seconds -> m Seconds
    forall a. IO a -> m a
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
    liftIO IO Seconds
    timingOffset
        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
    $ IO (Maybe (Seconds, Int)) -> ((Seconds, Int) -> IO ()) -> IO ()
    forall (m :: * -> *) a.
    Monad m =>
    m (Maybe a) -> (a -> m ()) -> m ()
    whenJustM (IORef (Maybe (Seconds, Int)) -> IO (Maybe (Seconds, Int))
    forall a. IORef a -> IO a
    readIORef IORef (Maybe (Seconds, Int))
    timingOverwrite) (((Seconds, Int) -> IO ()) -> IO ())
    -> ((Seconds, Int) -> IO ()) -> IO ()
    forall a b. (a -> b) -> a -> b
    $ \(Seconds
    t,Int
    n) ->
            if Bool
    overwrite Bool -> Bool -> Bool
    && Seconds
    start Seconds -> Seconds -> Bool
    forall a. Ord a => a -> a -> Bool
    < Seconds
    t then
                String -> IO ()
    putStr (String -> IO ()) -> String -> IO ()
    forall a b. (a -> b) -> a -> b
    $ Int -> Char -> String
    forall a. Int -> a -> [a]
    replicate Int
    n Char
    '\b' String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ Int -> Char -> String
    forall a. Int -> a -> [a]
    replicate Int
    n Char
    ' ' String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ Int -> Char -> String
    forall a. Int -> a -> [a]
    replicate Int
    n Char
    '\b'
            else
                String -> IO ()
    putStrLn String
    ""

        let out :: String -> m Int
    out String
    msg = IO Int -> m Int
    forall a. IO a -> m a
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
    liftIO (IO Int -> m Int) -> IO Int -> m Int
    forall a b. (a -> b) -> a -> b
    $ String -> IO ()
    putStr String
    msg IO () -> IO Int -> IO Int
    forall a b. IO a -> IO b -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
    >> Int -> IO Int
    forall a. a -> IO a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (String -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length String
    msg)
        Int
    undo1 <- String -> m Int
    forall {m :: * -> *}. MonadIO m => String -> m Int
    out (String -> m Int) -> String -> m Int
    forall a b. (a -> b) -> a -> b
    $ String
    msg String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ String
    "... "
        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
    $ Handle -> IO ()
    hFlush Handle
    stdout

        a
    res <- m a
    act
        Seconds
    end <- IO Seconds -> m Seconds
    forall a. IO a -> m a
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
    liftIO IO Seconds
    timingOffset
        let time :: Seconds
    time = Seconds
    end Seconds -> Seconds -> Seconds
    forall a. Num a => a -> a -> a
    - Seconds
    start
        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
    $ IORef [(String, Seconds)]
    -> ([(String, Seconds)] -> [(String, Seconds)]) -> IO ()
    forall a. IORef a -> (a -> a) -> IO ()
    modifyIORef IORef [(String, Seconds)]
    timingStore ((String
    msg,Seconds
    time)(String, Seconds) -> [(String, Seconds)] -> [(String, Seconds)]
    forall a. a -> [a] -> [a]
    :)

        String
    s <- String -> (String -> String) -> Maybe String -> String
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe String
    "" (\String
    x -> String
    " (" String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ String
    x String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ String
    ")") (Maybe String -> String) -> m (Maybe String) -> m String
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> IO (Maybe String) -> m (Maybe String)
    forall a. IO a -> m a
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
    liftIO IO (Maybe String)
    getStatsPeakAllocBytes
        Int
    undo2 <- String -> m Int
    forall {m :: * -> *}. MonadIO m => String -> m Int
    out (String -> m Int) -> String -> m Int
    forall a b. (a -> b) -> a -> b
    $ Seconds -> String
    showDuration Seconds
    time String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ String
    s

        Maybe (Seconds, Int)
    old <- IO (Maybe (Seconds, Int)) -> m (Maybe (Seconds, Int))
    forall a. IO a -> m a
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
    liftIO (IO (Maybe (Seconds, Int)) -> m (Maybe (Seconds, Int)))
    -> IO (Maybe (Seconds, Int)) -> m (Maybe (Seconds, Int))
    forall a b. (a -> b) -> a -> b
    $ IORef (Maybe (Seconds, Int)) -> IO (Maybe (Seconds, Int))
    forall a. IORef a -> IO a
    readIORef IORef (Maybe (Seconds, Int))
    timingOverwrite
        let next :: Seconds
    next = Seconds
    -> ((Seconds, Int) -> Seconds) -> Maybe (Seconds, Int) -> Seconds
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe (Seconds
    start Seconds -> Seconds -> Seconds
    forall a. Num a => a -> a -> a
    + Seconds
    1.0) (Seconds, Int) -> Seconds
    forall a b. (a, b) -> a
    fst Maybe (Seconds, Int)
    old
        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
    timingTerminal Bool -> Bool -> Bool
    && Bool
    overwrite Bool -> Bool -> Bool
    && Seconds
    end Seconds -> Seconds -> Bool
    forall a. Ord a => a -> a -> Bool
    < Seconds
    next then
            IORef (Maybe (Seconds, Int)) -> Maybe (Seconds, Int) -> IO ()
    forall a. IORef a -> a -> IO ()
    writeIORef IORef (Maybe (Seconds, Int))
    timingOverwrite (Maybe (Seconds, Int) -> IO ()) -> Maybe (Seconds, Int) -> IO ()
    forall a b. (a -> b) -> a -> b
    $ (Seconds, Int) -> Maybe (Seconds, Int)
    forall a. a -> Maybe a
    Just (Seconds
    next, Int
    undo1 Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    undo2)
         else do
            IORef (Maybe (Seconds, Int)) -> Maybe (Seconds, Int) -> IO ()
    forall a. IORef a -> a -> IO ()
    writeIORef IORef (Maybe (Seconds, Int))
    timingOverwrite Maybe (Seconds, Int)
    forall a. Maybe a
    Nothing
            String -> IO ()
    putStrLn String
    ""
        a -> m a
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure a
    res