Add overwriteGen to FrozenGen

Author: lehins

CHANGELOG.md

@@ -1,6 +1,6 @@
 # 1.3.0
 
-* Add `modifyGen` to the `FrozenGen` type class
+* Add `modifyGen` and `overwriteGen` to the `FrozenGen` type class
 * Add `splitFrozenM` and `splitMutableM`
 * Switch `randomM` and `randomRM` to use `FrozenGen` instead of `RandomGenM`
 * Deprecate `RandomGenM` in favor of a more powerful `FrozenGen`

src/System/Random/Internal.hs

@@ -315,6 +315,11 @@ class StatefulGen (MutableGen f m) m => FrozenGen f m where
   -- @since 1.3.0
   modifyGen :: MutableGen f m -> (f -> (a, f)) -> m a
 
+  -- | Overwrite the mutable generator with the supplied frozen one
+  --
+  -- @since 1.3.0
+  overwriteGen :: MutableGen f m -> f -> m ()
+  overwriteGen mg fg = modifyGen mg (const ((), fg))
 
 -- | Splits a pseudo-random number generator into two. Overwrites the mutable
 -- wrapper with one of the resulting generators and returns the other.
@@ -479,6 +484,8 @@ instance (RandomGen g, MonadState g m) => FrozenGen (StateGen g) m where
   thawGen (StateGen g) = StateGenM <$ put g
   modifyGen _ f = state (coerce f)
   {-# INLINE modifyGen #-}
+  overwriteGen _ f = put (coerce f)
+  {-# INLINE overwriteGen #-}
 
 -- | Splits a pseudo-random number generator into two. Updates the state with
 -- one of the resulting generators and returns the other.

src/System/Random/Stateful.hs

@@ -473,6 +473,8 @@ instance (RandomGen g, MonadIO m) => FrozenGen (IOGen g) m where
     g' `seq` writeIORef ref g'
     pure a
   {-# INLINE modifyGen #-}
+  overwriteGen (IOGenM ref) = liftIO . writeIORef ref . unIOGen
+  {-# INLINE overwriteGen #-}
 
 -- | Applies a pure operation to the wrapped pseudo-random number generator.
 --
@@ -538,6 +540,8 @@ instance RandomGen g => FrozenGen (STGen g) (ST s) where
     g' `seq` writeSTRef ref g'
     pure a
   {-# INLINE modifyGen #-}
+  overwriteGen (STGenM ref) = writeSTRef ref . unSTGen
+  {-# INLINE overwriteGen #-}
 
 
 -- | Applies a pure operation to the wrapped pseudo-random number generator.
@@ -639,6 +643,8 @@ instance RandomGen g => FrozenGen (TGen g) STM where
     g' `seq` writeTVar ref g'
     pure a
   {-# INLINE modifyGen #-}
+  overwriteGen (TGenM ref) = writeTVar ref . unTGen
+  {-# INLINE overwriteGen #-}
 
 
 -- | Applies a pure operation to the wrapped pseudo-random number generator.