{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

module Test.ImpSpec.Random where

import Control.Monad (replicateM)
import Control.Monad.Reader (MonadReader (ask), ReaderT (..))
import Data.ByteString (ByteString)
import Data.ByteString.Short (ShortByteString)
import qualified System.Random.Stateful as R
import qualified Test.QuickCheck as QC (Arbitrary (arbitrary))
import Test.QuickCheck.GenT (MonadGen (liftGen))

class R.StatefulGen g m => HasStatefulGen g m | m -> g where
  askStatefulGen :: m g
  default askStatefulGen :: MonadReader g m => m g
  askStatefulGen = m g
forall r (m :: * -> *). MonadReader r m => m r
ask

class HasGenEnv env g | env -> g where
  getGenEnv :: env -> g

instance HasGenEnv g g where
  getGenEnv :: g -> g
getGenEnv = g -> g
forall g. g -> g
id

instance
  (HasGenEnv env g, R.StatefulGen g (ReaderT env m), Monad m) =>
  HasStatefulGen g (ReaderT env m)
  where
  askStatefulGen :: ReaderT env m g
askStatefulGen = (env -> m g) -> ReaderT env m g
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT (g -> m g
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (g -> m g) -> (env -> g) -> env -> m g
forall b c a. (b -> c) -> (a -> b) -> a -> c
. env -> g
forall env g. HasGenEnv env g => env -> g
getGenEnv)

uniformM ::
  ( HasStatefulGen g m
  , R.Uniform a
  ) =>
  m a
uniformM :: forall g (m :: * -> *) a. (HasStatefulGen g m, Uniform a) => m a
uniformM = m g
forall g (m :: * -> *). HasStatefulGen g m => m g
askStatefulGen m g -> (g -> m a) -> m a
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= g -> m a
forall a g (m :: * -> *). (Uniform a, StatefulGen g m) => g -> m a
forall g (m :: * -> *). StatefulGen g m => g -> m a
R.uniformM
{-# INLINE uniformM #-}

uniformRM ::
  ( HasStatefulGen g m
  , R.UniformRange a
  ) =>
  (a, a) ->
  m a
uniformRM :: forall g (m :: * -> *) a.
(HasStatefulGen g m, UniformRange a) =>
(a, a) -> m a
uniformRM (a, a)
r = m g
forall g (m :: * -> *). HasStatefulGen g m => m g
askStatefulGen m g -> (g -> m a) -> m a
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (a, a) -> g -> m a
forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
forall g (m :: * -> *). StatefulGen g m => (a, a) -> g -> m a
R.uniformRM (a, a)
r
{-# INLINE uniformRM #-}

uniformListM ::
  ( HasStatefulGen g m
  , R.Uniform a
  ) =>
  Int ->
  m [a]
uniformListM :: forall g (m :: * -> *) a.
(HasStatefulGen g m, Uniform a) =>
Int -> m [a]
uniformListM Int
n = m g
forall g (m :: * -> *). HasStatefulGen g m => m g
askStatefulGen m g -> (g -> m [a]) -> m [a]
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> g -> m [a]
forall g (m :: * -> *) a.
(StatefulGen g m, Uniform a) =>
Int -> g -> m [a]
R.uniformListM Int
n
{-# INLINE uniformListM #-}

uniformListRM ::
  (HasStatefulGen g m, R.UniformRange a) =>
  (a, a) ->
  Int ->
  m [a]
uniformListRM :: forall g (m :: * -> *) a.
(HasStatefulGen g m, UniformRange a) =>
(a, a) -> Int -> m [a]
uniformListRM (a, a)
r Int
n = m g
forall g (m :: * -> *). HasStatefulGen g m => m g
askStatefulGen m g -> (g -> m [a]) -> m [a]
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> m a -> m [a]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
n (m a -> m [a]) -> (g -> m a) -> g -> m [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, a) -> g -> m a
forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
forall g (m :: * -> *). StatefulGen g m => (a, a) -> g -> m a
R.uniformRM (a, a)
r
{-# INLINE uniformListRM #-}

uniformByteStringM :: HasStatefulGen a m => Int -> m ByteString
uniformByteStringM :: forall a (m :: * -> *). HasStatefulGen a m => Int -> m ByteString
uniformByteStringM Int
n = m a
forall g (m :: * -> *). HasStatefulGen g m => m g
askStatefulGen m a -> (a -> m ByteString) -> m ByteString
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> a -> m ByteString
forall g (m :: * -> *). StatefulGen g m => Int -> g -> m ByteString
R.uniformByteStringM Int
n
{-# INLINE uniformByteStringM #-}

uniformShortByteStringM :: HasStatefulGen a m => Int -> m ShortByteString
uniformShortByteStringM :: forall a (m :: * -> *).
HasStatefulGen a m =>
Int -> m ShortByteString
uniformShortByteStringM Int
n = m a
forall g (m :: * -> *). HasStatefulGen g m => m g
askStatefulGen m a -> (a -> m ShortByteString) -> m ShortByteString
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> a -> m ShortByteString
forall g (m :: * -> *).
StatefulGen g m =>
Int -> g -> m ShortByteString
R.uniformShortByteString Int
n
{-# INLINE uniformShortByteStringM #-}

-- | Lifted version of `QC.arbitrary`.
arbitrary :: (QC.Arbitrary a, MonadGen m) => m a
arbitrary :: forall a (m :: * -> *). (Arbitrary a, MonadGen m) => m a
arbitrary = Gen a -> m a
forall a. Gen a -> m a
forall (g :: * -> *) a. MonadGen g => Gen a -> g a
liftGen Gen a
forall a. Arbitrary a => Gen a
QC.arbitrary