Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Make result lazy when mutating #161

Merged
merged 2 commits into from
Oct 28, 2024
+25 −4
Merged

Make result lazy when mutating #161

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
99073e7
Make IOGenM and STGenM lazy in the value generated
lehins Jan 27, 2024
5db2353
Make AtomicIOGenM lazy in the value generated
lehins Jan 27, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 25 additions & 4 deletions src/System/Random/Stateful.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
Expand Down Expand Up @@ -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
--
Expand Down Expand Up @@ -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 #-}
Expand All @@ -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.
Expand Down Expand Up @@ -526,7 +547,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.
Expand Down Expand Up @@ -600,7 +621,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
Expand Down