content: {-# LANGUAGE ConstraintKinds #-}
  {-# LANGUAGE NumericUnderscores #-}
  {-# LANGUAGE BangPatterns #-}
  -- | Functions to support the constant space traversal of a heap.
  module GHC.Debug.Trace ( traceFromM, TraceFunctions(..) ) where
  
  import           GHC.Debug.Types
  import GHC.Debug.Client.Monad
  import           GHC.Debug.Client.Query
  
  import qualified Data.IntMap as IM
  import Data.Array.BitArray.IO
  import Control.Monad.Reader
  import Data.IORef
  import Data.Word
  import System.IO
  
  newtype VisitedSet = VisitedSet (IM.IntMap (IOBitArray Word16))
  
  data TraceState = TraceState { visited :: !VisitedSet, n :: !Int }
  
  
  getKeyPair :: ClosurePtr -> (Int, Word16)
  getKeyPair cp =
    let BlockPtr raw_bk = applyBlockMask cp
        bk = fromIntegral raw_bk `div` 8
        offset = getBlockOffset cp `div` 8
    in (bk, fromIntegral offset)
  
  checkVisit :: ClosurePtr -> IORef TraceState -> IO Bool
  checkVisit cp mref = do
    st <- readIORef mref
    let VisitedSet v = visited st
        num_visited = n st
        (bk, offset) = getKeyPair cp
    case IM.lookup bk v of
      Nothing -> do
        na <- newArray (0, fromIntegral (blockMask `div` 8)) False
        writeArray na offset True
        writeIORef mref (TraceState (VisitedSet (IM.insert bk na v)) (num_visited + 1))
        when (num_visited `mod` 10_000 == 0) $ hPutStrLn stderr ("Traced: " ++ show num_visited)
        return False
      Just bm -> do
        res <- readArray bm offset
        unless res (writeArray bm offset True)
        return res
  
  
  
  data TraceFunctions m =
        TraceFunctions { papTrace :: !(GenPapPayload ClosurePtr -> m DebugM ())
        , stackTrace :: !(GenStackFrames ClosurePtr -> m DebugM ())
        , closTrace :: !(ClosurePtr -> SizedClosure -> m DebugM () -> m DebugM ())
        , visitedVal :: !(ClosurePtr -> (m DebugM) ())
        , conDescTrace :: !(ConstrDesc -> m DebugM ())
        }
  
  
  
  
  type C m = (MonadTrans m, Monad (m DebugM))
  
  -- | A generic heap traversal function which will use a small amount of
  -- memory linear in the heap size. Using this function with appropiate
  -- accumulation functions you should be able to traverse quite big heaps in
  -- not a huge amount of memory.
  traceFromM :: C m => TraceFunctions m-> [ClosurePtr] -> m DebugM ()
  traceFromM k cps = do
    st <- lift (unsafeLiftIO (newIORef (TraceState (VisitedSet IM.empty) 1)))
    runReaderT (mapM_ (traceClosureFromM k) cps) st
  {-# INLINE traceFromM #-}
  {-# INLINE traceClosureFromM #-}
  
  traceClosureFromM :: C m
                    => TraceFunctions m
                    -> ClosurePtr
                    -> ReaderT (IORef TraceState) (m DebugM) ()
  traceClosureFromM !k = go
    where
      go cp = do
        mref <- ask
        b <- lift $ lift $ unsafeLiftIO (checkVisit cp mref)
        if b
          then lift $ visitedVal k cp
          else do
          sc <- lift $ lift $ dereferenceClosure cp
          ReaderT $ \st -> closTrace k cp sc
           (runReaderT (() <$ quadtraverse gop gocd gos go sc) st)
  
  
      gos st = do
        st' <- lift $ lift $ dereferenceStack st
        lift $ stackTrace k st'
        () <$ traverse go st'
  
      gocd d = do
        cd <- lift $ lift $ dereferenceConDesc d
        lift $ conDescTrace k cd
  
      gop p = do
        p' <- lift $ lift $ dereferencePapPayload p
        lift $ papTrace k p'
        () <$ traverse go p'
  