Retainers.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE NamedFieldPuns #-}
module GHC.Debug.Brick.Action.Retainers (
  searchWithCurrentFiltersIncremental,
  RetainerMap,
  retainerTreeMode,
  topRetainersFromRetainerMap,
  retainerTree,
  retainerTreeM,
) where

import Brick.BChan
import Brick.Types
import Brick.Widgets.IOTree
import Control.Monad.IO.Class (liftIO)
import Data.IORef
import qualified Data.List as List
import qualified Data.Map.Strict as Map
import qualified Data.Text as T
import Lens.Micro
import Lens.Micro.Extras

import GHC.Debug.Brick.ClosureDetails
import GHC.Debug.Brick.IOTree (mkIOTree)
import GHC.Debug.Brick.Lib
import GHC.Debug.Brick.Model
import GHC.Debug.Brick.Render.Closure
import GHC.Debug.Brick.UI.Async
import GHC.Debug.Retainers (SearchLimit)
import GHC.Debug.Brick.UI.Filter

searchWithCurrentFiltersIncremental :: Debuggee -> [UIFilter] -> EventM n OperationalState ()
searchWithCurrentFiltersIncremental dbg uifilters = do
  outside_os <- get
  let mClosFilter = uiFiltersToFilter uifilters
  info_map_ref <- liftIO $ newIORef Map.empty
  let tree = retainerTreeM dbg (readIORef info_map_ref)
  let
    add_new_sample :: [ClosureDetails] -> IO ()
    add_new_sample new_retainer_stack_sample = do
      -- We expect this retainer stack to be non-empty.
      -- The first entry is the "Leaf" node in the tree, while the tail of '[ClosureDetails]'
      -- is the stack of retainers of the first 'ClosureDetails'.
      case List.uncons new_retainer_stack_sample of
        Just (k, v) -> do
          let kk = toPtr (_closure k)
          modifyIORef info_map_ref (Map.insert kk (k, v))
        Nothing -> return ()

    ret_add_sample cps = do
      let cps' = zipWith (\n cp -> (T.pack (show n),cp)) [0 :: Int ..] cps
      res <- mapM (completeClosureDetails dbg) cps'
      -- 1. Add the sample to the IOTree
      add_new_sample res
      -- 2. Update the IOTree
      writeBChan (view event_chan outside_os) (NewSampleEvent $ RetainerSample res)

  let searchLimit = _resultSize outside_os

  put $ outside_os
    & treeMode .~ retainerTreeMode searchLimit tree []

  asyncAction "Searching for closures" (liftIO $ retainersOfIncremental searchLimit mClosFilter Nothing dbg ret_add_sample) $ \() -> do
    -- We finished the traversal, finalise the search results
    result <- liftIO $ readIORef info_map_ref
    sendEvent $ CacheSearch (DoRetainerSearch searchLimit uifilters) result
    os <- get
    put (os & resetFooter)

retainerTreeMode :: SearchLimit -> IOTree RetainerNode Name -> [RetainerNode] -> TreeMode
retainerTreeMode _limit tree retainers =
  Retainer (renderClosureDetails . retainerNodeClosure) (setIOTreeRoots retainers tree)

topRetainersFromRetainerMap :: RetainerMap -> [RetainerNode]
topRetainersFromRetainerMap retainerMap =
  fmap (RetainerStackRoot . fst) $ Map.elems retainerMap

retainerTree :: Debuggee -> RetainerMap -> IOTree RetainerNode Name
retainerTree dbg retainerStacks =
  retainerTreeM dbg (pure retainerStacks)

retainerTreeM :: Debuggee -> IO RetainerMap -> IOTree RetainerNode Name
retainerTreeM dbg get_info_map_ref =
  let
    lookup_c dbg' (RetainerStackRoot dc'@(ClosureDetails dc _ _)) = do
      let ptr = toPtr dc
      info_map <- get_info_map_ref
      let results = maybe [] snd $ Map.lookup ptr info_map
      -- We are also looking up the children of the object we are retaining,
      -- and displaying them prior to the retainer stack
      cs <- getChildren dbg' dc'
      results' <- liftIO $ mapM (\(l, c) -> getClosureDetails dbg' l (ListFullClosure c)) (withIndex results)
      return $ map RetainerStackChild (cs ++ results')
    lookup_c dbg' (RetainerStackRoot dc) = fmap (map RetainerStackChild) $ getChildren dbg' dc
    lookup_c dbg' (RetainerStackChild dc) = fmap (map RetainerStackChild) $ getChildren dbg' dc
  in
    mkIOTree dbg [] lookup_c (renderInlineClosureDesc . retainerNodeClosure) id

  where
    withIndex vs = zipWith (\ n v -> (T.show n, _closure v) ) [(1 :: Int) .. ] vs