-- | Helpful functions that should have been included in System.Random.

module System.Random.XRayCustom
  ( withRandomGenIORef
  , replicateRandom
  ) where

import Prelude

import Data.IORef
import System.Random

-- | Like 'getStdRandom', except use a given 'IORef StdGen' instead of the
-- global 'theStdGen'.
withRandomGenIORef :: (RandomGen g) => IORef g -> (g -> (a, g)) -> IO a
withRandomGenIORef :: IORef g -> (g -> (a, g)) -> IO a
withRandomGenIORef IORef g
ioRef g -> (a, g)
f = IORef g -> (g -> (g, a)) -> IO a
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef g
ioRef ((a, g) -> (g, a)
forall b a. (b, a) -> (a, b)
swap ((a, g) -> (g, a)) -> (g -> (a, g)) -> g -> (g, a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. g -> (a, g)
f)
  where swap :: (b, a) -> (a, b)
swap (b
x, a
y) = (a
y, b
x)

-- | Runs the supplied function a specified number of times, each time using
-- the new 'RandomGen' value.
replicateRandom :: (RandomGen g) => Int -> g -> (g -> (a, g)) -> ([a], g)
replicateRandom :: Int -> g -> (g -> (a, g)) -> ([a], g)
replicateRandom Int
n g
gen g -> (a, g)
f = Int -> ([a], g) -> ([a], g)
replicateRandom' Int
n ([], g
gen)
 where
  replicateRandom' :: Int -> ([a], g) -> ([a], g)
replicateRandom' Int
n' ([a]
xs, g
gen')
    | Int
n' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0
    = ([a]
xs, g
gen')
    | Bool
otherwise
    = let ~(a
x, g
gen'') = g -> (a, g)
f g
gen' in Int -> ([a], g) -> ([a], g)
replicateRandom' (Int
n' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
xs, g
gen'')