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 `splitGen` and `splitMutableGen`
 * Switch `randomM` and `randomRM` to use `FrozenGen` instead of `RandomGenM`
 * Deprecate `RandomGenM` in favor of a more powerful `FrozenGen`

src/System/Random/Internal.hs

@@ -309,11 +309,17 @@ class StatefulGen (MutableGen f m) m => FrozenGen f m where
   -- @since 1.2.0
   thawGen :: f -> m (MutableGen f m)
 
-  -- | Apply a pure function to the frozen generator.
+  -- | Apply a pure function to the frozen pseudo-random number generator.
   --
   -- @since 1.3.0
   modifyGen :: MutableGen f m -> (f -> (a, f)) -> m a
 
+  -- | Overwrite contents of the mutable pseudo-random number 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.
@@ -478,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 #-}
 
 -- | Runs a monadic generating action in the `State` monad using a pure
 -- pseudo-random number generator.

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.

test/Spec/Stateful.hs

@@ -64,6 +64,22 @@ withMutableGenSpec toIO frozen =
     r' <- withMutableGen_ frozen action
     pure $ x == y && r == r'
 
+overwriteMutableGenSpec ::
+     forall f m. (FrozenGen f m, Eq f, Show f)
+  => (forall a. m a -> IO a)
+  -> f
+  -> Property IO
+overwriteMutableGenSpec toIO frozen =
+  forAll $ \n -> monadic $ toIO $ do
+    let action = uniformListM (abs n + 1) -- Non-empty
+    ((r1, r2), frozen') :: ((String, String), f) <- withMutableGen frozen $ \mutable -> do
+      r1 <- action mutable
+      overwriteGen mutable frozen
+      r2 <- action mutable
+      modifyGen mutable (const ((), frozen))
+      pure (r1, r2)
+    pure $ r1 == r2 && frozen == frozen'
+
 splitMutableGenSpec ::
      forall f m. (RandomGen f, FrozenGen f m, Eq f, Show f)
   => (forall a. m a -> IO a)
@@ -86,6 +102,8 @@ statefulSpecFor toIO toStdGen =
     (showsTypeRep (typeRep (Proxy :: Proxy f)) "")
     [ testProperty "withMutableGen" $
       forAll $ \(f :: f) -> withMutableGenSpec toIO f
+    , testProperty "overwriteGen" $
+      forAll $ \(f :: f) -> overwriteMutableGenSpec toIO f
     , testProperty "splitGen" $
       forAll $ \(f :: f) -> splitMutableGenSpec toIO f
     , testGroup