ghc-heap: Allow more control about decoding CCS fieldswip/t19038

We have to be careful not to decode too much, too eagerly, as in
ghc-debug this will lead to references to memory locations outside of
the currently copied closure.

Fixes #19038

5 files changed, 34 insertions, 15 deletions

diff --git a/libraries/ghc-heap/GHC/Exts/Heap.hs b/libraries/ghc-heap/GHC/Exts/Heap.hs
index 2dfe788406..10a9ea8be9 100644
--- a/libraries/ghc-heap/GHC/Exts/Heap.hs
+++ b/libraries/ghc-heap/GHC/Exts/Heap.hs
@@ -71,6 +71,7 @@ import GHC.Exts.Heap.InfoTable
 #endif
 import GHC.Exts.Heap.Utils
 import qualified GHC.Exts.Heap.FFIClosures as FFIClosures
+import qualified GHC.Exts.Heap.ProfInfo.PeekProfInfo as PPI
 
 import Control.Monad
 import Data.Bits
@@ -170,13 +171,19 @@ getClosureDataFromHeapObject x = do
 getClosureDataFromHeapRep :: ByteArray# -> Ptr StgInfoTable -> [b] -> IO (GenClosure b)
 getClosureDataFromHeapRep heapRep infoTablePtr pts = do
   itbl <- peekItbl infoTablePtr
-  getClosureDataFromHeapRepPrim (dataConNames infoTablePtr) itbl heapRep pts
+  getClosureDataFromHeapRepPrim (dataConNames infoTablePtr) PPI.peekTopCCS itbl heapRep pts
 
 getClosureDataFromHeapRepPrim
     :: IO (String, String, String)
     -- ^ A continuation used to decode the constructor description field,
     -- in ghc-debug this code can lead to segfaults because dataConNames
     -- will dereference a random part of memory.
+    -> (Ptr a -> IO (Maybe CostCentreStack))
+    -- ^ A continuation which is used to decode a cost centre stack
+    -- In ghc-debug, this code will need to call back into the debuggee to
+    -- fetch the representation of the CCS before decoding it. Using
+    -- `peekTopCCS` for this argument can lead to segfaults in ghc-debug as
+    -- the CCS argument will point outside the copied closure.
     -> StgInfoTable
     -- ^ The `StgInfoTable` of the closure, extracted from the heap
     -- representation.
@@ -191,7 +198,7 @@ getClosureDataFromHeapRepPrim
     -- `b` is some representation of a pointer e.g. `Any` or `Ptr Any`.
     -> IO (GenClosure b)
     -- ^ Heap representation of the closure.
-getClosureDataFromHeapRepPrim getConDesc itbl heapRep pts = do
+getClosureDataFromHeapRepPrim getConDesc decodeCCS itbl heapRep pts = do
     let -- heapRep as a list of words.
         rawHeapWords :: [Word]
         rawHeapWords = [W# (indexWordArray# heapRep i) | I# i <- [0.. end] ]
@@ -343,7 +350,7 @@ getClosureDataFromHeapRepPrim getConDesc itbl heapRep pts = do
                 }
         TSO | [ u_lnk, u_gbl_lnk, tso_stack, u_trec, u_blk_ex, u_bq] <- pts
                 -> withArray rawHeapWords (\ptr -> do
-                    fields <- FFIClosures.peekTSOFields ptr
+                    fields <- FFIClosures.peekTSOFields decodeCCS ptr
                     pure $ TSOClosure
                         { info = itbl
                         , link = u_lnk
diff --git a/libraries/ghc-heap/GHC/Exts/Heap/FFIClosures_ProfilingDisabled.hsc b/libraries/ghc-heap/GHC/Exts/Heap/FFIClosures_ProfilingDisabled.hsc
index eabb098a15..69c88db57d 100644
--- a/libraries/ghc-heap/GHC/Exts/Heap/FFIClosures_ProfilingDisabled.hsc
+++ b/libraries/ghc-heap/GHC/Exts/Heap/FFIClosures_ProfilingDisabled.hsc
@@ -30,8 +30,8 @@ data TSOFields = TSOFields {
 }
 
 -- | Get non-pointer fields from @StgTSO_@ (@TSO.h@)
-peekTSOFields :: Ptr tsoPtr -> IO TSOFields
-peekTSOFields ptr = do
+peekTSOFields :: (Ptr a -> IO (Maybe CostCentreStack)) -> Ptr tsoPtr -> IO TSOFields
+peekTSOFields decodeCCS ptr = do
     what_next' <- (#peek struct StgTSO_, what_next) ptr
     why_blocked' <- (#peek struct StgTSO_, why_blocked) ptr
     flags' <- (#peek struct StgTSO_, flags) ptr
@@ -40,7 +40,7 @@ peekTSOFields ptr = do
     dirty' <- (#peek struct StgTSO_, dirty) ptr
     alloc_limit' <- (#peek struct StgTSO_, alloc_limit) ptr
     tot_stack_size' <- (#peek struct StgTSO_, tot_stack_size) ptr
-    tso_prof' <- peekStgTSOProfInfo ptr
+    tso_prof' <- peekStgTSOProfInfo decodeCCS ptr
 
     return TSOFields {
         tso_what_next = parseWhatNext what_next',
diff --git a/libraries/ghc-heap/GHC/Exts/Heap/FFIClosures_ProfilingEnabled.hsc b/libraries/ghc-heap/GHC/Exts/Heap/FFIClosures_ProfilingEnabled.hsc
index 124940d1cd..17bf3c8334 100644
--- a/libraries/ghc-heap/GHC/Exts/Heap/FFIClosures_ProfilingEnabled.hsc
+++ b/libraries/ghc-heap/GHC/Exts/Heap/FFIClosures_ProfilingEnabled.hsc
@@ -30,8 +30,8 @@ data TSOFields = TSOFields {
 }
 
 -- | Get non-pointer fields from @StgTSO_@ (@TSO.h@)
-peekTSOFields :: Ptr tsoPtr -> IO TSOFields
-peekTSOFields ptr = do
+peekTSOFields :: (Ptr a -> IO (Maybe CostCentreStack)) -> Ptr tsoPtr -> IO TSOFields
+peekTSOFields decodeCCS ptr = do
     what_next' <- (#peek struct StgTSO_, what_next) ptr
     why_blocked' <- (#peek struct StgTSO_, why_blocked) ptr
     flags' <- (#peek struct StgTSO_, flags) ptr
@@ -40,7 +40,7 @@ peekTSOFields ptr = do
     dirty' <- (#peek struct StgTSO_, dirty) ptr
     alloc_limit' <- (#peek struct StgTSO_, alloc_limit) ptr
     tot_stack_size' <- (#peek struct StgTSO_, tot_stack_size) ptr
-    tso_prof' <- peekStgTSOProfInfo ptr
+    tso_prof' <- peekStgTSOProfInfo decodeCCS ptr
 
     return TSOFields {
         tso_what_next = parseWhatNext what_next',
diff --git a/libraries/ghc-heap/GHC/Exts/Heap/ProfInfo/PeekProfInfo_ProfilingDisabled.hsc b/libraries/ghc-heap/GHC/Exts/Heap/ProfInfo/PeekProfInfo_ProfilingDisabled.hsc
index 0f574b4f03..e9b5b5e43d 100644
--- a/libraries/ghc-heap/GHC/Exts/Heap/ProfInfo/PeekProfInfo_ProfilingDisabled.hsc
+++ b/libraries/ghc-heap/GHC/Exts/Heap/ProfInfo/PeekProfInfo_ProfilingDisabled.hsc
@@ -1,5 +1,6 @@
 module GHC.Exts.Heap.ProfInfo.PeekProfInfo_ProfilingDisabled(
     peekStgTSOProfInfo
+  , peekTopCCS
 ) where
 
 import Prelude
@@ -8,5 +9,8 @@ import GHC.Exts.Heap.ProfInfo.Types
 
 -- | This implementation is used when PROFILING is undefined.
 -- It always returns 'Nothing', because there is no profiling info available.
-peekStgTSOProfInfo :: Ptr tsoPtr -> IO (Maybe StgTSOProfInfo)
-peekStgTSOProfInfo _ = return Nothing
+peekStgTSOProfInfo :: (Ptr a -> IO (Maybe CostCentreStack)) -> Ptr tsoPtr -> IO (Maybe StgTSOProfInfo)
+peekStgTSOProfInfo _ _ = return Nothing
+
+peekTopCCS :: Ptr a -> IO (Maybe CostCentreStack)
+peekTopCCS _ = return Nothing
diff --git a/libraries/ghc-heap/GHC/Exts/Heap/ProfInfo/PeekProfInfo_ProfilingEnabled.hsc b/libraries/ghc-heap/GHC/Exts/Heap/ProfInfo/PeekProfInfo_ProfilingEnabled.hsc
index 363d73d79a..964b1f0b45 100644
--- a/libraries/ghc-heap/GHC/Exts/Heap/ProfInfo/PeekProfInfo_ProfilingEnabled.hsc
+++ b/libraries/ghc-heap/GHC/Exts/Heap/ProfInfo/PeekProfInfo_ProfilingEnabled.hsc
@@ -3,6 +3,7 @@
 
 module GHC.Exts.Heap.ProfInfo.PeekProfInfo_ProfilingEnabled(
     peekStgTSOProfInfo
+    , peekTopCCS
 ) where
 
 #if __GLASGOW_HASKELL__ >= 811
@@ -33,16 +34,20 @@ import           Prelude
 type AddressSet = IntSet
 type AddressMap = IntMap
 
-peekStgTSOProfInfo :: Ptr a -> IO (Maybe StgTSOProfInfo)
-peekStgTSOProfInfo tsoPtr = do
+peekStgTSOProfInfo :: (Ptr b -> IO (Maybe CostCentreStack)) -> Ptr a -> IO (Maybe StgTSOProfInfo)
+peekStgTSOProfInfo decodeCCS tsoPtr = do
     cccs_ptr <- peekByteOff tsoPtr cccsOffset
-    costCenterCacheRef <- newIORef IntMap.empty
-    cccs' <- peekCostCentreStack IntSet.empty costCenterCacheRef cccs_ptr
+    cccs' <- decodeCCS cccs_ptr
 
     return $ Just StgTSOProfInfo {
         cccs = cccs'
     }
 
+peekTopCCS :: Ptr b -> IO (Maybe CostCentreStack)
+peekTopCCS cccs_ptr = do
+  costCenterCacheRef <- newIORef IntMap.empty
+  peekCostCentreStack IntSet.empty costCenterCacheRef cccs_ptr
+
 cccsOffset :: Int
 cccsOffset = (#const OFFSET_StgTSO_cccs) + (#size StgHeader)
 
@@ -162,4 +167,7 @@ import GHC.Exts.Heap.ProfInfo.Types
 
 peekStgTSOProfInfo :: Ptr a -> IO (Maybe StgTSOProfInfo)
 peekStgTSOProfInfo _ = return Nothing
+
+peekTopCCS :: Ptr a -> IO (Maybe CostCentreStack)
+peekTopCCS _ = return Nothing
 #endif