Thunks.hs

module GHC.Debug.Thunks where

import Control.Monad.Identity (IdentityT(..))
import Control.Monad.RWS.Strict
import qualified Data.Map.Strict as Map
import qualified Data.Map.Monoidal.Strict as MMap

import GHC.Debug.Async
import GHC.Debug.Client.Monad
import GHC.Debug.Client.Query
import GHC.Debug.Profile.Types
import GHC.Debug.Trace
import GHC.Debug.Types
import qualified GHC.Debug.ParTrace as Par

type ThunkCensusBySourceLocation = Map.Map (Maybe SourceInformation) CensusStats

thunkAnalysis :: [ClosurePtr] -> DebugM ThunkCensusBySourceLocation
thunkAnalysis rroots = (\(_, r, _) -> r) <$> runRWST (traceFromM funcs rroots) () (Map.empty)
  where
    funcs = justClosures closAccum

    getSourceLoc c = getSourceInfo (tableId (info (noSize c)))

    closAccum  :: ClosurePtr
               -> SizedClosure
               -> (RWST () () ThunkCensusBySourceLocation DebugM) ()
               -> (RWST () () ThunkCensusBySourceLocation DebugM) ()
    closAccum cp sc k = do
          case (noSize sc) of
            ThunkClosure {} ->  do
              loc <- lift $ getSourceLoc sc
              modify' (Map.insertWith (<>) loc (mkCS cp 1))
              k
            _ -> k

thunkAnalysisIncremental :: [ClosurePtr] -> (Maybe SourceInformation -> CensusStats -> DebugM ()) -> DebugM ()
thunkAnalysisIncremental rroots add_new_sample = runIdentityT $ traceFromM funcs rroots
  where
    funcs = justClosures closAccum

    getSourceLoc c = getSourceInfo (tableId (info (noSize c)))

    closAccum  :: ClosurePtr
                -> SizedClosure
                -> IdentityT DebugM ()
                -> IdentityT DebugM ()
    closAccum cp sc k = do
      case noSize sc of
        ThunkClosure {} ->  do
          loc <- lift $ getSourceLoc sc
          lift $ add_new_sample loc (mkCS cp 1)
          k
        _ -> k

thunkAnalysisAsync :: [ClosurePtr] -> DebugM (AsyncTrace ThunkCensusBySourceLocation)
thunkAnalysisAsync rroots = do
  asyncTrace <- Par.asyncTraceParFromWithState funcs (map (Par.ClosurePtrWithInfo ()) rroots)
  pure $ fmap MMap.getMonoidalMap asyncTrace
  where
    funcs = Par.justClosuresPar closAccum

    getSourceLoc c = getSourceInfo (tableId (info (noSize c)))

    closAccum ::
      Par.WorkerId ->
      ClosurePtr ->
      SizedClosure ->
      () ->
      DebugM ((), MMap.MonoidalMap (Maybe SourceInformation) CensusStats, DebugM () -> DebugM ())
    closAccum _ cp sc _ = do
          case (noSize sc) of
            ThunkClosure {} ->  do
              loc <- getSourceLoc sc
              pure ((), MMap.singleton loc (mkCS cp 1), id)
            _ ->
              pure ((), MMap.empty, id)