Count.hs

module GHC.Debug.Count (parCount, count, asyncCount, asyncParCount) where

import GHC.Debug.Async
import           GHC.Debug.Types
import GHC.Debug.Client.Monad
import           GHC.Debug.Profile
import           GHC.Debug.Trace
import           GHC.Debug.ParTrace hiding (TraceFunctionsIO(..))
import Control.Monad.State
import Control.Monad.Reader
import Data.IORef

parCount :: [ClosurePtr] -> DebugM CensusStats
parCount = traceParFromWithState (justClosuresPar clos) . map (ClosurePtrWithInfo ())
  where
    clos :: WorkerId -> ClosurePtr -> SizedClosure -> ()
              -> DebugM ((), CensusStats, DebugM () -> DebugM ())
    clos _ cp sc _ = do
      return ((), mkCS cp (dcSize sc), id)

asyncParCount :: [ClosurePtr] -> DebugM (AsyncTrace CensusStats)
asyncParCount cps = asyncTraceParFromWithState (justClosuresPar clos) (map (ClosurePtrWithInfo ()) cps)
  where
    clos :: WorkerId -> ClosurePtr -> SizedClosure -> ()
              -> DebugM ((), CensusStats, DebugM () -> DebugM ())
    clos _ cp sc _ = do
      return ((), mkCS cp (dcSize sc), id)

-- | Simple statistics about a heap, total objects, size and maximum object
-- size
count :: [ClosurePtr] -> DebugM CensusStats
count cps = snd <$> runStateT (traceFromM funcs cps) mempty
  where
    funcs = justClosures closAccum

    closAccum  :: ClosurePtr
               -> SizedClosure
               ->  (StateT CensusStats DebugM) ()
               ->  (StateT CensusStats DebugM) ()
    closAccum cp s k = do
      modify' (go cp s)
      k

    go :: ClosurePtr -> SizedClosure -> CensusStats -> CensusStats
    go cp sc cs = mkCS cp (dcSize sc) <> cs


data CountResult = CountResult { countCensus :: !CensusStats } deriving (Show)

-- | An async version of count which the request can be inspected whilst the action is running.
asyncCount :: [ClosurePtr] -> DebugM (IO CountResult, DebugM ())
asyncCount cps = do
  cr <- unsafeLiftIO $ newIORef (CountResult mempty)
  let runCensus = runReaderT (traceFromM funcs cps) cr
  return (readIORef cr, runCensus)

  where
    funcs = justClosures closAccum

    closAccum  :: ClosurePtr
               -> SizedClosure
               ->  (ReaderT (IORef CountResult) DebugM) ()
               ->  (ReaderT (IORef CountResult) DebugM) ()
    closAccum cp s k = do
      cr <- ask
      lift $ unsafeLiftIO $ modifyIORef' cr (go cp s)
      k

    go :: ClosurePtr -> SizedClosure -> CountResult -> CountResult
    go cp sc cr = cr { countCensus = mkCS cp (dcSize sc) <> countCensus cr }