Add IORef, MVar and TVar instances

nothunks.cabal

@@ -26,6 +26,7 @@ library
     build-depends:    base       >= 4.12 && < 5
                     , bytestring >= 0.10 && < 0.11
                     , containers >= 0.5  && < 0.7
+                    , stm        >= 2.5  && < 2.6
                     , text       >= 1.2  && < 1.3
                     , time       >= 1.5  && < 1.11
                     , vector     >= 0.12 && < 0.13
@@ -49,6 +50,7 @@ test-suite nothunks-test
 
                       -- Dependencies shared with the lib
                     , containers
+                    , stm
 
                       -- Whatever is bundled with ghc
                     , ghc-prim

src/NoThunks/Class.hs

@@ -60,10 +60,13 @@ import Data.Word
 import GHC.Stack
 import Numeric.Natural
 
+import qualified Control.Concurrent.MVar       as MVar
+import qualified Control.Concurrent.STM.TVar   as TVar
 import qualified Data.ByteString               as BS.Strict
 import qualified Data.ByteString.Lazy          as BS.Lazy
 import qualified Data.ByteString.Lazy.Internal as BS.Lazy.Internal
 import qualified Data.IntMap                   as IntMap
+import qualified Data.IORef                    as IORef
 import qualified Data.Map                      as Map
 import qualified Data.Set                      as Set
 import qualified Data.Text                     as Text.Strict
@@ -456,6 +459,37 @@ deriving via OnlyCheckWhnf Word16 instance NoThunks Word16
 deriving via OnlyCheckWhnf Word32 instance NoThunks Word32
 deriving via OnlyCheckWhnf Word64 instance NoThunks Word64
 
+{-------------------------------------------------------------------------------
+  Mutable Vars
+-------------------------------------------------------------------------------}
+
+instance NoThunks a => NoThunks (IORef.IORef a) where
+    showTypeOf _ = "IORef"
+    wNoThunks ctx ref = do
+        val <- IORef.readIORef ref
+        noThunks ctx val
+
+instance NoThunks a => NoThunks (MVar.MVar a) where
+    showTypeOf _ = "MVar"
+    wNoThunks ctx ref = do
+        val <- MVar.tryReadMVar ref
+        maybe (return Nothing) (noThunks ctx) val
+
+instance NoThunks a => NoThunks (TVar.TVar a) where
+    showTypeOf _ = "TVar"
+    wNoThunks ctx ref = do
+        -- An alternative is to use
+        --
+        --   val <- STM.atomically $ TVar.readTVar ref
+        --
+        -- but that would cause nested atomically failures with
+        -- unsafeNoThunks. Fortunately, readTVarIO doesn't make a transaction.
+        --
+        -- See related tests.
+        --
+        val <- TVar.readTVarIO ref
+        noThunks ctx val
+
 {-------------------------------------------------------------------------------
   Time
 -------------------------------------------------------------------------------}

test/Test/NoThunks/Class.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE BangPatterns        #-}
 {-# LANGUAGE DataKinds           #-}
 {-# LANGUAGE DeriveFunctor       #-}
 {-# LANGUAGE DeriveGeneric       #-}
@@ -38,6 +39,11 @@ import Test.Tasty.Hedgehog
 import qualified Data.Sequence          as Seq
 import qualified Data.Sequence.Internal as Seq.Internal
 
+import qualified Control.Concurrent.MVar       as MVar
+import qualified Control.Concurrent.STM        as STM
+import qualified Control.Concurrent.STM.TVar   as TVar
+import qualified Data.IORef                    as IORef
+
 import Hedgehog
 import Hedgehog.Internal.Report (Result (..), reportStatus)
 import Hedgehog.Internal.Region (displayRegion)
@@ -81,6 +87,11 @@ tests = testGroup "NoThunks.Class" [
         , testProperty "ThunkFreeFn"   $ testWithModel agreeOnContext $ Proxy @(ThunkFree "->" (Int -> Int))
         , testProperty "ThunkFreeIO"   $ testWithModel agreeOnContext $ Proxy @(ThunkFree "IO" (IO ()))
         ]
+    , testGroup "MutableVars" [
+          checkRef (Proxy :: Proxy IORef.IORef)
+        , checkRef (Proxy :: Proxy MVar.MVar)
+        , checkRef (Proxy :: Proxy TVar.TVar)
+        ]
     ]
 
 -- | When using @InspectHeap@ we don't get a context, so merely check if
@@ -542,6 +553,95 @@ sanityCheckIO = checkNF False $ \k -> do
     n <- liftIO $ ack 5 <$> randomRIO (0, 10)
     k (print (notStrict b n 6) :: IO ())
 
+{-------------------------------------------------------------------------------
+  Mutable Vars
+-------------------------------------------------------------------------------}
+
+checkRef :: forall ref. (IsRef ref, NoThunks (ref Int)) => Proxy ref -> TestTree
+checkRef p = testGroup (show (typeRep p)) [
+      testProperty "NotNF"           checkRefNotNF
+    , testProperty "NF"              checkRefNF
+    , testProperty "NotNFPure"       checkRefNotNFPure
+    , testProperty "NFPure"          checkRefNFPure
+    , testProperty "NotNFAtomically" checkRefNotNFAtomically
+    , testProperty "NFAtomically"    checkRefNFAtomically
+    ]
+  where
+    checkRefNotNF :: Property
+    checkRefNotNF = checkNFClass False $ \k -> do
+        ref <- liftIO (newRef (if ack 3 3 > 0 then x else x) :: IO (ref Int))
+        k ref
+      where
+        x :: Int
+        x = 0
+
+    checkRefNF :: Property
+    checkRefNF = checkNFClass True $ \k -> do
+        ! ref <- liftIO (newRef x :: IO (ref Int))
+        k ref
+      where
+        x :: Int
+        !x = 0
+
+    checkRefNotNFPure :: Property
+    checkRefNotNFPure = unsafeCheckNF False $ \k -> do
+        ref <- liftIO (newRef (if ack 3 3 > 0 then x else x) :: IO (ref Int))
+        k ref
+      where
+        x :: Int
+        x = 0
+
+    checkRefNFPure :: Property
+    checkRefNFPure = unsafeCheckNF True $ \k -> do
+        ! ref <- liftIO (newRef x :: IO (ref Int))
+        k ref
+      where
+        x :: Int
+        !x = 0
+
+    checkRefNotNFAtomically :: Property
+    checkRefNotNFAtomically = unsafeCheckNFAtomically False $ \k -> do
+        ref <- liftIO (newRef (if ack 3 3 > 0 then x else x) :: IO (ref Int))
+        k ref
+      where
+        x :: Int
+        x = 0
+
+    checkRefNFAtomically :: Property
+    checkRefNFAtomically = unsafeCheckNFAtomically True $ \k -> do
+        ! ref <- liftIO (newRef x :: IO (ref Int))
+        k ref
+      where
+        x :: Int
+        !x = 0
+
+class Typeable ref => IsRef ref where newRef :: a -> IO (ref a)
+
+instance IsRef IORef.IORef where newRef = IORef.newIORef
+instance IsRef MVar.MVar   where newRef = MVar.newMVar
+instance IsRef TVar.TVar   where newRef = TVar.newTVarIO
+
+checkNFClass :: NoThunks a => Bool -> ((a -> PropertyT IO ()) -> PropertyT IO ()) -> Property
+checkNFClass expectedNF k = withTests 1 $ property $ k $ \x -> do
+    nf <- liftIO $ noThunks [] x
+    isNothing nf === expectedNF
+
+{-# NOINLINE unsafeCheckNF #-}
+unsafeCheckNF :: NoThunks a => Bool -> ((a -> PropertyT IO ()) -> PropertyT IO ()) -> Property
+unsafeCheckNF expectedNF k = withTests 1 $ property $ k $ \x -> do
+    let nf = unsafeNoThunks x
+    isNothing nf === expectedNF
+
+{-# NOINLINE unsafeCheckNFAtomically #-}
+unsafeCheckNFAtomically :: NoThunks a => Bool -> ((a -> PropertyT IO ()) -> PropertyT IO ()) -> Property
+unsafeCheckNFAtomically expectedNF k = withTests 1 $ property $ k $ \x -> do
+    tvar <- liftIO (TVar.newTVarIO True)
+    true <- liftIO $ STM.atomically $ do
+        val <- TVar.readTVar tvar
+        -- the $! is essential to trigger NestedAtomically exception.
+        return $! val && isNothing (unsafeNoThunks x)
+    true === expectedNF
+
 {-------------------------------------------------------------------------------
   Hedgehog auxiliary
 -------------------------------------------------------------------------------}