BlockCache.hs

{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE BinaryLiterals #-}
{-# LANGUAGE RankNTypes #-}
-- The BlockCache stores the currently fetched blocks
-- and is consulted first to avoid requesting too much
-- from the debuggee. The BlockCache can either be populated
-- via a call to RequestBlocks or on demand on a cache miss.

module GHC.Debug.Client.BlockCache(BlockCache, BlockCacheRequest(..)
                                  , handleBlockReq, emptyBlockCache, bcSize, addBlocks) where

import GHC.Debug.Types.Ptr
import GHC.Debug.Types
import qualified Data.IntMap.Strict as IM
import GHC.Word
import Data.Hashable
import Data.IORef
import Data.Bits
import Data.List (foldl')
import Data.Binary
import Control.Tracer
import Control.Monad (guard)
import Data.Bifunctor

newtype BlockCache = BlockCache (IM.IntMap RawBlock)

instance Binary BlockCache where
  get = BlockCache . IM.fromList <$> get
  put (BlockCache hm) = put (IM.toList hm)

emptyBlockCache :: BlockCache
emptyBlockCache = BlockCache IM.empty

addBlock :: RawBlock -> BlockCache -> BlockCache
addBlock rb@(RawBlock (BlockPtr bp) _ _) (BlockCache bc) =
  BlockCache (IM.insert (fromIntegral bp) rb bc)


addBlocks :: [RawBlock] -> BlockCache -> BlockCache
addBlocks bc bs = foldl' (flip addBlock) bs bc

lookupClosure :: ClosurePtr -> BlockCache -> Maybe RawBlock
lookupClosure cp@(ClosurePtr cpAddr) (BlockCache b) = do
  -- See: Note [Block groups and the block cache]
  (_, rb) <- IM.lookupLE (fromIntegral cpAddr) b
  guard $ rawBlockContains cp rb
  pure rb

bcSize :: BlockCache -> Int
bcSize (BlockCache b) = IM.size b

_bcKeys :: BlockCache -> [ClosurePtr]
_bcKeys (BlockCache b) = map (mkClosurePtr . fromIntegral) (IM.keys b) -- keys of an IntMap are sorted

data BlockCacheRequest a where
  LookupClosure :: ClosurePtr -> BlockCacheRequest RawClosure
  PopulateBlockCache :: BlockCacheRequest [RawBlock]

deriving instance Show (BlockCacheRequest a)
deriving instance Eq (BlockCacheRequest a)

instance Hashable (BlockCacheRequest a) where
  hashWithSalt s (LookupClosure cpt) = s `hashWithSalt` (1 :: Int) `hashWithSalt` cpt
  hashWithSalt s PopulateBlockCache  = s `hashWithSalt` (2 :: Int)

handleBlockReq :: Tracer IO String -> (forall a . Request a -> IO a) -> IORef BlockCache -> BlockCacheRequest resp -> IO resp
handleBlockReq _ do_req ref (LookupClosure cp) = do
  bc <- readIORef ref
  let mrb = lookupClosure cp bc
  rb <- case mrb of
               Nothing -> do
                  rb <- do_req (RequestBlock cp)
                  atomicModifyIORef' ref (\bc' -> (addBlock rb bc', ()))
                  return rb
               Just rb -> do
                 return rb
  return (extractFromBlock cp rb)
handleBlockReq tracer do_req ref PopulateBlockCache = do
  blocks <- do_req RequestAllBlocks
--  mapM_ (\rb -> print ("NEW", rawBlockAddr rb)) blocks
  traceWith tracer $ "Populating block cache with " ++ show (length blocks) ++ " blocks"
  atomicModifyIORef' ref ((,()) . addBlocks blocks)
  return blocks

-- Note [Block groups and the block cache]
-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
--
-- Block groups on the haskell heap normally only consist of one block,
-- but can be longer, eg, for large/pinned objects, compact regions, or for
-- segments on the nonmoving heap.
--
-- The stub always returns an entire block group as a RawBlock. This means
-- we have to be careful when looking up the enclosing block of a closure.
-- It is not enough to use the block mask, since only the head block of each
-- group is held in the cache. And the block mask may give us a non-head block.
-- Instead we find the block group with the greatest starting address less than
-- or equal to the closure pointer.
-- If the enclosing block is cached, then it will be this one.