From 99073e75e80ce890109119e627930ca1fe358737 Mon Sep 17 00:00:00 2001
From: Alexey Kuleshevich <alexey@kuleshevi.ch>
Date: Sat, 27 Jan 2024 15:29:47 +0100
Subject: [PATCH 1/2] Make IOGenM and STGenM lazy in the value generated

---
 src/System/Random/Stateful.hs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/System/Random/Stateful.hs b/src/System/Random/Stateful.hs
index a1a3f190..4dd8839d 100644
--- a/src/System/Random/Stateful.hs
+++ b/src/System/Random/Stateful.hs
@@ -526,7 +526,7 @@ applyIOGen :: MonadIO m => (g -> (a, g)) -> IOGenM g -> m a
 applyIOGen f (IOGenM ref) = liftIO $ do
   g <- readIORef ref
   case f g of
-    (!a, !g') -> a <$ writeIORef ref g'
+    (a, !g') -> a <$ writeIORef ref g'
 {-# INLINE applyIOGen #-}
 
 -- | Wraps an 'STRef' that holds a pure pseudo-random number generator.
@@ -600,7 +600,7 @@ applySTGen :: (g -> (a, g)) -> STGenM g s -> ST s a
 applySTGen f (STGenM ref) = do
   g <- readSTRef ref
   case f g of
-    (!a, !g') -> a <$ writeSTRef ref g'
+    (a, !g') -> a <$ writeSTRef ref g'
 {-# INLINE applySTGen #-}
 
 -- | Runs a monadic generating action in the `ST` monad using a pure

From 5db23532a76e28697a2da5abd24fbb1ef766aa3f Mon Sep 17 00:00:00 2001
From: Alexey Kuleshevich <alexey@kuleshevi.ch>
Date: Sat, 27 Jan 2024 17:11:59 +0100
Subject: [PATCH 2/2] Make AtomicIOGenM lazy in the value generated

---
 src/System/Random/Stateful.hs | 25 +++++++++++++++++++++++--
 1 file changed, 23 insertions(+), 2 deletions(-)

diff --git a/src/System/Random/Stateful.hs b/src/System/Random/Stateful.hs
index 4dd8839d..f45ca0ac 100644
--- a/src/System/Random/Stateful.hs
+++ b/src/System/Random/Stateful.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE CPP #-}
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE FunctionalDependencies #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
@@ -143,6 +144,10 @@ import Data.STRef
 import Foreign.Storable
 import System.Random
 import System.Random.Internal
+#if __GLASGOW_HASKELL__ >= 808
+import GHC.IORef (atomicModifyIORef2Lazy)
+#endif
+
 
 -- $introduction
 --
@@ -414,7 +419,7 @@ instance (RandomGen g, MonadIO m) => FrozenGen (AtomicGen g) m where
   type MutableGen (AtomicGen g) m = AtomicGenM g
   freezeGen = fmap AtomicGen . liftIO . readIORef . unAtomicGenM
   modifyGen (AtomicGenM ioRef) f =
-    liftIO $ atomicModifyIORef' ioRef $ \g ->
+    liftIO $ atomicModifyIORefHS ioRef $ \g ->
       case f (AtomicGen g) of
         (a, AtomicGen g') -> (g', a)
   {-# INLINE modifyGen #-}
@@ -436,11 +441,27 @@ instance (RandomGen g, MonadIO m) => ThawedGen (AtomicGen g) m where
 -- @since 1.2.0
 applyAtomicGen :: MonadIO m => (g -> (a, g)) -> AtomicGenM g -> m a
 applyAtomicGen op (AtomicGenM gVar) =
-  liftIO $ atomicModifyIORef' gVar $ \g ->
+  liftIO $ atomicModifyIORefHS gVar $ \g ->
     case op g of
       (a, g') -> (g', a)
 {-# INLINE applyAtomicGen #-}
 
+-- HalfStrict version of atomicModifyIORef, i.e. strict in the modifcation of the contents
+-- of the IORef, but not in the result produced.
+atomicModifyIORefHS :: IORef a -> (a -> (a, b)) -> IO b
+atomicModifyIORefHS ref f = do
+#if __GLASGOW_HASKELL__ >= 808
+  (_old, (_new, res)) <- atomicModifyIORef2Lazy ref $ \old ->
+    case f old of
+      r@(!_new, _res) -> r
+  pure res
+#else
+  atomicModifyIORef ref $ \old ->
+    case f old of
+      r@(!_new, _res) -> r
+#endif
+{-# INLINE atomicModifyIORefHS #-}
+
 -- | Wraps an 'IORef' that holds a pure pseudo-random number generator.
 --
 -- *   'IOGenM' is safe in the presence of exceptions, but not concurrency.