From e1815c6e8ce7a51236bce3d1770033098d50ad96 Mon Sep 17 00:00:00 2001
From: Matthew Pickering <matthewtpickering@gmail.com>
Date: Wed, 3 Apr 2024 11:23:19 +0100
Subject: [PATCH] Fix snapshotting

We also need to find all the RequestBlock requests and add those into
the block cache, as well as the results of RequestAllBlocks.
---
 client/src/GHC/Debug/Client/Monad/Simple.hs |  4 +---
 client/src/GHC/Debug/Client/RequestCache.hs | 16 ++++++++++++++++
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/client/src/GHC/Debug/Client/Monad/Simple.hs b/client/src/GHC/Debug/Client/Monad/Simple.hs
index 6171fa65..7ba7cb7b 100644
--- a/client/src/GHC/Debug/Client/Monad/Simple.hs
+++ b/client/src/GHC/Debug/Client/Monad/Simple.hs
@@ -123,9 +123,7 @@ instance DebugMonad DebugM where
 
 initBlockCacheFromReqCache :: RequestCache -> BlockCache
 initBlockCacheFromReqCache new_req_cache  =
-  case lookupReq RequestAllBlocks new_req_cache of
-        Just bs -> addBlocks bs emptyBlockCache
-        Nothing -> emptyBlockCache
+  addBlocks (lookupBlocks new_req_cache) emptyBlockCache
 
 
 
diff --git a/client/src/GHC/Debug/Client/RequestCache.hs b/client/src/GHC/Debug/Client/RequestCache.hs
index 43529625..9e37cd7e 100644
--- a/client/src/GHC/Debug/Client/RequestCache.hs
+++ b/client/src/GHC/Debug/Client/RequestCache.hs
@@ -4,6 +4,7 @@
 module GHC.Debug.Client.RequestCache(RequestCache
                                     , cacheReq
                                     , lookupReq
+                                    , lookupBlocks
                                     , emptyRequestCache
                                     , clearMovableRequests
                                     , putCache
@@ -33,6 +34,21 @@ lookupReq req (RequestCache rc) = coerceResult <$> HM.lookup (AnyReq req) rc
     coerceResult :: AnyResp -> resp
     coerceResult (AnyResp a _) = unsafeCoerce a
 
+lookupBlocks :: RequestCache -> [RawBlock]
+lookupBlocks c@(RequestCache rc) =
+  let all_blocks = case lookupReq RequestAllBlocks c of
+                        Just bs -> bs
+                        Nothing -> []
+
+      get_block :: AnyReq -> AnyResp -> [RawBlock] -> [RawBlock]
+      get_block (AnyReq (RequestBlock {})) (AnyResp resp _) bs = unsafeCoerce resp : bs
+      get_block _ _ bs = bs
+
+      individual_blocks = HM.foldrWithKey get_block [] rc
+
+  in (all_blocks ++ individual_blocks)
+
+
 emptyRequestCache :: RequestCache
 emptyRequestCache = RequestCache HM.empty
 
-- 
GitLab