Random.hs

{-# LANGUAGE DefaultSignatures          #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MagicHash                  #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE QuantifiedConstraints      #-}
{-# LANGUAGE RankNTypes                 #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE StandaloneDeriving         #-}
{-# LANGUAGE TypeFamilies               #-}
{-# LANGUAGE UnboxedTuples              #-}
{-# LANGUAGE UndecidableInstances       #-}
-- |
module Random2 where

import Control.Monad.Primitive
import Control.Monad.ST
import Control.Monad.State.Strict
import Control.Monad.Trans.Reader
import Data.Primitive.ByteArray
import Data.Primitive.MutVar
import Data.Word
import Data.STRef

import Data.Primitive.MutVar
import GHC.Exts (State#)


----------------------------------------------------------------
-- Random
----------------------------------------------------------------

class (Monad1 (RandM g)) => PRNG g where
  type RandM g :: * -> * -> *
  runRand      :: Rand g a -> g -> (a, g)
  freezeGen    :: Rand g g
  thawGen      :: g -> Rand g ()
  genWord64    :: Rand g Word64


data family MutPRNG g :: * -> *


-- | Monad for sequencing PRNG transitions that allow to use 
newtype MRand g s a = MRand
  { unMRand :: RandM g s a
  }

instance Monad1 (RandM g) => Functor (MRand g s) where
  fmap f = MRand . fmap1 f . unMRand

instance Monad1 (RandM g) => Applicative (MRand g s) where
  pure = MRand . pure1
  MRand f <*> MRand g = MRand $ ap1 f g

instance Monad1 (RandM g) => Monad (MRand g s) where
  return = pure
  MRand m >>= f = MRand $ bind1 m (unMRand . f)

instance Monad1 (RandM g) => PrimMonad (MRand g s) where
  type PrimState (MRand g s) = s
  primitive = MRand . primitive1


-- | Monad for sequencing generator states which hides in-place
--   mutation
newtype Rand g a = Rand
  { unRand :: forall s. MRand g s a
  }

instance Monad1 (RandM g) => Functor (Rand g) where
  fmap f (Rand m) = Rand (fmap f m)

instance Monad1 (RandM g) => Applicative (Rand g) where
  pure a = Rand (pure a)
  Rand f <*> Rand g = Rand (f <*> g)
instance Monad1 (RandM g) => Monad (Rand g) where
  return = pure
  Rand m >>= f = Rand $ m >>= unRand . f




class Monad1 m where
  fmap1      :: (a -> b) -> m s a -> m s b 
  pure1      :: a -> m s a
  ap1        :: m s (a -> b) -> m s a -> m s b
  bind1      :: m s a -> (a -> m s b) -> m s b
  primitive1 :: (State# s -> (# State# s, a #)) -> m s a

-- | Base monad for pure PRNG
newtype RandPure g s a = RandPure
  { unRandPure :: StateT g (ST s) a
  }
  deriving (Functor,Applicative,Monad,MonadState g)

runRandPure :: (RandM g ~ RandPure g) => Rand g a -> g -> (a, g)
runRandPure m g = runST (runStateT (unRandPure (unMRand (unRand m))) g)

instance PrimMonad (RandPure g s) where
  type PrimState (RandPure g s) = s
  primitive = RandPure . primitive

instance Monad1 (RandPure g) where
  fmap1      = fmap
  pure1      = pure
  ap1        = (<*>)
  bind1      = (>>=)
  primitive1 = primitive



-- | Base monad for stateful PRNG
newtype RandST g s a = RandST
  { unRandST :: ReaderT (MutPRNG g s) (ST s) a
  }
  deriving (Functor,Applicative,Monad)

instance PrimMonad (RandST g s) where
  type PrimState (RandST g s) = s
  primitive = RandST . primitive

instance Monad1 (RandST g) where
  fmap1      = fmap
  pure1      = pure
  ap1        = (<*>)
  bind1      = (>>=)
  primitive1 = primitive


----------------------------------------------------------------
-- Lifting to MonadRandom
----------------------------------------------------------------

-- class Monad m => MonadRandom g m where



----------------------------------------------------------------
-- Examples
----------------------------------------------------------------

newtype RNG'Pure = RNG'Pure Word64

instance PRNG RNG'Pure where
  type RandM RNG'Pure = RandPure RNG'Pure
  genWord64 = Rand $ MRand $ RandPure $ state $ \(RNG'Pure w) -> (w, RNG'Pure (w+1))
  runRand   = runRandPure
  freezeGen = Rand $ MRand   get
  thawGen g = Rand $ MRand $ put g

newtype RNG'Mutable = RNG'Mutable Word64

newtype instance MutPRNG RNG'Mutable s = RNG'Ref (STRef s Word64)

instance PRNG RNG'Mutable where
  type RandM RNG'Mutable = RandST RNG'Mutable
  genWord64 = Rand $ MRand $ RandST $ ReaderT $ \(RNG'Ref ref) -> do
    w <- readSTRef ref
    writeSTRef ref $! w + 1
    return w
  freezeGen = Rand $ MRand $ RandST $ do
    RNG'Ref ref <- ask
    lift $ RNG'Mutable <$> readSTRef ref
  thawGen (RNG'Mutable w) = Rand $ MRand $ RandST $ do
    RNG'Ref ref <- ask
    lift $ writeSTRef ref w
  runRand m (RNG'Mutable w) = runST $ do
    ref <- newSTRef w
    a   <- runReaderT (unRandST (unMRand (unRand m))) (RNG'Ref ref)
    w'  <- readSTRef ref
    return (a, RNG'Mutable w')


internalMutability :: PRNG g => Rand g Word64
internalMutability = Rand $ do
  w1  <- unRand genWord64
  ref <- newMutVar w1
  w2  <- unRand genWord64
  modifyMutVar' ref (+w2)
  readMutVar ref