Strings.hs

{-# LANGUAGE NumericUnderscores #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE BangPatterns #-}
module GHC.Debug.Strings (
    -- * Show analysis
    stringProgram,
    arrWordsProgram,
    -- * Analsis scripts
    CensusByByteString,
    arrWordsAnalysis,
    arrWordsAnalysisAsync,
    CensusByStrings,
    stringAnalysis,
    stringAnalysisAsync,
    decodeString,
    ) where

import Control.Monad.RWS
import Control.Monad.Trans.Reader (ReaderT)
import qualified Control.Monad.Trans.Reader as ReaderT
import Data.ByteString.Lazy (ByteString)
import qualified Data.ByteString.Lazy as BS (length)
import Data.Char
import qualified Data.Foldable as F
import Data.IORef
import Data.List
import qualified Data.Map as Map
import Data.Ord
import qualified Data.Set as S
import qualified Data.Text as T
import qualified Data.Text.IO as T
import GHC.Debug.Async (AsyncTrace, withAsyncTrace, waitForAsyncTraceResultM)
import GHC.Debug.Client
import GHC.Debug.Client.Monad.Class (unsafeLiftIO)
import GHC.Debug.Profile.Types
import GHC.Debug.Trace
import GHC.Debug.Types.Ptr
import Control.Monad.Identity

-- | Find all the strings and then print out how many duplicates there are
stringProgram :: Debuggee -> IO ()
arrWordsProgram :: Debuggee -> IO ()
stringProgram = programX length stringAnalysis
arrWordsProgram = programX (fromIntegral . BS.length) arrWordsAnalysis

programX :: Show a => (a -> Int) -> ([ClosurePtr] -> DebugM (Map.Map a (S.Set b))) -> Debuggee -> IO ()
programX sizeOf anal e = do
  pause e
  res <- runTrace e $ do
    precacheBlocks
    rs <- gcRoots
    res <- anal rs
    return res
  printResult (Map.map (\s -> Count (S.size s)) res)
  printResult (Map.mapWithKey (\k s -> Count (fromIntegral (sizeOf k) * (S.size s))) res)
  return ()

type CensusByStrings = Map.Map String (S.Set ClosurePtr)

-- | Find the parents of Bin nodes
stringAnalysis :: [ClosurePtr] -> DebugM CensusByStrings
stringAnalysis rroots =
  waitForAsyncTraceResultM =<< stringAnalysisAsync rroots

-- | Find the parents of Bin nodes
--
stringAnalysisAsync :: [ClosurePtr] -> DebugM (AsyncTrace CensusByStrings)
stringAnalysisAsync rroots = do
  stringCensus <- unsafeLiftIO $ newIORef Map.empty
  withAsyncTrace (readIORef stringCensus) (ReaderT.runReaderT (traceFromM (funcs stringCensus) rroots) False)
  where
    funcs stringCensus = justClosures (closAccum stringCensus)

    -- First time we have visited a closure
    closAccum  :: IORef CensusByStrings
                -> ClosurePtr
                -> SizedClosure
                -> ReaderT Bool DebugM ()
                -> ReaderT Bool DebugM ()
    closAccum ref cp sc k = do
      case noSize sc of
        ConstrClosure _ _ _ _ cd -> do
          cd' <- lift $ dereferenceConDesc cd
          case cd' of
            ConstrDesc _ _ cd2 | cd2 == ":" -> do
              process cp sc
            _ -> local (const False) k
        _  -> local (const False) k
      where
        process :: ClosurePtr -> SizedClosure
                -> ReaderT Bool DebugM ()
        process p_cp clos = do
          clos' <- lift $ hextraverse pure pure pure dereferenceConDesc return return (noSize clos)
          checked <- lift $ check_bin clos'
          if checked
            then do
              parent_is_cons <- ask
              if parent_is_cons
                then local (const True) k
                else do
                  ds <- lift $ decodeString p_cp
                  lift $ unsafeLiftIO $ modifyIORef' ref (Map.insertWith (<>) ds (S.singleton p_cp))
                  local (const True) k
            else local (const False) k

        process_2 p_cp = do
          cp' <- dereferenceClosure p_cp
          case noSize cp' of
            (ConstrClosure _ _ _ _ cd) -> do
              (ConstrDesc _ _ cn) <- dereferenceConDesc cd
              return (cn == "C#")
            (IndClosure _ _ i) ->
              process_2 i
            _ -> return False

        check_bin (ConstrClosure _ _ [h,_] _ (ConstrDesc _ _ ":")) = process_2 h
        check_bin (IndClosure _ _ i) = do
          sizedI <- dereferenceClosure i
          clos' <- hextraverse pure pure pure dereferenceConDesc return return (noSize sizedI)
          check_bin clos'
        check_bin _ = return False


decodeString :: ClosurePtr -> DebugM String
decodeString cp = do
  cp' <- dereferenceClosure cp
  case noSize cp' of
    (IndClosure _ _ i) -> decodeString i
    (ConstrClosure _ _ [p,ps] _ _) -> do
      go p ps
    _ -> return []
    where
  go headp tailp = do
    cp'' <- dereferenceClosure headp
    case noSize cp'' of
      (IndClosure _ _ i) -> go i tailp
      (ConstrClosure _ _ _ [w] _) -> do
        (chr (fromIntegral w):) <$> decodeString tailp
      _ -> return []


printResult :: Show a => Map.Map a Count -> IO [a]
printResult m = do
  putStrLn $ "TOTAL: " ++ show total
  mapM_ show_line top10
  return (map fst top10)
  where
    show_line (k, Count v) = T.putStrLn (T.pack (show k) <> ": " <> T.pack (show v))
    top10 = take 1000 $ reverse (sortBy (comparing snd) (Map.toList m))
    total = F.fold (Map.elems m)

type CensusByByteString = Map.Map ByteString (S.Set ClosurePtr)

-- | Find how many distinct ArrWords there are
arrWordsAnalysis :: [ClosurePtr] -> DebugM CensusByByteString
arrWordsAnalysis rroots =
  waitForAsyncTraceResultM =<< arrWordsAnalysisAsync rroots

arrWordsAnalysisAsync :: [ClosurePtr] -> DebugM (AsyncTrace CensusByByteString)
arrWordsAnalysisAsync rroots = do
  byteStringCensus <- unsafeLiftIO $ newIORef Map.empty
  withAsyncTrace (readIORef byteStringCensus) (runIdentityT (traceFromM (funcs byteStringCensus) rroots))
  where
    funcs ref = justClosures (closAccum ref)

    -- First time we have visited a closure
    closAccum  ::
      IORef CensusByByteString ->
      ClosurePtr ->
      SizedClosure ->
      IdentityT DebugM () ->
      IdentityT DebugM ()
    closAccum ref cp sc k = do
          case (noSize sc) of
            ArrWordsClosure _ _ _ p ->  do
              lift $ unsafeLiftIO $ modifyIORef' ref (Map.insertWith (<>) (arrWordsBS p) (S.singleton cp))
              k
            _ -> k