Retainers.hs

{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE FlexibleContexts #-}
-- | Functions for computing retainers
module GHC.Debug.Retainers
  ( SearchLimit
  , findRetainers
  , findRetainersIncremental
  , findRetainersOf
  , findRetainersOfConstructor
  , findRetainersOfConstructorExact
  , findRetainersOfInfoTable
  , addLocationToStack
  , displayRetainerStack
  , addLocationToStack'
  , displayRetainerStack'
  , findRetainersOfArrWords
  , EraRange(..)
  , profHeaderInEraRange
  , ClosureFilter(..)
  , profHeaderReferencesCCS
  , findRetainersOfEra) where

import Prelude hiding (filter)
import GHC.Debug.Client
import Control.Monad.State
import GHC.Debug.Trace
import GHC.Debug.Types.Graph
import Control.Monad

import qualified Data.Set as Set
import Control.Monad.RWS
import Data.Word

-- | Optional limit for the number of results of a search.
-- If no search limit is given, the search is unlimited.
type SearchLimit = Maybe Int

-- | Check the search limit, and if it has been reached, then stop the execution.
ifSearchLimitNotReached :: (MonadState s m) => (s -> SearchLimit) -> m () -> m ()
ifSearchLimitNotReached getLimit act = do
  limit <- get
  case getLimit limit of
    Just 0 -> pure ()
    Nothing ->
      act
    Just _ ->
      act

-- | Indicate that a result value has been found and the search limit, if there is any,
-- needs to be updated.
addedOneSearchResult :: (MonadState s m) => ((SearchLimit -> SearchLimit) -> s -> s) -> m ()
addedOneSearchResult updateLimit = do
  st <- get
  put (updateLimit (subtract 1 <$>) st)

data EraRange
  = EraRange { startEra :: Word64, endEra :: Word64} -- inclusive
  deriving (Eq, Ord, Show)

inEraRange :: Word64 -> Maybe EraRange -> Bool
inEraRange _ Nothing = True
inEraRange n (Just (EraRange s e)) = s <= n && n <= e

profHeaderReferencesCCS :: Maybe ProfHeaderWithPtr -> Set.Set CCSPtr -> Bool
profHeaderReferencesCCS Nothing _ = False
profHeaderReferencesCCS (Just profHeader) f = ccs profHeader `Set.member` f

profHeaderInEraRange :: Maybe (ProfHeader a) -> Maybe EraRange -> Bool
profHeaderInEraRange Nothing _ = True
profHeaderInEraRange (Just ph) eras
  = case hp ph of
      EraWord w -> w `inEraRange` eras
      _ -> True -- Don't filter if no era profiling

data ClosureFilter
 = ConstructorDescFilter (ConstrDesc -> Bool)
 | InfoFilter (StgInfoTable -> Bool)
 | InfoPtrFilter (InfoTablePtr -> Bool)
 | InfoSourceFilter (SourceInformation -> Bool)
 | SizeFilter (Size -> Bool)
 | ProfHeaderFilter (Maybe ProfHeaderWithPtr -> Bool)
 | AddressFilter (ClosurePtr -> Bool)
 | AndFilter ClosureFilter ClosureFilter
 | OrFilter ClosureFilter ClosureFilter
 | NotFilter ClosureFilter
 | PureFilter Bool

matchesFilter :: ClosureFilter -> ClosurePtr -> SizedClosure -> [ClosurePtr] -> DebugM Bool
matchesFilter filter ptr sc parents = case filter of
  ConstructorDescFilter p -> case noSize sc of
    ConstrClosure _ _ _ _ cd -> do
      cd' <- dereferenceConDesc cd
      return $ p cd'
    _ -> pure False
  InfoFilter p -> pure $ p (decodedTable (info (noSize sc)))
  InfoPtrFilter p -> pure $ p (tableId (info (noSize sc)))
  InfoSourceFilter p -> do
    loc <- getSourceInfo (tableId (info (noSize sc)))
    case loc of
      Nothing -> return False
      Just cur_loc -> pure $ p cur_loc
  SizeFilter p -> pure $ p (dcSize sc)
  ProfHeaderFilter p -> pure $ p (profHeader $ noSize sc)
  AddressFilter p -> pure $ p ptr
  AndFilter f1 f2 -> do
    r1 <- matchesFilter f1 ptr sc parents
    case r1 of
      False -> pure False
      True -> matchesFilter f2 ptr sc parents
  OrFilter f1 f2 -> do
    r1 <- matchesFilter f1 ptr sc parents
    case r1 of
      True -> pure True
      False -> matchesFilter f2 ptr sc parents
  NotFilter f1  -> do
    r1 <- matchesFilter f1 ptr sc parents
    pure (not r1)
  PureFilter b -> pure b

findRetainersOf :: SearchLimit
                -> [ClosurePtr]
                -> [ClosurePtr]
                -> DebugM [[ClosurePtr]]
findRetainersOf limit cps bads =
  findRetainers limit (AddressFilter (`Set.member` bad_set)) cps
  where
    bad_set = Set.fromList bads

findRetainersOfConstructor :: SearchLimit
                           -> [ClosurePtr] -> String -> DebugM [[ClosurePtr]]
findRetainersOfConstructor limit rroots con_name =
  findRetainers limit (ConstructorDescFilter ((== con_name) . name)) rroots

findRetainersOfConstructorExact
  :: SearchLimit
  -> [ClosurePtr] -> String -> DebugM [[ClosurePtr]]
findRetainersOfConstructorExact limit rroots clos_name =
  findRetainers limit (InfoSourceFilter ((== clos_name) . infoName)) rroots

findRetainersOfEra
  :: SearchLimit
  -> EraRange
  -> [ClosurePtr] -> DebugM [[ClosurePtr]]
findRetainersOfEra limit eras rroots =
  findRetainers limit filter rroots
  where
    filter = ProfHeaderFilter (`profHeaderInEraRange` (Just eras))

findRetainersOfArrWords
  :: SearchLimit
  -> [ClosurePtr] -> Size -> DebugM [[ClosurePtr]]
findRetainersOfArrWords limit rroots lim =
  findRetainers limit filter rroots
  where
    -- TODO : this is the size of the entire closure, not the size of the ArrWords
    filter = AndFilter (InfoFilter ((== ARR_WORDS) . tipe))
                       (SizeFilter (>= lim))

findRetainersOfInfoTable
  :: SearchLimit
  -> [ClosurePtr] -> InfoTablePtr -> DebugM [[ClosurePtr]]
findRetainersOfInfoTable limit rroots info_ptr =
  findRetainers limit (InfoPtrFilter (== info_ptr)) rroots

-- | From the given roots, find any path to one of the given pointers.
-- Note: This function can be quite slow! The first argument is a limit to
-- how many paths to find. You should normally set this to a small number
-- such as 10.
findRetainers :: SearchLimit
  -> ClosureFilter
  -> [ClosurePtr] -> DebugM [[ClosurePtr]]
findRetainers limit filter rroots = (\(_, r, _) -> snd r) <$> runRWST (traceFromM funcs rroots) [] (limit, [])
  where
    funcs = justClosures closAccum
    -- Add clos
    closAccum  :: ClosurePtr
               -> SizedClosure
               -> RWST [ClosurePtr] () (SearchLimit, [[ClosurePtr]]) DebugM ()
               -> RWST [ClosurePtr] () (SearchLimit, [[ClosurePtr]]) DebugM ()
    closAccum _ (noSize -> WeakClosure {}) _ = return ()
    closAccum cp sc k = ifSearchLimitNotReached (\ (searchLimit, _) -> searchLimit) $ do
      ctx <- ask
      b <- lift $ matchesFilter filter cp sc ctx
      when b $ do
        addOne (cp: ctx)
        addedOneSearchResult ( \ upd (searchLimit, r) -> ( upd searchLimit, r))
      local (cp:) k

    addOne :: [ClosurePtr] -> RWST [ClosurePtr] () (SearchLimit, [[ClosurePtr]]) DebugM ()
    addOne cp = do
      modify' (\ (searchLimit, cps) -> ( searchLimit, cp : cps))

-- |  From the given roots, find retainer paths to closures which satisfy the filter.
-- The callback is called immediately when each sample is found.
findRetainersIncremental :: SearchLimit
  -> ClosureFilter
  -> [ClosurePtr]
  -> ([ClosurePtr] -> DebugM ()) -- ^ What to do with each sample when we find it.
  -> DebugM ()
findRetainersIncremental limit filter rroots add_new_sample = () <$ execRWST (traceFromM funcs rroots) [] limit
  where
    funcs = justClosures closAccum
    -- Add clos
    closAccum  :: ClosurePtr
               -> SizedClosure
               -> RWST [ClosurePtr] () SearchLimit DebugM ()
               -> RWST [ClosurePtr] () SearchLimit DebugM ()
    closAccum _ (noSize -> WeakClosure {}) _ = return ()
    closAccum cp sc k = ifSearchLimitNotReached id $ do
      ctx <- ask
      b <- lift $ matchesFilter filter cp sc ctx
      when b $ do
        lift $ add_new_sample (cp : ctx)
        addedOneSearchResult ($)
      local (cp:) k

addLocationToStack :: [ClosurePtr] -> DebugM [(SizedClosureP, Maybe SourceInformation)]
addLocationToStack r = do
  cs <- dereferenceClosures r
  cs' <- mapM dereferenceToClosurePtr cs
  locs <- mapM getSourceLoc cs'
  return $ (zip cs' locs)
  where
    getSourceLoc c = getSourceInfo (tableId (info (noSize c)))

addLocationToStack' :: [ClosurePtr] -> DebugM [(ClosurePtr, SizedClosureP, Maybe SourceInformation)]
addLocationToStack' r = do
  cs <- dereferenceClosures r
  cs' <- mapM dereferenceToClosurePtr cs
  locs <- mapM getSourceLoc cs'
  return $ (zip3 r cs' locs)
  where
    getSourceLoc c = getSourceInfo (tableId (info (noSize c)))

displayRetainerStack :: [(String, [(SizedClosureP, Maybe SourceInformation)])] -> IO ()
displayRetainerStack rs = do
      let disp (d, l) =
            (ppClosure  (\_ -> show) 0 . noSize $ d) ++  " <" ++ maybe "nl" tdisplay l ++ ">"
            where
              tdisplay sl = infoName sl ++ ":" ++ infoType sl ++ ":" ++ infoModule sl ++ ":" ++ infoPosition sl
          do_one k (l, stack) = do
            putStrLn (show k ++ "-------------------------------------")
            print l
            mapM (putStrLn . disp) stack
      zipWithM_ do_one [0 :: Int ..] rs

displayRetainerStack' :: [(String, [(ClosurePtr, SizedClosureP, Maybe SourceInformation)])] -> IO ()
displayRetainerStack' rs = do
      let disp (p, d, l) =
            show p ++ ": " ++ (ppClosure  (\_ -> show) 0 . noSize $ d) ++  " <" ++ maybe "nl" tdisplay l ++ ">"
            where
              tdisplay sl = infoName sl ++ ":" ++ infoType sl ++ ":" ++ infoModule sl ++ ":" ++ infoPosition sl
          do_one k (l, stack) = do
            putStrLn (show k ++ "-------------------------------------")
            print l
            mapM (putStrLn . disp) stack
      zipWithM_ do_one [0 :: Int ..] rs