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

import qualified Data.HashMap.Strict 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 = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ forall a. a -> IO (IORef a)
newIORef Bool
False

{-# NOINLINE useTimings #-}
useTimings :: Bool
useTimings :: Bool
useTimings = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ 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 = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ forall a. a -> IO (IORef a)
newIORef 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 forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ forall a. Category -> Category -> IO a -> IO a
timedIO Category
c Category
i forall a b. (a -> b) -> a -> b
$ 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 forall a. Eq a => a -> a -> Bool
== Category
"Hint"
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
quiet forall a b. (a -> b) -> a -> b
$ IO () -> IO ()
whenLoud forall a b. (a -> b) -> a -> b
$ do
        Category -> IO ()
putStr forall a b. (a -> b) -> a -> b
$ Category
"# " forall a. [a] -> [a] -> [a]
++ Category
c forall a. [a] -> [a] -> [a]
++ Category
" of " forall a. [a] -> [a] -> [a]
++ Category
i forall a. [a] -> [a] -> [a]
++ Category
"... "
        Handle -> IO ()
hFlush Handle
stdout
    (Seconds
time, a
x) <- forall (m :: * -> *) a. MonadIO m => m a -> m (Seconds, a)
duration IO a
x
    forall a. IORef a -> (a -> a) -> IO ()
atomicModifyIORef'_ IORef (HashMap (Category, Category) Seconds)
timings forall a b. (a -> b) -> a -> b
$ forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
Map.insertWith forall a. Num a => a -> a -> a
(+) (Category
c, Category
i) Seconds
time
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
quiet forall a b. (a -> b) -> a -> b
$ IO () -> IO ()
whenLoud forall a b. (a -> b) -> a -> b
$ Category -> IO ()
putStrLn forall a b. (a -> b) -> a -> b
$ Category
"took " forall a. [a] -> [a] -> [a]
++ Seconds -> Category
showDuration Seconds
time
    forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x

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

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

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