Add ability to split PrimGen

idontgetoutmuch · Mar 13, 2020 · b918908 · b918908
1 parent 9564db4
commit b918908
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/System/Random.hs b/System/Random.hs
@@ -71,6 +71,7 @@ module System.Random
   , PrimGen
   , runPrimGenST
   , runPrimGenIO
+  , splitPrimGen
   , atomicPrimGen
 
   -- ** The global random number generator
@@ -360,6 +361,17 @@ atomicPrimGen op (PrimGen gVar) =
     case op g of
       (a, g') -> (g', a)
 
+
+-- | Split `PrimGen` into atomically updated current generator and a newly created that is
+-- returned.
+--
+-- @since 1.2
+splitPrimGen ::
+     (RandomGen g, PrimMonad m)
+  => PrimGen (PrimState m) g
+  -> m (PrimGen (PrimState m) g)
+splitPrimGen = atomicPrimGen split >=> restore
+
 runPrimGenST :: RandomGen g => g -> (forall s . PrimGen s g -> ST s a) -> (a, g)
 runPrimGenST g action = runST $ do
   primGen :: PrimGen s g <- restore g