From 7e17707ce9b24062070dc1148ae18f8e7c8e7864 Mon Sep 17 00:00:00 2001
From: Joris Dral <joris@well-typed.com>
Date: Mon, 25 Nov 2024 18:24:15 +0100
Subject: [PATCH] Re-enable `NoThunks` tests (#444)

Tests are re-enabled, but in cases where I expect failures related to
`StrictMVar`s, we ignore the thunk and keep checking further into the `MVar`
contents.
---
 lsm-tree.cabal                                |  1 +
 .../Database/LSMTree/Extras/NoThunks.hs       | 32 +++++++++++++++----
 .../Database/LSMTree/Normal/StateMachine.hs   |  3 +-
 3 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/lsm-tree.cabal b/lsm-tree.cabal
index 66a6f06ab..3d2c8af83 100644
--- a/lsm-tree.cabal
+++ b/lsm-tree.cabal
@@ -318,6 +318,7 @@ library extras
     , contra-tracer
     , deepseq
     , fs-api
+    , fs-sim
     , io-classes:strict-mvar
     , io-classes:strict-stm
     , lsm-tree
diff --git a/src-extras/Database/LSMTree/Extras/NoThunks.hs b/src-extras/Database/LSMTree/Extras/NoThunks.hs
index f37396f9e..244b4a81a 100644
--- a/src-extras/Database/LSMTree/Extras/NoThunks.hs
+++ b/src-extras/Database/LSMTree/Extras/NoThunks.hs
@@ -68,6 +68,8 @@ import           KMerge.Heap
 import           NoThunks.Class
 import           System.FS.API
 import           System.FS.BlockIO.API
+import           System.FS.IO
+import           System.FS.Sim.MockFS
 import           Test.QuickCheck (Property, Testable (..), counterexample)
 import           Unsafe.Coerce
 
@@ -542,7 +544,7 @@ instance (NoThunks a, Typeable s, Typeable a) => NoThunks (MutableHeap s a) wher
 -- a)@, can not be satisfied for arbitrary @m@\/@s@, and must be instantiated
 -- for a concrete @m@\/@s@, like @IO@\/@RealWorld@.
 class ( forall a. NoThunks a => NoThunks (StrictTVar m a)
-      , forall a. NoThunks a => NoThunks (StrictMVar m a)
+      , forall a. (NoThunks a, Typeable a) => NoThunks (StrictMVar m a)
       ) => NoThunksIOLike' m s
 
 instance NoThunksIOLike' IO RealWorld
@@ -564,11 +566,29 @@ instance NoThunks a => NoThunks (StrictTVar IO a) where
 #endif
 #endif
 
-instance NoThunks a => NoThunks (StrictMVar IO a) where
-  showTypeOf (_ :: Proxy (StrictMVar IO a)) = "StrictMVar IO"
-  wNoThunks ctx var = do
-      x <- readMVar var
-      noThunks ctx x
+-- TODO: in some cases, strict-mvar functions leave thunks behind, in particular
+-- modifyMVarMasked and modifyMVarMasked_. So in some specific cases we evaluate
+-- the contents of the MVar to WHNF, and keep checking nothunks from there. See
+-- lsm-tree#444.
+instance (NoThunks a, Typeable a) => NoThunks (StrictMVar IO a) where
+  showTypeOf (p :: Proxy (StrictMVar IO a)) = show $ typeRep p
+  wNoThunks ctx var
+    | Just (Proxy :: Proxy (MergingRunState IO HandleIO))
+        <- gcast (Proxy @a)
+    = workAroundCheck
+    | Just (Proxy :: Proxy (MergingRunState IO HandleMock))
+        <- gcast (Proxy @a)
+    = workAroundCheck
+    | otherwise
+    = properCheck
+    where
+      properCheck = do
+        x <- readMVar var
+        noThunks ctx x
+
+      workAroundCheck = do
+        !x <- readMVar var
+        noThunks ctx x
 
 {-------------------------------------------------------------------------------
   vector
diff --git a/test/Test/Database/LSMTree/Normal/StateMachine.hs b/test/Test/Database/LSMTree/Normal/StateMachine.hs
index 82251ffec..305ceedaf 100644
--- a/test/Test/Database/LSMTree/Normal/StateMachine.hs
+++ b/test/Test/Database/LSMTree/Normal/StateMachine.hs
@@ -957,8 +957,7 @@ runIO action lookUp = ReaderT $ \(session, handler) -> do
     x <- aux (unwrapSession session) handler action
     case session of
       WrapSession sesh ->
-        -- TODO: Re-enable NoThunks assertions. See lsm-tree#444.
-        const id (assertNoThunks sesh) $ pure ()
+        assertNoThunks sesh $ pure ()
     pure x
   where
     aux ::