{-# LANGUAGE BangPatterns #-}
module GHC.Debug.CostCentres
  ( findAllChildrenOfCC
  , findExactlyByCC
  , findAllCCSPayloads
  , traverseCCSPayloads
  -- * Helper functions for working with `IndexTable`'s
  , flattenIndexTable
  , traverseIndexTable
  , foldIndexTable
  -- * Efficient representation of CCSPtr sets
  , CCSSet(..)
  , memberCCSSet
  ) where

import qualified Data.Set as Set
import Data.IntSet (IntSet)
import qualified Data.IntSet as IntSet

import GHC.Debug.Client
import GHC.Debug.Types.Ptr (CCSPtr(..))
import Data.Coerce (coerce)

newtype CCSSet = CCSSet IntSet

memberCCSSet :: CCSPtr -> CCSSet -> Bool
memberCCSSet :: CCSPtr -> CCSSet -> Bool
memberCCSSet (CCSPtr Word64
k) CCSSet
set = Key -> IntSet -> Bool
IntSet.member (Word64 -> Key
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
k) (CCSSet -> IntSet
forall a b. Coercible a b => a -> b
coerce CCSSet
set)

-- | Find all Cost Centre Stacks that reference precisely the cost centre with the given id.
findExactlyByCC :: (CCPayload -> Bool) -> DebugM (Set.Set CCSPtr)
findExactlyByCC :: (CCPayload -> Bool) -> DebugM (Set CCSPtr)
findExactlyByCC CCPayload -> Bool
isRelevantCC = do
  CCSPtr
ccsMain <- DebugM CCSPtr
requestCCSMain
  Set CCSPtr -> CCSPtr -> DebugM (Set CCSPtr)
collectNode Set CCSPtr
forall a. Set a
Set.empty CCSPtr
ccsMain
  where
    collectNode :: Set.Set CCSPtr -> CCSPtr -> DebugM (Set.Set CCSPtr)
    collectNode :: Set CCSPtr -> CCSPtr -> DebugM (Set CCSPtr)
collectNode !Set CCSPtr
seen !CCSPtr
ccsPtr = do
      CCSPayload
ccsPl <- CCSPtr -> DebugM CCSPayload
dereferenceCCS CCSPtr
ccsPtr
      CCPayload
ccPl <- CCPtr -> DebugM CCPayload
dereferenceCC (CCSPayload -> CCPtr
forall ccsPtr ccPtr. GenCCSPayload ccsPtr ccPtr -> ccPtr
ccsCc CCSPayload
ccsPl)
      let newSeen :: Set CCSPtr
newSeen = if CCPayload -> Bool
isRelevantCC CCPayload
ccPl
            then CCSPtr
ccsPtr CCSPtr -> Set CCSPtr -> Set CCSPtr
forall a. Ord a => a -> Set a -> Set a
`Set.insert` Set CCSPtr
seen
            else Set CCSPtr
seen
      (CCPtr -> CCSPtr -> Bool -> Set CCSPtr -> DebugM (Set CCSPtr))
-> Set CCSPtr -> Maybe IndexTablePtr -> DebugM (Set CCSPtr)
forall a.
(CCPtr -> CCSPtr -> Bool -> a -> DebugM a)
-> a -> Maybe IndexTablePtr -> DebugM a
foldIndexTable (\CCPtr
_ CCSPtr
ptr Bool
backEdge !Set CCSPtr
seen' -> do
        if Bool
backEdge
          then Set CCSPtr -> DebugM (Set CCSPtr)
forall a. a -> DebugM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Set CCSPtr
seen'
          else Set CCSPtr -> CCSPtr -> DebugM (Set CCSPtr)
collectNode Set CCSPtr
seen' CCSPtr
ptr)
        Set CCSPtr
newSeen
        (CCSPayload -> Maybe IndexTablePtr
forall ccsPtr ccPtr.
GenCCSPayload ccsPtr ccPtr -> Maybe IndexTablePtr
ccsIndexTable CCSPayload
ccsPl)

-- | Find all cost centre stack parts that are transitively children of the cost
-- centre with the given id.
findAllChildrenOfCC :: (CCPayload -> Bool) -> DebugM (Set.Set CCSPtr)
findAllChildrenOfCC :: (CCPayload -> Bool) -> DebugM (Set CCSPtr)
findAllChildrenOfCC CCPayload -> Bool
isRelevantCC = do
  CCSPtr
ccsMain <- DebugM CCSPtr
requestCCSMain
  Set CCSPtr -> CCSPtr -> DebugM (Set CCSPtr)
findCostCentre Set CCSPtr
forall a. Set a
Set.empty CCSPtr
ccsMain
  where
    findCostCentre :: Set.Set CCSPtr -> CCSPtr -> DebugM (Set.Set CCSPtr)
    findCostCentre :: Set CCSPtr -> CCSPtr -> DebugM (Set CCSPtr)
findCostCentre !Set CCSPtr
seen !CCSPtr
ccsPtr = do
      CCSPayload
ccsPl <- CCSPtr -> DebugM CCSPayload
dereferenceCCS CCSPtr
ccsPtr
      CCPayload
ccPl <- CCPtr -> DebugM CCPayload
dereferenceCC (CCSPayload -> CCPtr
forall ccsPtr ccPtr. GenCCSPayload ccsPtr ccPtr -> ccPtr
ccsCc CCSPayload
ccsPl)
      if CCPayload -> Bool
isRelevantCC CCPayload
ccPl
        then Set CCSPtr -> CCSPtr -> DebugM (Set CCSPtr)
collectNodes Set CCSPtr
seen CCSPtr
ccsPtr
        else
          (CCPtr -> CCSPtr -> Bool -> Set CCSPtr -> DebugM (Set CCSPtr))
-> Set CCSPtr -> Maybe IndexTablePtr -> DebugM (Set CCSPtr)
forall a.
(CCPtr -> CCSPtr -> Bool -> a -> DebugM a)
-> a -> Maybe IndexTablePtr -> DebugM a
foldIndexTable  (\CCPtr
_ CCSPtr
ptr Bool
backEdge !Set CCSPtr
seen' -> do
            if Bool
backEdge
              then Set CCSPtr -> DebugM (Set CCSPtr)
forall a. a -> DebugM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Set CCSPtr
seen'
              else Set CCSPtr -> CCSPtr -> DebugM (Set CCSPtr)
findCostCentre Set CCSPtr
seen' CCSPtr
ptr)
            Set CCSPtr
seen
            (CCSPayload -> Maybe IndexTablePtr
forall ccsPtr ccPtr.
GenCCSPayload ccsPtr ccPtr -> Maybe IndexTablePtr
ccsIndexTable CCSPayload
ccsPl)

    collectNodes :: Set.Set CCSPtr -> CCSPtr -> DebugM (Set.Set CCSPtr)
    collectNodes :: Set CCSPtr -> CCSPtr -> DebugM (Set CCSPtr)
collectNodes !Set CCSPtr
seen !CCSPtr
ccsPtr = do
      CCSPayload
ccsPl <- CCSPtr -> DebugM CCSPayload
dereferenceCCS CCSPtr
ccsPtr
      (CCPtr -> CCSPtr -> Bool -> Set CCSPtr -> DebugM (Set CCSPtr))
-> Set CCSPtr -> Maybe IndexTablePtr -> DebugM (Set CCSPtr)
forall a.
(CCPtr -> CCSPtr -> Bool -> a -> DebugM a)
-> a -> Maybe IndexTablePtr -> DebugM a
foldIndexTable  (\CCPtr
_ CCSPtr
ptr Bool
backEdge !Set CCSPtr
seen' -> do
        let seen'' :: Set CCSPtr
seen'' = CCSPtr
ptr CCSPtr -> Set CCSPtr -> Set CCSPtr
forall a. Ord a => a -> Set a -> Set a
`Set.insert` Set CCSPtr
seen'
        if Bool
backEdge
          then Set CCSPtr -> DebugM (Set CCSPtr)
forall a. a -> DebugM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Set CCSPtr
seen''
          else Set CCSPtr -> CCSPtr -> DebugM (Set CCSPtr)
collectNodes Set CCSPtr
seen'' CCSPtr
ptr)
        (CCSPtr
ccsPtr CCSPtr -> Set CCSPtr -> Set CCSPtr
forall a. Ord a => a -> Set a -> Set a
`Set.insert` Set CCSPtr
seen)
        (CCSPayload -> Maybe IndexTablePtr
forall ccsPtr ccPtr.
GenCCSPayload ccsPtr ccPtr -> Maybe IndexTablePtr
ccsIndexTable CCSPayload
ccsPl)

findAllCCSPayloads :: DebugM CCSSet
findAllCCSPayloads :: DebugM CCSSet
findAllCCSPayloads = do
  CCSPtr
ccsMain <- DebugM CCSPtr
requestCCSMain
  IntSet -> CCSSet
CCSSet (IntSet -> CCSSet) -> DebugM IntSet -> DebugM CCSSet
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IntSet -> CCSPtr -> DebugM IntSet
collectNodes IntSet
IntSet.empty CCSPtr
ccsMain
  where
    collectNodes :: IntSet.IntSet -> CCSPtr -> DebugM (IntSet.IntSet)
    collectNodes :: IntSet -> CCSPtr -> DebugM IntSet
collectNodes !IntSet
seen ccsPtr :: CCSPtr
ccsPtr@(CCSPtr Word64
w_) = do
      CCSPayload
ccsPl <- CCSPtr -> DebugM CCSPayload
dereferenceCCS CCSPtr
ccsPtr
      (CCPtr -> CCSPtr -> Bool -> IntSet -> DebugM IntSet)
-> IntSet -> Maybe IndexTablePtr -> DebugM IntSet
forall a.
(CCPtr -> CCSPtr -> Bool -> a -> DebugM a)
-> a -> Maybe IndexTablePtr -> DebugM a
foldIndexTable  (\CCPtr
_ ptr :: CCSPtr
ptr@(CCSPtr Word64
w) Bool
backEdge !IntSet
seen' -> do
        let seen'' :: IntSet
seen'' = Word64 -> Key
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
w Key -> IntSet -> IntSet
`IntSet.insert` IntSet
seen'
        if Bool
backEdge
          then IntSet -> DebugM IntSet
forall a. a -> DebugM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure IntSet
seen''
          else IntSet -> CCSPtr -> DebugM IntSet
collectNodes IntSet
seen'' CCSPtr
ptr)
        (Word64 -> Key
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
w_ Key -> IntSet -> IntSet
`IntSet.insert` IntSet
seen)
        (CCSPayload -> Maybe IndexTablePtr
forall ccsPtr ccPtr.
GenCCSPayload ccsPtr ccPtr -> Maybe IndexTablePtr
ccsIndexTable CCSPayload
ccsPl)

traverseCCSPayloads :: DebugM ()
traverseCCSPayloads :: DebugM ()
traverseCCSPayloads = do
  CCSPtr
ccsMain <- DebugM CCSPtr
requestCCSMain
  CCSPtr -> DebugM ()
collectNodes CCSPtr
ccsMain
  where
    collectNodes :: CCSPtr -> DebugM ()
    collectNodes :: CCSPtr -> DebugM ()
collectNodes CCSPtr
ccsPtr = do
      CCSPayload
ccsPl <- CCSPtr -> DebugM CCSPayload
dereferenceCCS CCSPtr
ccsPtr
      (CCPtr -> CCSPtr -> Bool -> () -> DebugM ())
-> () -> Maybe IndexTablePtr -> DebugM ()
forall a.
(CCPtr -> CCSPtr -> Bool -> a -> DebugM a)
-> a -> Maybe IndexTablePtr -> DebugM a
foldIndexTable  (\CCPtr
_ CCSPtr
ptr Bool
backEdge () -> do
        if Bool
backEdge
          then () -> DebugM ()
forall a. a -> DebugM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          else CCSPtr -> DebugM ()
collectNodes CCSPtr
ptr)
        ()
        (CCSPayload -> Maybe IndexTablePtr
forall ccsPtr ccPtr.
GenCCSPayload ccsPtr ccPtr -> Maybe IndexTablePtr
ccsIndexTable CCSPayload
ccsPl)

traverseIndexTable :: Maybe IndexTablePtr -> (CCPtr -> CCSPtr -> Bool -> DebugM a) -> DebugM [a]
traverseIndexTable :: forall a.
Maybe IndexTablePtr
-> (CCPtr -> CCSPtr -> Bool -> DebugM a) -> DebugM [a]
traverseIndexTable Maybe IndexTablePtr
Nothing CCPtr -> CCSPtr -> Bool -> DebugM a
_ = [a] -> DebugM [a]
forall a. a -> DebugM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
traverseIndexTable (Just IndexTablePtr
ptr) CCPtr -> CCSPtr -> Bool -> DebugM a
f = do
  IndexTable
idxTable <- IndexTablePtr -> DebugM IndexTable
dereferenceIndexTable IndexTablePtr
ptr
  a
x <- CCPtr -> CCSPtr -> Bool -> DebugM a
f (IndexTable -> CCPtr
itCostCentre IndexTable
idxTable) (IndexTable -> CCSPtr
itCostCentreStack IndexTable
idxTable) (IndexTable -> Bool
itBackEdge IndexTable
idxTable)
  [a]
rest <- Maybe IndexTablePtr
-> (CCPtr -> CCSPtr -> Bool -> DebugM a) -> DebugM [a]
forall a.
Maybe IndexTablePtr
-> (CCPtr -> CCSPtr -> Bool -> DebugM a) -> DebugM [a]
traverseIndexTable (IndexTable -> Maybe IndexTablePtr
itNext IndexTable
idxTable) CCPtr -> CCSPtr -> Bool -> DebugM a
f
  pure $ a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
rest

foldIndexTable :: (CCPtr -> CCSPtr -> Bool -> a -> DebugM a) -> a -> Maybe IndexTablePtr -> DebugM a
foldIndexTable :: forall a.
(CCPtr -> CCSPtr -> Bool -> a -> DebugM a)
-> a -> Maybe IndexTablePtr -> DebugM a
foldIndexTable CCPtr -> CCSPtr -> Bool -> a -> DebugM a
_f a
acc Maybe IndexTablePtr
Nothing = a -> DebugM a
forall a. a -> DebugM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
acc
foldIndexTable CCPtr -> CCSPtr -> Bool -> a -> DebugM a
f a
acc (Just IndexTablePtr
ptr) = do
  IndexTable
idxTable <- IndexTablePtr -> DebugM IndexTable
dereferenceIndexTable IndexTablePtr
ptr
  a
acc' <- CCPtr -> CCSPtr -> Bool -> a -> DebugM a
f (IndexTable -> CCPtr
itCostCentre IndexTable
idxTable) (IndexTable -> CCSPtr
itCostCentreStack IndexTable
idxTable) (IndexTable -> Bool
itBackEdge IndexTable
idxTable) a
acc
  (CCPtr -> CCSPtr -> Bool -> a -> DebugM a)
-> a -> Maybe IndexTablePtr -> DebugM a
forall a.
(CCPtr -> CCSPtr -> Bool -> a -> DebugM a)
-> a -> Maybe IndexTablePtr -> DebugM a
foldIndexTable CCPtr -> CCSPtr -> Bool -> a -> DebugM a
f a
acc' (IndexTable -> Maybe IndexTablePtr
itNext IndexTable
idxTable)

-- | Flatten an optional index table pointer into a list of CCS Payloads.
flattenIndexTable :: Maybe IndexTablePtr -> DebugM [CCSPayload]
flattenIndexTable :: Maybe IndexTablePtr -> DebugM [CCSPayload]
flattenIndexTable Maybe IndexTablePtr
root = Maybe IndexTablePtr
-> (CCPtr -> CCSPtr -> Bool -> DebugM CCSPayload)
-> DebugM [CCSPayload]
forall a.
Maybe IndexTablePtr
-> (CCPtr -> CCSPtr -> Bool -> DebugM a) -> DebugM [a]
traverseIndexTable Maybe IndexTablePtr
root (\CCPtr
_ CCSPtr
ccsPtr Bool
_ -> CCSPtr -> DebugM CCSPayload
dereferenceCCS CCSPtr
ccsPtr)