{-# LANGUAGE TemplateHaskell #-}

module Polysemy.Random
  ( -- * Effect
    Random (..)

    -- * Actions
  , random
  , randomR

    -- * Interpretations
  , runRandom
  , runRandomIO
  ) where

import           Polysemy
import           Polysemy.State
import qualified System.Random as R

------------------------------------------------------------------------------
-- | An effect capable of providing 'R.Random' values.
data Random m a where
  Random :: R.Random x => Random m x
  RandomR :: R.Random x => (x, x) -> Random m x

makeSem ''Random


------------------------------------------------------------------------------
-- | Run a 'Random' effect with an explicit 'R.RandomGen'.
runRandom
    :: forall q r a
     . R.RandomGen q
    => q
    -> Sem (Random ': r) a
    -> Sem r (q, a)
runRandom :: q -> Sem (Random : r) a -> Sem r (q, a)
runRandom q
q = q -> Sem (State q : r) a -> Sem r (q, a)
forall s (r :: [(* -> *) -> * -> *]) a.
s -> Sem (State s : r) a -> Sem r (s, a)
runState q
q (Sem (State q : r) a -> Sem r (q, a))
-> (Sem (Random : r) a -> Sem (State q : r) a)
-> Sem (Random : r) a
-> Sem r (q, a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall (rInitial :: [(* -> *) -> * -> *]) x.
 Random (Sem rInitial) x -> Sem (State q : r) x)
-> Sem (Random : r) a -> Sem (State q : r) a
forall (e1 :: (* -> *) -> * -> *) (e2 :: (* -> *) -> * -> *)
       (r :: [(* -> *) -> * -> *]) a.
FirstOrder e1 "reinterpret" =>
(forall (rInitial :: [(* -> *) -> * -> *]) x.
 e1 (Sem rInitial) x -> Sem (e2 : r) x)
-> Sem (e1 : r) a -> Sem (e2 : r) a
reinterpret (\case
  Random (Sem rInitial) x
Random -> do
    ~(x
a, q
q') <- (q -> (x, q)) -> Sem (State q : r) (x, q)
forall s a (r :: [(* -> *) -> * -> *]).
Member (State s) r =>
(s -> a) -> Sem r a
gets @q q -> (x, q)
forall a g. (Random a, RandomGen g) => g -> (a, g)
R.random
    q -> Sem (State q : r) ()
forall s (r :: [(* -> *) -> * -> *]).
Member (State s) r =>
s -> Sem r ()
put q
q'
    x -> Sem (State q : r) x
forall (f :: * -> *) a. Applicative f => a -> f a
pure x
a
  RandomR r -> do
    ~(x
a, q
q') <- forall a (r :: [(* -> *) -> * -> *]).
Member (State q) r =>
(q -> a) -> Sem r a
forall s a (r :: [(* -> *) -> * -> *]).
Member (State s) r =>
(s -> a) -> Sem r a
gets @q ((q -> (x, q)) -> Sem (State q : r) (x, q))
-> (q -> (x, q)) -> Sem (State q : r) (x, q)
forall a b. (a -> b) -> a -> b
$ (x, x) -> q -> (x, q)
forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
R.randomR (x, x)
r
    q -> Sem (State q : r) ()
forall s (r :: [(* -> *) -> * -> *]).
Member (State s) r =>
s -> Sem r ()
put q
q'
    x -> Sem (State q : r) x
forall (f :: * -> *) a. Applicative f => a -> f a
pure x
a
                                       )
{-# INLINE runRandom #-}


------------------------------------------------------------------------------
-- | Run a 'Random' effect by using the 'IO' random generator.
runRandomIO :: Member (Embed IO) r => Sem (Random ': r) a -> Sem r a
runRandomIO :: Sem (Random : r) a -> Sem r a
runRandomIO Sem (Random : r) a
m = do
  StdGen
q <- IO StdGen -> Sem r StdGen
forall (m :: * -> *) (r :: [(* -> *) -> * -> *]) a.
Member (Embed m) r =>
m a -> Sem r a
embed @IO IO StdGen
forall (m :: * -> *). MonadIO m => m StdGen
R.newStdGen
  (StdGen, a) -> a
forall a b. (a, b) -> b
snd ((StdGen, a) -> a) -> Sem r (StdGen, a) -> Sem r a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StdGen -> Sem (Random : r) a -> Sem r (StdGen, a)
forall q (r :: [(* -> *) -> * -> *]) a.
RandomGen q =>
q -> Sem (Random : r) a -> Sem r (q, a)
runRandom StdGen
q Sem (Random : r) a
m
{-# INLINE runRandomIO #-}