ArrWords.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE OverloadedStrings #-}
module GHC.Debug.Brick.Action.ArrWords (
  arrWordsAction,
  renderArrWordsHeader,
  arrWordsTreeMode,
  arrWordsTree,
) where

import Brick
import Brick.Widgets.Border
import Control.Monad (forM)
import Lens.Micro.Platform
import Data.Text (pack)
import qualified Data.ByteString.Lazy as BS
import qualified Data.Set as S

import Brick.Widgets.IOTree (IOTree, setIOTreeRoots)
import qualified Data.List as List
import qualified Data.Map.Strict as Map
import qualified Data.Ord as Ord
import qualified Data.Set as Set
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.Histogram
import GHC.Debug.Brick.Render.Utils
import GHC.Debug.Retainers (SearchLimit)
import qualified GHC.Debug.Strings as GD

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

    put
      (outside_os
        & resetFooter
        & treeMode .~ arrWordsTreeMode (_resultSize outside_os) (arrWordsTree dbg) Map.empty
      )

    runWithAsyncSampleReporter dbg "Counting ARR_WORDS" ArrWordsSample
      (arrWordsAnalysisAsync Nothing dbg)
      (\ result -> do
        sendEvent $ CacheSearch DoArrWordsSearch result
        os <- get
        put (os & resetFooter))

arrWordsTreeMode :: SearchLimit -> IOTree ArrWordsLine Name -> Map.Map BS.ByteString (Set.Set ClosurePtr) -> TreeMode
arrWordsTreeMode limit tree payload =
  let
    sorted_res = [(k, v) | (k, v) <- List.sortOn ((\(k, v) -> Ord.Down $ fromIntegral (BS.length k) * Set.size v)) (Map.toList payload)]
    display_res = maybe id take limit sorted_res
    top_closure = [ArrWordsCountLine k (fromIntegral (BS.length k)) (Set.size v) v | (k, v) <- display_res]
  in
    ArrWords (renderArrWordsHeader payload) (setIOTreeRoots top_closure tree)

arrWordsTree :: Debuggee -> IOTree ArrWordsLine Name
arrWordsTree dbg =
  let
    g_children d (ArrWordsCountLine _ _ _ cs) = do
        cs' <- run d $ forM (S.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) -> ArrWordsFieldLine <$> getClosureDetails d (pack lbl) child) children'
    g_children d (ArrWordsFieldLine c) = map ArrWordsFieldLine <$> getChildren d c
  in
    mkIOTree dbg [] g_children renderArrWordsLines id

renderArrWordsHeader :: GD.CensusByByteString -> ArrWordsLine -> Widget Name
renderArrWordsHeader census c =
  renderHeaderWithSummary (renderHeaderPane c) (renderArrWordsHistogram census)

renderArrWordsHistogram :: GD.CensusByByteString -> Widget Name
renderArrWordsHistogram census =
  borderWithLabel (txt "Histogram") $ hLimit 100 $ mkSimpleArrWordsHistogram census

mkSimpleArrWordsHistogram :: GD.CensusByByteString -> Widget Name
mkSimpleArrWordsHistogram census =
  histogram 8 (concatMap (\(k, bs) -> let sz = BS.length k in replicate (length bs) (Size (fromIntegral sz))) (Map.toList census))

renderArrWordsLines :: ArrWordsLine -> [Widget n]
renderArrWordsLines (ArrWordsCountLine k l n _) = [strLabel (show n), vSpace, renderBytes l, vSpace, strWrap (take 100 $ show k)]
renderArrWordsLines (ArrWordsFieldLine cd) = renderInlineClosureDesc cd

renderHeaderPane :: ArrWordsLine -> Widget Name
renderHeaderPane (ArrWordsCountLine b 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 b)
  ]
renderHeaderPane (ArrWordsFieldLine c) = renderClosureDetails c