{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE UndecidableInstances #-}
module Polysemy.RandomFu
(
RandomFu (..)
, sampleRVar
, getRandomPrim
, sampleDist
, runRandomSource
, runRandomIO
, runRandomIOPureMT
, absorbMonadRandom
)
where
import Polysemy
import Polysemy.MTL
import Data.IORef ( newIORef )
import qualified Data.Random as R
import qualified Data.Random.Source as R
import qualified Data.Random.Internal.Source as R
import qualified Data.Random.Source.PureMT as R
import Control.Monad.IO.Class ( MonadIO(..) )
data RandomFu m r where
SampleRVar :: R.RVar t -> RandomFu m t
GetRandomPrim :: R.Prim t -> RandomFu m t
makeSem ''RandomFu
sampleDist
:: (Member RandomFu r, R.Distribution d t) => d t -> Sem r t
sampleDist = sampleRVar . R.rvar
{-# INLINEABLE sampleDist #-}
runRandomSource
:: forall s r a
. R.RandomSource (Sem r) s
=> s
-> Sem (RandomFu ': r) a
-> Sem r a
runRandomSource source = interpret $ \case
SampleRVar rv -> R.runRVar (R.sample rv) source
GetRandomPrim pt -> R.runRVar (R.getRandomPrim pt) source
{-# INLINEABLE runRandomSource #-}
runRandomIO
:: forall r a
. MonadIO (Sem r)
=> Sem (RandomFu ': r) a
-> Sem r a
runRandomIO = interpret $ \case
SampleRVar rv -> liftIO $ R.sample rv
GetRandomPrim pt -> liftIO $ R.getRandomPrim pt
{-# INLINEABLE runRandomIO #-}
runRandomIOPureMT
:: MonadIO (Sem r)
=> R.PureMT
-> Sem (RandomFu ': r) a
-> Sem r a
runRandomIOPureMT source re =
liftIO (newIORef source) >>= flip runRandomSource re
{-# INLINEABLE runRandomIOPureMT #-}
absorbMonadRandom
:: Member RandomFu r => (R.MonadRandom (Sem r) => Sem r a) -> Sem r a
absorbMonadRandom = absorb @R.MonadRandom
{-# INLINEABLE absorbMonadRandom #-}
instance ReifiableConstraint1 (R.MonadRandom) where
data Dict1 R.MonadRandom m = MonadRandom
{
getRandomPrim_ :: forall t. R.Prim t -> m t
}
reifiedInstance = Sub Dict
$(R.monadRandom [d|
instance ( Monad m
, Reifies s' (Dict1 R.MonadRandom m)
) => R.MonadRandom (ConstrainedAction R.MonadRandom m s') where
getRandomPrim t = ConstrainedAction
$ getRandomPrim_ (reflect $ Proxy @s') t
{-# INLINEABLE getRandomPrim #-}
|])
instance Member RandomFu r => IsCanonicalEffect R.MonadRandom r where
canonicalDictionary = MonadRandom getRandomPrim
{-# INLINEABLE canonicalDictionary #-}