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 :: CombIx -> K -> Profile Word64 -> Profile Word64
addSample :: CombIx -> K -> Profile Word64 -> Profile Word64
addSample CombIx
c K
k (Prof Int
count ProfTrie Word64 Int
trie Map Word64 Reference
refs) =
  Int
-> ProfTrie Word64 Int -> Map Word64 Reference -> Profile Word64
forall k. Int -> ProfTrie k Int -> Map k Reference -> Profile k
Prof
    (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
count)
    ([Word64] -> ProfTrie Word64 Int -> ProfTrie Word64 Int
forall k. Ord k => [k] -> ProfTrie k Int -> ProfTrie k Int
addPath ((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
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
    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 :: [(CombIx, K)] -> Profile Word64 -> Profile Word64
addSamples :: [(CombIx, K)] -> Profile Word64 -> Profile Word64
addSamples [(CombIx, K)]
ts Profile Word64
p = (Profile Word64 -> (CombIx, K) -> Profile Word64)
-> Profile Word64 -> [(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' (((CombIx, K) -> Profile Word64 -> Profile Word64)
-> Profile Word64 -> (CombIx, K) -> Profile Word64
forall a b c. (a -> b -> c) -> b -> a -> c
flip (((CombIx, K) -> Profile Word64 -> Profile Word64)
 -> Profile Word64 -> (CombIx, K) -> Profile Word64)
-> ((CombIx -> K -> Profile Word64 -> Profile Word64)
    -> (CombIx, K) -> Profile Word64 -> Profile Word64)
-> (CombIx -> K -> Profile Word64 -> Profile Word64)
-> Profile Word64
-> (CombIx, K)
-> Profile Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (CombIx -> K -> Profile Word64 -> Profile Word64)
-> (CombIx, K) -> Profile Word64 -> Profile Word64
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((CombIx -> K -> Profile Word64 -> Profile Word64)
 -> Profile Word64 -> (CombIx, K) -> Profile Word64)
-> (CombIx -> K -> Profile Word64 -> Profile Word64)
-> Profile Word64
-> (CombIx, K)
-> Profile Word64
forall a b. (a -> b) -> a -> b
$ CombIx -> K -> Profile Word64 -> Profile Word64
addSample) Profile Word64
p [(CombIx, K)]
ts

-- 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 [(CombIx, K)]
  | Final [(CombIx, K)]

readInput :: TVar TickComm -> IO (Bool, [(CombIx, K)])
readInput :: TVar TickComm -> IO (Bool, [(CombIx, K)])
readInput TVar TickComm
input =
  STM (Bool, [(CombIx, K)]) -> IO (Bool, [(CombIx, K)])
forall a. STM a -> IO a
atomically (STM (Bool, [(CombIx, K)]) -> IO (Bool, [(CombIx, K)]))
-> STM (Bool, [(CombIx, K)]) -> IO (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, [(CombIx, K)]))
-> STM (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, [(CombIx, K)])
forall a. STM a
retry
      TickComm
Finished -> (Bool, [(CombIx, K)]) -> STM (Bool, [(CombIx, K)])
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool
True, [])
      Ticks [(CombIx, K)]
ts -> (Bool
False, [(CombIx, K)]
ts) (Bool, [(CombIx, K)]) -> STM () -> STM (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 [(CombIx, K)]
ts -> (Bool
True, [(CombIx, K)]
ts) (Bool, [(CombIx, K)]) -> STM () -> STM (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, [(CombIx, K)]
ts) <- TVar TickComm -> IO (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
$ [(CombIx, K)] -> Profile Word64 -> Profile Word64
addSamples [(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 -> CombIx -> K -> IO ()
enqueue :: TVar TickComm -> CombIx -> K -> IO ()
enqueue TVar TickComm
comm 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 -> [(CombIx, K)] -> TickComm
Ticks [(CombIx
c, K
k)]
    TickComm
Finished -> [(CombIx, K)] -> TickComm
Final [(CombIx
c, K
k)]
    Ticks [(CombIx, K)]
ts -> [(CombIx, K)] -> TickComm
Ticks ((CombIx
c, K
k) (CombIx, K) -> [(CombIx, K)] -> [(CombIx, K)]
forall a. a -> [a] -> [a]
: [(CombIx, K)]
ts)
    Final [(CombIx, K)]
ts -> [(CombIx, K)] -> TickComm
Final ((CombIx
c, K
k) (CombIx, K) -> [(CombIx, K)] -> [(CombIx, K)]
forall a. a -> [a] -> [a]
: [(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 [(CombIx, K)]
ts -> [(CombIx, K)] -> TickComm
Final [(CombIx, K)]
ts
    Final [(CombIx, K)]
ts -> [(CombIx, K)] -> TickComm
Final [(CombIx, K)]
ts

data ProfileComm
  = PC
      (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 $ (CombIx -> K -> IO ())
-> IO () -> IO (Profile Word64) -> ProfileComm
PC (TVar TickComm -> 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)