Strings.hs

{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE DeriveGeneric #-}
module GHC.Debug.Brick.Action.Strings (
  stringsAction,
  stringTreeMode,
  stringTree,
  renderStringHeader,
  renderStringHeaderPane,
  stringTotalsFromCensus,
  stringHeapBytes,
  StringTotals(..)
) where

import Brick
import Brick.Widgets.Border (borderWithLabel)
import Brick.Widgets.IOTree
import Control.Monad (forM)
import qualified Data.List as List
import qualified Data.Map.Strict as Map
import qualified Data.Ord as Ord
import Data.Semigroup (Sum (..))
import qualified Data.Set as Set
import Data.Text (pack)
import GHC.Debug.Brick.Action.Async
import GHC.Debug.Brick.ClosureDetails
import GHC.Debug.Brick.IOTree
import GHC.Debug.Brick.Lib as GD
import GHC.Debug.Brick.Model
import GHC.Debug.Brick.Render.Closure
import GHC.Debug.Brick.Render.Header (renderHeaderWithSummary)
import GHC.Debug.Brick.Render.Utils
import GHC.Debug.Retainers (SearchLimit)
import qualified GHC.Debug.Strings as GS
import GHC.Generics (Generic, Generically (..))
import Lens.Micro.Platform

stringsAction :: Debuggee -> EventM n OperationalState ()
stringsAction dbg = do
  outside_os <- get

  put (outside_os
      & resetFooter
      & treeMode .~ stringTreeMode (_resultSize outside_os) (stringTree dbg) Map.empty
      )

  runWithAsyncSampleReporter dbg "Counting strings" StringSample
    (stringsAnalysisAsync Nothing dbg)
    (\ result -> do
      sendEvent $ CacheSearch DoStringSearch result
      os <- get
      put (os & resetFooter))

stringTreeMode :: SearchLimit -> IOTree StringCountLine Name -> GS.CensusByStrings -> TreeMode
stringTreeMode limit tree census =
  let
    sorted_res = [(k, v ) | (k, v) <- (List.sortOn (Ord.Down . (Set.size . snd)) (Map.toList census))]
    display_res = maybe id take limit sorted_res
    top_closure = [StringCountLine k (stringHeapBytes k) (length v) v | (k, v) <- display_res]
    totals = stringTotalsFromCensus census
  in
    Strings (renderStringHeader totals) (setIOTreeRoots top_closure tree)

stringTree :: Debuggee -> IOTree StringCountLine Name
stringTree dbg =
  let
    g_children d (StringCountLine _ _ _ cs) = do
      cs' <- run d $ forM (Set.toList cs) $ \c -> do
        c' <- GD.dereferenceClosure c
        return $ ListFullClosure $ Closure c c'
      children' <- traverse (traverse (fillListItem d)) $ zipWith (\n c -> (show @Int n, c)) [0..] cs'
      mapM (\(lbl, child) -> StringFieldLine <$> getClosureDetails d (pack lbl) child) children'
    g_children d (StringFieldLine c) = map StringFieldLine <$> getChildren d c
  in
    mkIOTree dbg [] g_children renderStringLines id

renderStringHeader :: StringTotals -> StringCountLine -> Widget Name
renderStringHeader totals line =
  renderHeaderWithSummary
    (renderStringHeaderPane line)
    (renderTotalsPane totals)

renderStringHeaderPane :: StringCountLine -> Widget Name
renderStringHeaderPane (StringCountLine k l n _) = vBox
  [ labelled "Count     " $ vLimit 1 $ str (show n)
  , labelled "Size      " $ vLimit 1 $ renderBytes l
  , labelled "Total Size" $ vLimit 1 $ renderBytes (n * l)
  , strWrap (take 100 $ show k)
  ]
renderStringHeaderPane (StringFieldLine c) = renderClosureDetails c

renderTotalsPane :: StringTotals -> Widget Name
renderTotalsPane (StringTotals distinct totalCount totalSize) =
  borderWithLabel (txt "Summary") $ hLimit 32 $ vBox
    [ labelled "Distinct Strings" $ vLimit 1 $ str (show $ getSum distinct)
    , labelled "Total Count"      $ vLimit 1 $ str (show $ getSum totalCount)
    , labelled "Total Size"       $ vLimit 1 $ renderBytes $ getSum totalSize
    ]

data StringTotals = StringTotals
  { stDistinct   :: !(Sum Int)
  , stTotalCount :: !(Sum Int)
  , stTotalSize  :: !(Sum Int)
  }
  deriving stock (Show, Eq, Ord, Generic)
  deriving (Semigroup, Monoid) via (Generically StringTotals)

stringTotalsFromCensus :: GS.CensusByStrings -> StringTotals
stringTotalsFromCensus = Map.foldlWithKey' go mempty
  where
    go acc s closures =
      let count = Set.size closures
          perStringBytes = stringHeapBytes s
          total = count * perStringBytes
      in acc <> StringTotals 1 (Sum count) (Sum total)

stringHeapBytes :: String -> Int
stringHeapBytes s = length s * stringConsCellBytes

stringConsCellBytes :: Int
stringConsCellBytes = 32

renderStringLines :: StringCountLine -> [Widget n]
renderStringLines (StringCountLine k l n _) = [strLabel (show n), vSpace, renderBytes l, vSpace, strWrap (take 100 $ show k)]
renderStringLines (StringFieldLine cd) = renderInlineClosureDesc cd