module Unison.Runtime.Profiling where

import Control.Concurrent
import Control.Concurrent.STM
import Data.Foldable
import Data.Map.Strict qualified as M
import Data.Word
import Unison.Codebase.Runtime.Profile
import Unison.Runtime.MCode
import Unison.Runtime.Stack

addSample :: Bool -> CombIx -> K -> Profile Word64 -> Profile Word64
addSample :: Bool -> CombIx -> K -> Profile Word64 -> Profile Word64
addSample Bool
wait CombIx
c K
k (Prof (Int, Int)
count ProfTrie Word64 (Int, Int)
trie Map Word64 Reference
refs) =
  (Int, Int)
-> ProfTrie Word64 (Int, Int)
-> Map Word64 Reference
-> Profile Word64
forall k.
(Int, Int) -> ProfTrie k (Int, Int) -> Map k Reference -> Profile k
Prof
    (Bool -> (Int, Int) -> (Int, Int)
forall {a} {b}. (Num a, Num b) => Bool -> (a, b) -> (a, b)
inc Bool
wait (Int, Int)
count)
    (Bool
-> [Word64]
-> ProfTrie Word64 (Int, Int)
-> ProfTrie Word64 (Int, Int)
forall k.
Ord k =>
Bool -> [k] -> ProfTrie k (Int, Int) -> ProfTrie k (Int, Int)
addPath Bool
wait ((Word64, Reference) -> Word64
forall a b. (a, b) -> a
fst ((Word64, Reference) -> Word64)
-> [(Word64, Reference)] -> [Word64]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Word64, Reference)]
cmbs) ProfTrie Word64 (Int, Int)
trie)
    (Map Word64 Reference
-> Map Word64 Reference -> Map Word64 Reference
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union Map Word64 Reference
refs (Map Word64 Reference -> Map Word64 Reference)
-> Map Word64 Reference -> Map Word64 Reference
forall a b. (a -> b) -> a -> b
$ [(Word64, Reference)] -> Map Word64 Reference
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Word64, Reference)]
cmbs)
  where
    inc :: Bool -> (a, b) -> (a, b)
inc Bool
b (a
m, b
n) = a -> b -> (a, b)
forall {a} {b}. a -> b -> (a, b)
pair (a
m a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) (if Bool
b then b
n b -> b -> b
forall a. Num a => a -> a -> a
+ b
1 else b
n)
    pair :: a -> b -> (a, b)
pair !a
m !b
n = (a
m, b
n)

    cixToPair :: CombIx -> (Word64, Reference)
cixToPair (CIx Reference
r Word64
i Word64
_) = (Word64
i, Reference
r)

    cmbs :: [(Word64, Reference)]
cmbs = [(Word64, Reference)] -> K -> [(Word64, Reference)]
combs [CombIx -> (Word64, Reference)
cixToPair CombIx
c] K
k

    combs :: [(Word64, Reference)] -> K -> [(Word64, Reference)]
combs [(Word64, Reference)]
acc K
KE = [(Word64, Reference)]
acc
    combs [(Word64, Reference)]
acc (CB Callback
_) = [(Word64, Reference)]
acc
    combs [(Word64, Reference)]
acc (AMark Int
_ AEnv
_ AffineRef
_ K
k) = [(Word64, Reference)] -> K -> [(Word64, Reference)]
combs [(Word64, Reference)]
acc K
k
    combs [(Word64, Reference)]
acc (Mark Int
_ EnumSet Word64
_ DEnv
_ K
k) = [(Word64, Reference)] -> K -> [(Word64, Reference)]
combs [(Word64, Reference)]
acc K
k
    combs [(Word64, Reference)]
acc (Local HEnv
_ Int
_ K
k) = [(Word64, Reference)] -> K -> [(Word64, Reference)]
combs [(Word64, Reference)]
acc K
k
    combs [(Word64, Reference)]
acc (Push Int
_ Int
_ CombIx
c Int
_ RSection Val
_ K
k) = [(Word64, Reference)] -> K -> [(Word64, Reference)]
combs (CombIx -> (Word64, Reference)
cixToPair CombIx
c (Word64, Reference)
-> [(Word64, Reference)] -> [(Word64, Reference)]
forall a. a -> [a] -> [a]
: [(Word64, Reference)]
acc) K
k

addSamples :: [(Bool, CombIx, K)] -> Profile Word64 -> Profile Word64
addSamples :: [(Bool, CombIx, K)] -> Profile Word64 -> Profile Word64
addSamples [(Bool, CombIx, K)]
ts Profile Word64
p = (Profile Word64 -> (Bool, CombIx, K) -> Profile Word64)
-> Profile Word64 -> [(Bool, CombIx, K)] -> Profile Word64
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (((Bool, CombIx, K) -> Profile Word64 -> Profile Word64)
-> Profile Word64 -> (Bool, CombIx, K) -> Profile Word64
forall a b c. (a -> b -> c) -> b -> a -> c
flip (((Bool, CombIx, K) -> Profile Word64 -> Profile Word64)
 -> Profile Word64 -> (Bool, CombIx, K) -> Profile Word64)
-> ((Bool -> CombIx -> K -> Profile Word64 -> Profile Word64)
    -> (Bool, CombIx, K) -> Profile Word64 -> Profile Word64)
-> (Bool -> CombIx -> K -> Profile Word64 -> Profile Word64)
-> Profile Word64
-> (Bool, CombIx, K)
-> Profile Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Bool -> CombIx -> K -> Profile Word64 -> Profile Word64)
-> (Bool, CombIx, K) -> Profile Word64 -> Profile Word64
forall {t} {t} {t} {t}. (t -> t -> t -> t) -> (t, t, t) -> t
uncurry3 ((Bool -> CombIx -> K -> Profile Word64 -> Profile Word64)
 -> Profile Word64 -> (Bool, CombIx, K) -> Profile Word64)
-> (Bool -> CombIx -> K -> Profile Word64 -> Profile Word64)
-> Profile Word64
-> (Bool, CombIx, K)
-> Profile Word64
forall a b. (a -> b) -> a -> b
$ Bool -> CombIx -> K -> Profile Word64 -> Profile Word64
addSample) Profile Word64
p [(Bool, CombIx, K)]
ts
  where
    uncurry3 :: (t -> t -> t -> t) -> (t, t, t) -> t
uncurry3 t -> t -> t -> t
f (t
x, t
y, t
z) = t -> t -> t -> t
f t
x t
y t
z

-- For communication between execution and a profiling thread. `Final`
-- indicates that execution is complete and the profiling thread should
-- write a final result to annother channel.
data TickComm
  = Empty
  | Finished
  | Ticks [(Bool, CombIx, K)]
  | Final [(Bool, CombIx, K)]

readInput :: TVar TickComm -> IO (Bool, [(Bool, CombIx, K)])
readInput :: TVar TickComm -> IO (Bool, [(Bool, CombIx, K)])
readInput TVar TickComm
input =
  STM (Bool, [(Bool, CombIx, K)]) -> IO (Bool, [(Bool, CombIx, K)])
forall a. STM a -> IO a
atomically (STM (Bool, [(Bool, CombIx, K)]) -> IO (Bool, [(Bool, CombIx, K)]))
-> STM (Bool, [(Bool, CombIx, K)])
-> IO (Bool, [(Bool, CombIx, K)])
forall a b. (a -> b) -> a -> b
$
    TVar TickComm -> STM TickComm
forall a. TVar a -> STM a
readTVar TVar TickComm
input STM TickComm
-> (TickComm -> STM (Bool, [(Bool, CombIx, K)]))
-> STM (Bool, [(Bool, CombIx, K)])
forall a b. STM a -> (a -> STM b) -> STM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      TickComm
Empty -> STM (Bool, [(Bool, CombIx, K)])
forall a. STM a
retry
      TickComm
Finished -> (Bool, [(Bool, CombIx, K)]) -> STM (Bool, [(Bool, CombIx, K)])
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool
True, [])
      Ticks [(Bool, CombIx, K)]
ts -> (Bool
False, [(Bool, CombIx, K)]
ts) (Bool, [(Bool, CombIx, K)])
-> STM () -> STM (Bool, [(Bool, CombIx, K)])
forall a b. a -> STM b -> STM a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ TVar TickComm -> TickComm -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar TickComm
input TickComm
Empty
      Final [(Bool, CombIx, K)]
ts -> (Bool
True, [(Bool, CombIx, K)]
ts) (Bool, [(Bool, CombIx, K)])
-> STM () -> STM (Bool, [(Bool, CombIx, K)])
forall a b. a -> STM b -> STM a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ TVar TickComm -> TickComm -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar TickComm
input TickComm
Finished

profileLoop ::
  TVar TickComm ->
  TMVar (Profile Word64) ->
  Profile Word64 ->
  IO ()
profileLoop :: TVar TickComm -> TMVar (Profile Word64) -> Profile Word64 -> IO ()
profileLoop TVar TickComm
input TMVar (Profile Word64)
output Profile Word64
prof = do
  (Bool
finish, [(Bool, CombIx, K)]
ts) <- TVar TickComm -> IO (Bool, [(Bool, CombIx, K)])
readInput TVar TickComm
input
  Profile Word64
prof <- Profile Word64 -> IO (Profile Word64)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Profile Word64 -> IO (Profile Word64))
-> Profile Word64 -> IO (Profile Word64)
forall a b. (a -> b) -> a -> b
$ [(Bool, CombIx, K)] -> Profile Word64 -> Profile Word64
addSamples [(Bool, CombIx, K)]
ts Profile Word64
prof
  if Bool -> Bool
not Bool
finish
    then TVar TickComm -> TMVar (Profile Word64) -> Profile Word64 -> IO ()
profileLoop TVar TickComm
input TMVar (Profile Word64)
output Profile Word64
prof
    else STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TMVar (Profile Word64) -> Profile Word64 -> STM ()
forall a. TMVar a -> a -> STM ()
putTMVar TMVar (Profile Word64)
output Profile Word64
prof

enqueue :: TVar TickComm -> Bool -> CombIx -> K -> IO ()
enqueue :: TVar TickComm -> Bool -> CombIx -> K -> IO ()
enqueue TVar TickComm
comm Bool
b CombIx
c K
k = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$
  TVar TickComm -> (TickComm -> TickComm) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar TVar TickComm
comm \case
    TickComm
Empty -> [(Bool, CombIx, K)] -> TickComm
Ticks [(Bool
b, CombIx
c, K
k)]
    TickComm
Finished -> [(Bool, CombIx, K)] -> TickComm
Final [(Bool
b, CombIx
c, K
k)]
    Ticks [(Bool, CombIx, K)]
ts -> [(Bool, CombIx, K)] -> TickComm
Ticks ((Bool
b, CombIx
c, K
k) (Bool, CombIx, K) -> [(Bool, CombIx, K)] -> [(Bool, CombIx, K)]
forall a. a -> [a] -> [a]
: [(Bool, CombIx, K)]
ts)
    Final [(Bool, CombIx, K)]
ts -> [(Bool, CombIx, K)] -> TickComm
Final ((Bool
b, CombIx
c, K
k) (Bool, CombIx, K) -> [(Bool, CombIx, K)] -> [(Bool, CombIx, K)]
forall a. a -> [a] -> [a]
: [(Bool, CombIx, K)]
ts)

finish :: TVar TickComm -> IO ()
finish :: TVar TickComm -> IO ()
finish TVar TickComm
comm = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$
  TVar TickComm -> (TickComm -> TickComm) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar TVar TickComm
comm \case
    TickComm
Empty -> TickComm
Finished
    TickComm
Finished -> TickComm
Finished
    Ticks [(Bool, CombIx, K)]
ts -> [(Bool, CombIx, K)] -> TickComm
Final [(Bool, CombIx, K)]
ts
    Final [(Bool, CombIx, K)]
ts -> [(Bool, CombIx, K)] -> TickComm
Final [(Bool, CombIx, K)]
ts

data ProfileComm
  = PC
      (Bool -> CombIx -> K -> IO ())
      (IO ())
      (IO (Profile Word64))

spawnProfiler :: IO ProfileComm
spawnProfiler :: IO ProfileComm
spawnProfiler = do
  TVar TickComm
input <- TickComm -> IO (TVar TickComm)
forall a. a -> IO (TVar a)
newTVarIO TickComm
Empty
  TMVar (Profile Word64)
output <- IO (TMVar (Profile Word64))
forall a. IO (TMVar a)
newEmptyTMVarIO
  ThreadId
_ <- IO () -> IO ThreadId
forkIO (IO () -> IO ThreadId) -> IO () -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ TVar TickComm -> TMVar (Profile Word64) -> Profile Word64 -> IO ()
profileLoop TVar TickComm
input TMVar (Profile Word64)
output Profile Word64
forall k. Profile k
emptyProfile
  pure $ (Bool -> CombIx -> K -> IO ())
-> IO () -> IO (Profile Word64) -> ProfileComm
PC (TVar TickComm -> Bool -> CombIx -> K -> IO ()
enqueue TVar TickComm
input) (TVar TickComm -> IO ()
finish TVar TickComm
input) (STM (Profile Word64) -> IO (Profile Word64)
forall a. STM a -> IO a
atomically (STM (Profile Word64) -> IO (Profile Word64))
-> STM (Profile Word64) -> IO (Profile Word64)
forall a b. (a -> b) -> a -> b
$ TMVar (Profile Word64) -> STM (Profile Word64)
forall a. TMVar a -> STM a
takeTMVar TMVar (Profile Word64)
output)