MWC.hs

{-# LANGUAGE GeneralizedNewtypeDeriving, MultiParamTypeClasses, ScopedTypeVariables, FlexibleInstances, UndecidableInstances, ExistentialQuantification, TypeApplications, StandaloneDeriving, DeriveFunctor, TypeOperators #-}

module Control.Effect.Random.MWC
  ( Random (..)
  , Control.Effect.Random.MWC.uniform
  , Control.Effect.Random.MWC.uniformR
  , RandomC (..)
  , runRandom
  , Variate
  , Has
  ) where

import Control.Carrier.Reader
import System.Random.MWC (Variate (..))
import qualified System.Random.MWC as MWC
import Control.Algebra
import Control.Monad.IO.Class

data Random m k
  = forall a . Variate a => Uniform (a -> m k)
  | forall a . Variate a => UniformR (a, a) (a -> m k)

deriving instance Functor m => Functor (Random m)

instance Effect Random where
  thread ctx handler (Uniform k) = Uniform (handler . (<$ ctx) . k)
  thread ctx handler (UniformR r k) = UniformR r (handler . (<$ ctx) . k)

uniform :: forall a sig m . (MWC.Variate a, Has Random sig m) => m a
uniform = send (Uniform pure)

uniformR :: forall a sig m . (MWC.Variate a, Has Random sig m) => (a, a) -> m a
uniformR r = send (UniformR r pure)

newtype RandomC m a = RandomC (ReaderC MWC.GenIO m a)
  deriving (Applicative, Functor, Monad, MonadIO)

instance (Algebra sig m, MonadIO m) => Algebra (Random :+: sig) (RandomC m) where
  alg (L act) = do
    gen <- RandomC ask
    case act of
      Uniform k    -> liftIO (MWC.uniform gen)    >>= k
      UniformR r k -> liftIO (MWC.uniformR r gen) >>= k
  alg (R other) = RandomC (handleCoercible other)

runRandom :: MonadIO m => RandomC m a -> m a
runRandom (RandomC r) = do
  gen <- liftIO MWC.createSystemRandom
  runReader gen r