{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}

-- |
-- Copyright: © 2021 IOHK
-- License: Apache-2.0
--
-- This module provides functions and types that extend those provided by
-- the 'Control.Monad.Random' module hierarchy.
--
module Control.Monad.Random.Extra
    (
    -- * Random number generator seeds
      StdGenSeed (..)
    , stdGenSeed
    , stdGenFromSeed
    , stdGenToSeed

    -- * Non-random contexts
    , NonRandom (..)

    ) where

import Prelude

import Control.Applicative
    ( Applicative (..) )
import Control.Monad.Random.Class
    ( MonadRandom (..) )
import Data.Aeson
    ( FromJSON (..), ToJSON (..), Value (Number) )
import Data.Aeson.Extra
    ( parseBoundedIntegral )
import Data.Bits
    ( (.|.) )
import Data.Coerce
    ( coerce )
import Data.Word
    ( Word64 )
import Data.Word.Odd
    ( Lit, OddWord )
import GHC.Generics
    ( Generic )
import Quiet
    ( Quiet (..) )
import System.Random
    ( Random (..), RandomGen (..) )
import System.Random.Internal
    ( StdGen (..) )
import System.Random.SplitMix
    ( seedSMGen', unseedSMGen )

import qualified Data.Bits as Bits

--------------------------------------------------------------------------------
-- Random number generator seeds
--------------------------------------------------------------------------------

-- | A seed for the standard random number generator.
--
-- This type is equivalent to the internal state of a 'StdGen', but provides a
-- representation that is more convenient for construction and serialization.
--
-- The number of possible seeds is identical to the number of valid states of
-- the 'StdGen' type, but unlike the 'StdGen' type, whose state has an internal
-- invariant that must not be broken, values of the 'StdGenSeed' type are
-- correct by construction.
--
newtype StdGenSeed = StdGenSeed
    { StdGenSeed -> Word127
unStdGenSeed :: Word127
    }
    deriving (StdGenSeed -> StdGenSeed -> Bool
(StdGenSeed -> StdGenSeed -> Bool)
-> (StdGenSeed -> StdGenSeed -> Bool) -> Eq StdGenSeed
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: StdGenSeed -> StdGenSeed -> Bool
$c/= :: StdGenSeed -> StdGenSeed -> Bool
== :: StdGenSeed -> StdGenSeed -> Bool
$c== :: StdGenSeed -> StdGenSeed -> Bool
Eq, StdGenSeed
StdGenSeed -> StdGenSeed -> Bounded StdGenSeed
forall a. a -> a -> Bounded a
maxBound :: StdGenSeed
$cmaxBound :: StdGenSeed
minBound :: StdGenSeed
$cminBound :: StdGenSeed
Bounded, (forall x. StdGenSeed -> Rep StdGenSeed x)
-> (forall x. Rep StdGenSeed x -> StdGenSeed) -> Generic StdGenSeed
forall x. Rep StdGenSeed x -> StdGenSeed
forall x. StdGenSeed -> Rep StdGenSeed x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep StdGenSeed x -> StdGenSeed
$cfrom :: forall x. StdGenSeed -> Rep StdGenSeed x
Generic, Eq StdGenSeed
Eq StdGenSeed
-> (StdGenSeed -> StdGenSeed -> Ordering)
-> (StdGenSeed -> StdGenSeed -> Bool)
-> (StdGenSeed -> StdGenSeed -> Bool)
-> (StdGenSeed -> StdGenSeed -> Bool)
-> (StdGenSeed -> StdGenSeed -> Bool)
-> (StdGenSeed -> StdGenSeed -> StdGenSeed)
-> (StdGenSeed -> StdGenSeed -> StdGenSeed)
-> Ord StdGenSeed
StdGenSeed -> StdGenSeed -> Bool
StdGenSeed -> StdGenSeed -> Ordering
StdGenSeed -> StdGenSeed -> StdGenSeed
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: StdGenSeed -> StdGenSeed -> StdGenSeed
$cmin :: StdGenSeed -> StdGenSeed -> StdGenSeed
max :: StdGenSeed -> StdGenSeed -> StdGenSeed
$cmax :: StdGenSeed -> StdGenSeed -> StdGenSeed
>= :: StdGenSeed -> StdGenSeed -> Bool
$c>= :: StdGenSeed -> StdGenSeed -> Bool
> :: StdGenSeed -> StdGenSeed -> Bool
$c> :: StdGenSeed -> StdGenSeed -> Bool
<= :: StdGenSeed -> StdGenSeed -> Bool
$c<= :: StdGenSeed -> StdGenSeed -> Bool
< :: StdGenSeed -> StdGenSeed -> Bool
$c< :: StdGenSeed -> StdGenSeed -> Bool
compare :: StdGenSeed -> StdGenSeed -> Ordering
$ccompare :: StdGenSeed -> StdGenSeed -> Ordering
$cp1Ord :: Eq StdGenSeed
Ord)
    deriving Int -> StdGenSeed -> ShowS
[StdGenSeed] -> ShowS
StdGenSeed -> String
(Int -> StdGenSeed -> ShowS)
-> (StdGenSeed -> String)
-> ([StdGenSeed] -> ShowS)
-> Show StdGenSeed
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [StdGenSeed] -> ShowS
$cshowList :: [StdGenSeed] -> ShowS
show :: StdGenSeed -> String
$cshow :: StdGenSeed -> String
showsPrec :: Int -> StdGenSeed -> ShowS
$cshowsPrec :: Int -> StdGenSeed -> ShowS
Show via (Quiet StdGenSeed)

type Word127 = OddWord Integer (Lit 127)

instance ToJSON StdGenSeed where
    toJSON :: StdGenSeed -> Value
toJSON = Value -> Value
forall a. ToJSON a => a -> Value
toJSON (Value -> Value) -> (StdGenSeed -> Value) -> StdGenSeed -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Scientific -> Value
Number (Scientific -> Value)
-> (StdGenSeed -> Scientific) -> StdGenSeed -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word127 -> Scientific
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word127 -> Scientific)
-> (StdGenSeed -> Word127) -> StdGenSeed -> Scientific
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StdGenSeed -> Word127
unStdGenSeed

instance FromJSON StdGenSeed where
    parseJSON :: Value -> Parser StdGenSeed
parseJSON = (Word127 -> StdGenSeed) -> Parser Word127 -> Parser StdGenSeed
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Word127 -> StdGenSeed
StdGenSeed (Parser Word127 -> Parser StdGenSeed)
-> (Value -> Parser Word127) -> Value -> Parser StdGenSeed
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Value -> Parser Word127
forall a. (Bounded a, Integral a) => String -> Value -> Parser a
parseBoundedIntegral String
"StdGenSeed"

-- | Creates a new 'StdGenSeed' from within a random monadic context.
--
stdGenSeed :: MonadRandom m => m StdGenSeed
stdGenSeed :: m StdGenSeed
stdGenSeed = do
    Word64
hi <- m Word64
forall (m :: * -> *) a. (MonadRandom m, Random a) => m a
getRandom
    Word64
lo <- m Word64
forall (m :: * -> *) a. (MonadRandom m, Random a) => m a
getRandom
    StdGenSeed -> m StdGenSeed
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StdGenSeed -> m StdGenSeed) -> StdGenSeed -> m StdGenSeed
forall a b. (a -> b) -> a -> b
$ Word127 -> StdGenSeed
StdGenSeed (Word127 -> StdGenSeed) -> Word127 -> StdGenSeed
forall a b. (a -> b) -> a -> b
$ Word127 -> Word127 -> Word127
forall a. Bits a => a -> a -> a
(.|.)
        (Word64 -> Word127
forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word64 @Word127 Word64
hi Word127 -> Int -> Word127
forall a. Bits a => a -> Int -> a
`Bits.shiftL` Int
63)
        (Word64 -> Word127
forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word64 @Word127 Word64
lo)

-- | Converts a 'StdGenSeed' value to a 'StdGen' value.
--
-- This function satisfies the following properties:
--
-- >>> stdGenFromSeed . stdGenToSeed == id
-- >>> stdGenToSeed . stdGenFromSeed == id
--
stdGenFromSeed :: StdGenSeed -> StdGen
stdGenFromSeed :: StdGenSeed -> StdGen
stdGenFromSeed
    = SMGen -> StdGen
StdGen
    (SMGen -> StdGen) -> (StdGenSeed -> SMGen) -> StdGenSeed -> StdGen
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word64, Word64) -> SMGen
seedSMGen'
    ((Word64, Word64) -> SMGen)
-> (StdGenSeed -> (Word64, Word64)) -> StdGenSeed -> SMGen
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (\Word127
s -> (,)
        (Word127 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word127 @Word64 (Word127
s Word127 -> Int -> Word127
forall a. Bits a => a -> Int -> a
`Bits.shiftR` Int
63))
        (Word127 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word127 @Word64 (Word127
s Word127 -> Int -> Word127
forall a. Bits a => a -> Int -> a
`Bits.shiftL` Int
1)))
    (Word127 -> (Word64, Word64))
-> (StdGenSeed -> Word127) -> StdGenSeed -> (Word64, Word64)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StdGenSeed -> Word127
unStdGenSeed

-- | Converts a 'StdGen' value to a 'StdGenSeed' value.
--
-- This function satisfies the following properties:
--
-- >>> stdGenFromSeed . stdGenToSeed == id
-- >>> stdGenToSeed . stdGenFromSeed == id
--
stdGenToSeed :: StdGen -> StdGenSeed
stdGenToSeed :: StdGen -> StdGenSeed
stdGenToSeed
    = Word127 -> StdGenSeed
StdGenSeed
    (Word127 -> StdGenSeed)
-> (StdGen -> Word127) -> StdGen -> StdGenSeed
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (\(Word64
a, Word64
b) -> Word127 -> Word127 -> Word127
forall a. Bits a => a -> a -> a
(.|.)
        (Word64 -> Word127
forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word64 @Word127 Word64
a Word127 -> Int -> Word127
forall a. Bits a => a -> Int -> a
`Bits.shiftL` Int
63)
        (Word64 -> Word127
forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word64 @Word127 Word64
b Word127 -> Int -> Word127
forall a. Bits a => a -> Int -> a
`Bits.shiftR` Int
1))
    ((Word64, Word64) -> Word127)
-> (StdGen -> (Word64, Word64)) -> StdGen -> Word127
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SMGen -> (Word64, Word64)
unseedSMGen
    (SMGen -> (Word64, Word64))
-> (StdGen -> SMGen) -> StdGen -> (Word64, Word64)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StdGen -> SMGen
unStdGen

--------------------------------------------------------------------------------
-- Non-random contexts
--------------------------------------------------------------------------------

-- | Provides a stateless context for computations that must be non-random.
--
-- This type is useful for testing functions that require a 'MonadRandom'
-- context, but when actual randomness is not required or even desired.
--
newtype NonRandom a = NonRandom
    { NonRandom a -> a
runNonRandom :: a }
    deriving (NonRandom a -> NonRandom a -> Bool
(NonRandom a -> NonRandom a -> Bool)
-> (NonRandom a -> NonRandom a -> Bool) -> Eq (NonRandom a)
forall a. Eq a => NonRandom a -> NonRandom a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: NonRandom a -> NonRandom a -> Bool
$c/= :: forall a. Eq a => NonRandom a -> NonRandom a -> Bool
== :: NonRandom a -> NonRandom a -> Bool
$c== :: forall a. Eq a => NonRandom a -> NonRandom a -> Bool
Eq, (forall x. NonRandom a -> Rep (NonRandom a) x)
-> (forall x. Rep (NonRandom a) x -> NonRandom a)
-> Generic (NonRandom a)
forall x. Rep (NonRandom a) x -> NonRandom a
forall x. NonRandom a -> Rep (NonRandom a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall a x. Rep (NonRandom a) x -> NonRandom a
forall a x. NonRandom a -> Rep (NonRandom a) x
$cto :: forall a x. Rep (NonRandom a) x -> NonRandom a
$cfrom :: forall a x. NonRandom a -> Rep (NonRandom a) x
Generic, Eq (NonRandom a)
Eq (NonRandom a)
-> (NonRandom a -> NonRandom a -> Ordering)
-> (NonRandom a -> NonRandom a -> Bool)
-> (NonRandom a -> NonRandom a -> Bool)
-> (NonRandom a -> NonRandom a -> Bool)
-> (NonRandom a -> NonRandom a -> Bool)
-> (NonRandom a -> NonRandom a -> NonRandom a)
-> (NonRandom a -> NonRandom a -> NonRandom a)
-> Ord (NonRandom a)
NonRandom a -> NonRandom a -> Bool
NonRandom a -> NonRandom a -> Ordering
NonRandom a -> NonRandom a -> NonRandom a
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall a. Ord a => Eq (NonRandom a)
forall a. Ord a => NonRandom a -> NonRandom a -> Bool
forall a. Ord a => NonRandom a -> NonRandom a -> Ordering
forall a. Ord a => NonRandom a -> NonRandom a -> NonRandom a
min :: NonRandom a -> NonRandom a -> NonRandom a
$cmin :: forall a. Ord a => NonRandom a -> NonRandom a -> NonRandom a
max :: NonRandom a -> NonRandom a -> NonRandom a
$cmax :: forall a. Ord a => NonRandom a -> NonRandom a -> NonRandom a
>= :: NonRandom a -> NonRandom a -> Bool
$c>= :: forall a. Ord a => NonRandom a -> NonRandom a -> Bool
> :: NonRandom a -> NonRandom a -> Bool
$c> :: forall a. Ord a => NonRandom a -> NonRandom a -> Bool
<= :: NonRandom a -> NonRandom a -> Bool
$c<= :: forall a. Ord a => NonRandom a -> NonRandom a -> Bool
< :: NonRandom a -> NonRandom a -> Bool
$c< :: forall a. Ord a => NonRandom a -> NonRandom a -> Bool
compare :: NonRandom a -> NonRandom a -> Ordering
$ccompare :: forall a. Ord a => NonRandom a -> NonRandom a -> Ordering
$cp1Ord :: forall a. Ord a => Eq (NonRandom a)
Ord, Int -> NonRandom a -> ShowS
[NonRandom a] -> ShowS
NonRandom a -> String
(Int -> NonRandom a -> ShowS)
-> (NonRandom a -> String)
-> ([NonRandom a] -> ShowS)
-> Show (NonRandom a)
forall a. Show a => Int -> NonRandom a -> ShowS
forall a. Show a => [NonRandom a] -> ShowS
forall a. Show a => NonRandom a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [NonRandom a] -> ShowS
$cshowList :: forall a. Show a => [NonRandom a] -> ShowS
show :: NonRandom a -> String
$cshow :: forall a. Show a => NonRandom a -> String
showsPrec :: Int -> NonRandom a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> NonRandom a -> ShowS
Show)

instance Functor NonRandom where
    fmap :: (a -> b) -> NonRandom a -> NonRandom b
fmap = (a -> b) -> NonRandom a -> NonRandom b
coerce

instance Applicative NonRandom where
    liftA2 :: (a -> b -> c) -> NonRandom a -> NonRandom b -> NonRandom c
liftA2 = (a -> b -> c) -> NonRandom a -> NonRandom b -> NonRandom c
coerce
    pure :: a -> NonRandom a
pure = a -> NonRandom a
forall a. a -> NonRandom a
NonRandom
    <*> :: NonRandom (a -> b) -> NonRandom a -> NonRandom b
(<*>) = NonRandom (a -> b) -> NonRandom a -> NonRandom b
coerce

instance Monad NonRandom where
    NonRandom a
m >>= :: NonRandom a -> (a -> NonRandom b) -> NonRandom b
>>= a -> NonRandom b
k = a -> NonRandom b
k (NonRandom a -> a
forall a. NonRandom a -> a
runNonRandom NonRandom a
m)

instance MonadRandom NonRandom where
    getRandom :: NonRandom a
getRandom = a -> NonRandom a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> NonRandom a) -> a -> NonRandom a
forall a b. (a -> b) -> a -> b
$ (a, NonRandomGen) -> a
forall a b. (a, b) -> a
fst ((a, NonRandomGen) -> a) -> (a, NonRandomGen) -> a
forall a b. (a -> b) -> a -> b
$ NonRandomGen -> (a, NonRandomGen)
forall a g. (Random a, RandomGen g) => g -> (a, g)
random NonRandomGen
NonRandomGen
    getRandomR :: (a, a) -> NonRandom a
getRandomR (a, a)
r = a -> NonRandom a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> NonRandom a) -> a -> NonRandom a
forall a b. (a -> b) -> a -> b
$ (a, NonRandomGen) -> a
forall a b. (a, b) -> a
fst ((a, NonRandomGen) -> a) -> (a, NonRandomGen) -> a
forall a b. (a -> b) -> a -> b
$ (a, a) -> NonRandomGen -> (a, NonRandomGen)
forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
randomR (a, a)
r NonRandomGen
NonRandomGen
    getRandomRs :: (a, a) -> NonRandom [a]
getRandomRs (a, a)
r = [a] -> NonRandom [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([a] -> NonRandom [a]) -> [a] -> NonRandom [a]
forall a b. (a -> b) -> a -> b
$ (a, a) -> NonRandomGen -> [a]
forall a g. (Random a, RandomGen g) => (a, a) -> g -> [a]
randomRs (a, a)
r NonRandomGen
NonRandomGen
    getRandoms :: NonRandom [a]
getRandoms = [a] -> NonRandom [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([a] -> NonRandom [a]) -> [a] -> NonRandom [a]
forall a b. (a -> b) -> a -> b
$ NonRandomGen -> [a]
forall a g. (Random a, RandomGen g) => g -> [a]
randoms NonRandomGen
NonRandomGen

-- | Provides a stateless and non-random implementation of 'RandomGen'
--
data NonRandomGen = NonRandomGen

instance RandomGen NonRandomGen where
    genRange :: NonRandomGen -> (Int, Int)
genRange NonRandomGen
NonRandomGen = (Int
forall a. Bounded a => a
minBound, Int
forall a. Bounded a => a
maxBound)
    next :: NonRandomGen -> (Int, NonRandomGen)
next NonRandomGen
NonRandomGen = (Int
0, NonRandomGen
NonRandomGen)
    split :: NonRandomGen -> (NonRandomGen, NonRandomGen)
split NonRandomGen
NonRandomGen = (NonRandomGen
NonRandomGen, NonRandomGen
NonRandomGen)