{-# LANGUAGE ImportQualifiedPost #-}

module Timing(
    timed, timedIO,
    startTimings,
    printTimings
    ) where

import Data.HashMap.Strict qualified as Map
import Control.Exception
import Data.IORef.Extra
import Data.Tuple.Extra
import Data.List.Extra
import Control.Monad
import System.Console.CmdArgs.Verbosity
import System.Time.Extra
import System.IO.Unsafe
import System.IO


type Category = String
type Item = String

{-# NOINLINE useTimingsRef #-}
useTimingsRef :: IORef Bool
useTimingsRef :: IORef Bool
useTimingsRef = IO (IORef Bool) -> IORef Bool
forall a. IO a -> a
unsafePerformIO (IO (IORef Bool) -> IORef Bool) -> IO (IORef Bool) -> IORef Bool
forall a b. (a -> b) -> a -> b
$ Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False

{-# NOINLINE useTimings #-}
useTimings :: Bool
useTimings :: Bool
useTimings = IO Bool -> Bool
forall a. IO a -> a
unsafePerformIO (IO Bool -> Bool) -> IO Bool -> Bool
forall a b. (a -> b) -> a -> b
$ IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef IORef Bool
useTimingsRef

{-# NOINLINE timings #-}
timings :: IORef (Map.HashMap (Category, Item) Seconds)
timings :: IORef (HashMap (Category, Category) Seconds)
timings = IO (IORef (HashMap (Category, Category) Seconds))
-> IORef (HashMap (Category, Category) Seconds)
forall a. IO a -> a
unsafePerformIO (IO (IORef (HashMap (Category, Category) Seconds))
 -> IORef (HashMap (Category, Category) Seconds))
-> IO (IORef (HashMap (Category, Category) Seconds))
-> IORef (HashMap (Category, Category) Seconds)
forall a b. (a -> b) -> a -> b
$ HashMap (Category, Category) Seconds
-> IO (IORef (HashMap (Category, Category) Seconds))
forall a. a -> IO (IORef a)
newIORef HashMap (Category, Category) Seconds
forall k v. HashMap k v
Map.empty

{-# NOINLINE timed #-}
timed :: Category -> Item -> a -> a
timed :: forall a. Category -> Category -> a -> a
timed Category
c Category
i a
x = if Bool -> Bool
not Bool
useTimings then a
x else IO a -> a
forall a. IO a -> a
unsafePerformIO (IO a -> a) -> IO a -> a
forall a b. (a -> b) -> a -> b
$ Category -> Category -> IO a -> IO a
forall a. Category -> Category -> IO a -> IO a
timedIO Category
c Category
i (IO a -> IO a) -> IO a -> IO a
forall a b. (a -> b) -> a -> b
$ a -> IO a
forall a. a -> IO a
evaluate a
x


timedIO :: Category -> Item -> IO a -> IO a
timedIO :: forall a. Category -> Category -> IO a -> IO a
timedIO Category
c Category
i IO a
x = if Bool -> Bool
not Bool
useTimings then IO a
x else do
    let quiet :: Bool
quiet = Category
c Category -> Category -> Bool
forall a. Eq a => a -> a -> Bool
== Category
"Hint"
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
quiet (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO ()
whenLoud (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
        Category -> IO ()
putStr (Category -> IO ()) -> Category -> IO ()
forall a b. (a -> b) -> a -> b
$ Category
"# " Category -> Category -> Category
forall a. [a] -> [a] -> [a]
++ Category
c Category -> Category -> Category
forall a. [a] -> [a] -> [a]
++ Category
" of " Category -> Category -> Category
forall a. [a] -> [a] -> [a]
++ Category
i Category -> Category -> Category
forall a. [a] -> [a] -> [a]
++ Category
"... "
        Handle -> IO ()
hFlush Handle
stdout
    (Seconds
time, a
x) <- IO a -> IO (Seconds, a)
forall (m :: * -> *) a. MonadIO m => m a -> m (Seconds, a)
duration IO a
x
    IORef (HashMap (Category, Category) Seconds)
-> (HashMap (Category, Category) Seconds
    -> HashMap (Category, Category) Seconds)
-> IO ()
forall a. IORef a -> (a -> a) -> IO ()
atomicModifyIORef'_ IORef (HashMap (Category, Category) Seconds)
timings ((HashMap (Category, Category) Seconds
  -> HashMap (Category, Category) Seconds)
 -> IO ())
-> (HashMap (Category, Category) Seconds
    -> HashMap (Category, Category) Seconds)
-> IO ()
forall a b. (a -> b) -> a -> b
$ (Seconds -> Seconds -> Seconds)
-> (Category, Category)
-> Seconds
-> HashMap (Category, Category) Seconds
-> HashMap (Category, Category) Seconds
forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
Map.insertWith Seconds -> Seconds -> Seconds
forall a. Num a => a -> a -> a
(+) (Category
c, Category
i) Seconds
time
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
quiet (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO ()
whenLoud (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Category -> IO ()
putStrLn (Category -> IO ()) -> Category -> IO ()
forall a b. (a -> b) -> a -> b
$ Category
"took " Category -> Category -> Category
forall a. [a] -> [a] -> [a]
++ Seconds -> Category
showDuration Seconds
time
    a -> IO a
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x

startTimings :: IO ()
startTimings :: IO ()
startTimings = do
    IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Bool
useTimingsRef Bool
True
    IORef (HashMap (Category, Category) Seconds)
-> HashMap (Category, Category) Seconds -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (HashMap (Category, Category) Seconds)
timings HashMap (Category, Category) Seconds
forall k v. HashMap k v
Map.empty

printTimings :: IO ()
printTimings :: IO ()
printTimings = do
    HashMap (Category, Category) Seconds
mp <- IORef (HashMap (Category, Category) Seconds)
-> IO (HashMap (Category, Category) Seconds)
forall a. IORef a -> IO a
readIORef IORef (HashMap (Category, Category) Seconds)
timings
    let items :: [(Category, [(Category, Seconds)])]
items = ((Category, [(Category, Seconds)]) -> Seconds)
-> [(Category, [(Category, Seconds)])]
-> [(Category, [(Category, Seconds)])]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn ([(Category, Seconds)] -> Seconds
forall {a}. [(a, Seconds)] -> Seconds
sumSnd ([(Category, Seconds)] -> Seconds)
-> ((Category, [(Category, Seconds)]) -> [(Category, Seconds)])
-> (Category, [(Category, Seconds)])
-> Seconds
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Category, [(Category, Seconds)]) -> [(Category, Seconds)]
forall a b. (a, b) -> b
snd) ([(Category, [(Category, Seconds)])]
 -> [(Category, [(Category, Seconds)])])
-> [(Category, [(Category, Seconds)])]
-> [(Category, [(Category, Seconds)])]
forall a b. (a -> b) -> a -> b
$
                [(Category, (Category, Seconds))]
-> [(Category, [(Category, Seconds)])]
forall k v. Ord k => [(k, v)] -> [(k, [v])]
groupSort ([(Category, (Category, Seconds))]
 -> [(Category, [(Category, Seconds)])])
-> [(Category, (Category, Seconds))]
-> [(Category, [(Category, Seconds)])]
forall a b. (a -> b) -> a -> b
$ (((Category, Category), Seconds)
 -> (Category, (Category, Seconds)))
-> [((Category, Category), Seconds)]
-> [(Category, (Category, Seconds))]
forall a b. (a -> b) -> [a] -> [b]
map (\((Category
a,Category
b),Seconds
c) -> (Category
a,(Category
b,Seconds
c))) ([((Category, Category), Seconds)]
 -> [(Category, (Category, Seconds))])
-> [((Category, Category), Seconds)]
-> [(Category, (Category, Seconds))]
forall a b. (a -> b) -> a -> b
$ HashMap (Category, Category) Seconds
-> [((Category, Category), Seconds)]
forall k v. HashMap k v -> [(k, v)]
Map.toList HashMap (Category, Category) Seconds
mp
    Category -> IO ()
putStrLn (Category -> IO ()) -> Category -> IO ()
forall a b. (a -> b) -> a -> b
$ [Category] -> Category
unlines ([Category] -> Category) -> [Category] -> Category
forall a b. (a -> b) -> a -> b
$ [Category] -> [[Category]] -> [Category]
forall a. [a] -> [[a]] -> [a]
intercalate [Category
""] ([[Category]] -> [Category]) -> [[Category]] -> [Category]
forall a b. (a -> b) -> a -> b
$ ((Category, [(Category, Seconds)]) -> [Category])
-> [(Category, [(Category, Seconds)])] -> [[Category]]
forall a b. (a -> b) -> [a] -> [b]
map (Category, [(Category, Seconds)]) -> [Category]
disp ([(Category, [(Category, Seconds)])] -> [[Category]])
-> [(Category, [(Category, Seconds)])] -> [[Category]]
forall a b. (a -> b) -> a -> b
$ [(Category, [(Category, Seconds)])]
items [(Category, [(Category, Seconds)])]
-> [(Category, [(Category, Seconds)])]
-> [(Category, [(Category, Seconds)])]
forall a. [a] -> [a] -> [a]
++ [(Category
"TOTAL", ((Category, [(Category, Seconds)]) -> (Category, Seconds))
-> [(Category, [(Category, Seconds)])] -> [(Category, Seconds)]
forall a b. (a -> b) -> [a] -> [b]
map (([(Category, Seconds)] -> Seconds)
-> (Category, [(Category, Seconds)]) -> (Category, Seconds)
forall b b' a. (b -> b') -> (a, b) -> (a, b')
second [(Category, Seconds)] -> Seconds
forall {a}. [(a, Seconds)] -> Seconds
sumSnd) [(Category, [(Category, Seconds)])]
items)]
    where
        sumSnd :: [(a, Seconds)] -> Seconds
sumSnd = [Seconds] -> Seconds
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Seconds] -> Seconds)
-> ([(a, Seconds)] -> [Seconds]) -> [(a, Seconds)] -> Seconds
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, Seconds) -> Seconds) -> [(a, Seconds)] -> [Seconds]
forall a b. (a -> b) -> [a] -> [b]
map (a, Seconds) -> Seconds
forall a b. (a, b) -> b
snd

        disp :: (Category, [(Category, Seconds)]) -> [Category]
disp (Category
cat,[(Category, Seconds)]
xs) =
                (Category
"Timing " Category -> Category -> Category
forall a. [a] -> [a] -> [a]
++ Category
cat) Category -> [Category] -> [Category]
forall a. a -> [a] -> [a]
:
                [Category
"  " Category -> Category -> Category
forall a. [a] -> [a] -> [a]
++ Seconds -> Category
showDuration Seconds
b Category -> Category -> Category
forall a. [a] -> [a] -> [a]
++ Category
" " Category -> Category -> Category
forall a. [a] -> [a] -> [a]
++ Category
a | (Category
a,Seconds
b) <- [(Category, Seconds)]
xs2] [Category] -> [Category] -> [Category]
forall a. [a] -> [a] -> [a]
++
                [Category
"  " Category -> Category -> Category
forall a. [a] -> [a] -> [a]
++ Seconds -> Category
showDuration ([(Category, Seconds)] -> Seconds
forall {a}. [(a, Seconds)] -> Seconds
sumSnd [(Category, Seconds)]
xs2) Category -> Category -> Category
forall a. [a] -> [a] -> [a]
++ Category
" TOTAL"]
            where
                xs2 :: [(Category, Seconds)]
xs2 = ([(Category, Seconds)], [(Category, Seconds)])
-> [(Category, Seconds)]
f (([(Category, Seconds)], [(Category, Seconds)])
 -> [(Category, Seconds)])
-> ([(Category, Seconds)], [(Category, Seconds)])
-> [(Category, Seconds)]
forall a b. (a -> b) -> a -> b
$ Int
-> [(Category, Seconds)]
-> ([(Category, Seconds)], [(Category, Seconds)])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
9 ([(Category, Seconds)]
 -> ([(Category, Seconds)], [(Category, Seconds)]))
-> [(Category, Seconds)]
-> ([(Category, Seconds)], [(Category, Seconds)])
forall a b. (a -> b) -> a -> b
$ ((Category, Seconds) -> Seconds)
-> [(Category, Seconds)] -> [(Category, Seconds)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Seconds -> Seconds
forall a. Num a => a -> a
negate (Seconds -> Seconds)
-> ((Category, Seconds) -> Seconds)
-> (Category, Seconds)
-> Seconds
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Category, Seconds) -> Seconds
forall a b. (a, b) -> b
snd) [(Category, Seconds)]
xs
                f :: ([(Category, Seconds)], [(Category, Seconds)])
-> [(Category, Seconds)]
f ([(Category, Seconds)]
xs,[(Category, Seconds)]
ys)
                    | [(Category, Seconds)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Category, Seconds)]
ys Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 = [(Category, Seconds)]
xs [(Category, Seconds)]
-> [(Category, Seconds)] -> [(Category, Seconds)]
forall a. [a] -> [a] -> [a]
++ [(Category, Seconds)]
ys
                    | Bool
otherwise = [(Category, Seconds)]
xs [(Category, Seconds)]
-> [(Category, Seconds)] -> [(Category, Seconds)]
forall a. [a] -> [a] -> [a]
++ [(Category
"Other items (" Category -> Category -> Category
forall a. [a] -> [a] -> [a]
++ Int -> Category
forall a. Show a => a -> Category
show ([(Category, Seconds)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Category, Seconds)]
ys) Category -> Category -> Category
forall a. [a] -> [a] -> [a]
++ Category
")", [(Category, Seconds)] -> Seconds
forall {a}. [(a, Seconds)] -> Seconds
sumSnd [(Category, Seconds)]
ys)]