module Foundation.Random.ChaChaDRG
    ( State(..)
    , keySize
    ) where

import           Foundation.Class.Storable (peek)
import           Basement.Imports
import           Basement.Types.OffsetSize
import           Basement.Monad
import           Foundation.Random.Class
import           Foundation.Random.DRG
import qualified Basement.UArray as A
import qualified Basement.UArray.Mutable as A
import           GHC.ST
import qualified Foreign.Marshal.Alloc (alloca)

-- | RNG based on ChaCha core.
--
-- The algorithm is identical to the arc4random found in recent BSDs,
-- namely a ChaCha core provide 64 bytes of random from 32 bytes of
-- key.
newtype State = State (UArray Word8)

instance RandomGen State where
    randomNew :: forall (m :: * -> *). MonadRandom m => m State
randomNew = UArray Word8 -> State
State forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *).
MonadRandom m =>
CountOf Word8 -> m (UArray Word8)
getRandomBytes CountOf Word8
keySize
    randomNewFrom :: UArray Word8 -> Maybe State
randomNewFrom UArray Word8
bs
        | forall ty. UArray ty -> CountOf ty
A.length UArray Word8
bs forall a. Eq a => a -> a -> Bool
== CountOf Word8
keySize = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ UArray Word8 -> State
State UArray Word8
bs
        | Bool
otherwise              = forall a. Maybe a
Nothing
    randomGenerate :: CountOf Word8 -> State -> (UArray Word8, State)
randomGenerate = CountOf Word8 -> State -> (UArray Word8, State)
generate
    randomGenerateWord64 :: State -> (Word64, State)
randomGenerateWord64 = State -> (Word64, State)
generateWord64
    randomGenerateF32 :: State -> (Float, State)
randomGenerateF32 = State -> (Float, State)
generateF32
    randomGenerateF64 :: State -> (Double, State)
randomGenerateF64 = State -> (Double, State)
generateF64

keySize :: CountOf Word8
keySize :: CountOf Word8
keySize = CountOf Word8
32

generate :: CountOf Word8 -> State -> (UArray Word8, State)
generate :: CountOf Word8 -> State -> (UArray Word8, State)
generate CountOf Word8
n (State UArray Word8
key) = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
    MUArray Word8 s
dst    <- forall (prim :: * -> *) ty.
(PrimMonad prim, PrimType ty) =>
CountOf ty -> prim (MUArray ty (PrimState prim))
A.newPinned CountOf Word8
n
    MUArray Word8 s
newKey <- forall (prim :: * -> *) ty.
(PrimMonad prim, PrimType ty) =>
CountOf ty -> prim (MUArray ty (PrimState prim))
A.newPinned CountOf Word8
keySize
    forall (prim :: * -> *) ty a.
(PrimMonad prim, PrimType ty) =>
MUArray ty (PrimState prim) -> (Ptr ty -> prim a) -> prim a
A.withMutablePtr MUArray Word8 s
dst        forall a b. (a -> b) -> a -> b
$ \Ptr Word8
dstP    ->
        forall (prim :: * -> *) ty a.
(PrimMonad prim, PrimType ty) =>
MUArray ty (PrimState prim) -> (Ptr ty -> prim a) -> prim a
A.withMutablePtr MUArray Word8 s
newKey forall a b. (a -> b) -> a -> b
$ \Ptr Word8
newKeyP ->
        forall ty (prim :: * -> *) a.
(PrimMonad prim, PrimType ty) =>
UArray ty -> (Ptr ty -> prim a) -> prim a
A.withPtr UArray Word8
key           forall a b. (a -> b) -> a -> b
$ \Ptr Word8
keyP    -> do
            Word32
_ <- forall (prim :: * -> *) a. PrimMonad prim => IO a -> prim a
unsafePrimFromIO forall a b. (a -> b) -> a -> b
$ Ptr Word8 -> Ptr Word8 -> Ptr Word8 -> CountOf Word8 -> IO Word32
c_rngv1_generate Ptr Word8
newKeyP Ptr Word8
dstP Ptr Word8
keyP CountOf Word8
n
            forall (m :: * -> *) a. Monad m => a -> m a
return ()
    (,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (prim :: * -> *) ty.
PrimMonad prim =>
MUArray ty (PrimState prim) -> prim (UArray ty)
A.unsafeFreeze MUArray Word8 s
dst
        forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (UArray Word8 -> State
State forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (prim :: * -> *) ty.
PrimMonad prim =>
MUArray ty (PrimState prim) -> prim (UArray ty)
A.unsafeFreeze MUArray Word8 s
newKey)

generateWord64 :: State -> (Word64, State)
generateWord64 :: State -> (Word64, State)
generateWord64 (State UArray Word8
key) = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ forall (prim :: * -> *) a. PrimMonad prim => IO a -> prim a
unsafePrimFromIO forall a b. (a -> b) -> a -> b
$
    forall a b. Storable a => (Ptr a -> IO b) -> IO b
Foreign.Marshal.Alloc.alloca forall a b. (a -> b) -> a -> b
$ \Ptr Word64
dst -> do
        MUArray Word8 RealWorld
newKey <- forall (prim :: * -> *) ty.
(PrimMonad prim, PrimType ty) =>
CountOf ty -> prim (MUArray ty (PrimState prim))
A.newPinned CountOf Word8
keySize
        forall (prim :: * -> *) ty a.
(PrimMonad prim, PrimType ty) =>
MUArray ty (PrimState prim) -> (Ptr ty -> prim a) -> prim a
A.withMutablePtr MUArray Word8 RealWorld
newKey forall a b. (a -> b) -> a -> b
$ \Ptr Word8
newKeyP ->
          forall ty (prim :: * -> *) a.
(PrimMonad prim, PrimType ty) =>
UArray ty -> (Ptr ty -> prim a) -> prim a
A.withPtr UArray Word8
key           forall a b. (a -> b) -> a -> b
$ \Ptr Word8
keyP  ->
            Ptr Word8 -> Ptr Word64 -> Ptr Word8 -> IO Word32
c_rngv1_generate_word64 Ptr Word8
newKeyP Ptr Word64
dst Ptr Word8
keyP forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> forall (m :: * -> *) a. Monad m => a -> m a
return ()
        (,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Storable a => Ptr a -> IO a
peek Ptr Word64
dst forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (UArray Word8 -> State
State forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (prim :: * -> *) ty.
PrimMonad prim =>
MUArray ty (PrimState prim) -> prim (UArray ty)
A.unsafeFreeze MUArray Word8 RealWorld
newKey)

generateF32 :: State -> (Float, State)
generateF32 :: State -> (Float, State)
generateF32 (State UArray Word8
key) = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ forall (prim :: * -> *) a. PrimMonad prim => IO a -> prim a
unsafePrimFromIO forall a b. (a -> b) -> a -> b
$
    forall a b. Storable a => (Ptr a -> IO b) -> IO b
Foreign.Marshal.Alloc.alloca forall a b. (a -> b) -> a -> b
$ \Ptr Float
dst -> do
        MUArray Word8 RealWorld
newKey <- forall (prim :: * -> *) ty.
(PrimMonad prim, PrimType ty) =>
CountOf ty -> prim (MUArray ty (PrimState prim))
A.newPinned CountOf Word8
keySize
        forall (prim :: * -> *) ty a.
(PrimMonad prim, PrimType ty) =>
MUArray ty (PrimState prim) -> (Ptr ty -> prim a) -> prim a
A.withMutablePtr MUArray Word8 RealWorld
newKey forall a b. (a -> b) -> a -> b
$ \Ptr Word8
newKeyP ->
          forall ty (prim :: * -> *) a.
(PrimMonad prim, PrimType ty) =>
UArray ty -> (Ptr ty -> prim a) -> prim a
A.withPtr UArray Word8
key           forall a b. (a -> b) -> a -> b
$ \Ptr Word8
keyP  ->
            Ptr Word8 -> Ptr Float -> Ptr Word8 -> IO Word32
c_rngv1_generate_f32 Ptr Word8
newKeyP Ptr Float
dst Ptr Word8
keyP forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> forall (m :: * -> *) a. Monad m => a -> m a
return ()
        (,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Storable a => Ptr a -> IO a
peek Ptr Float
dst forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (UArray Word8 -> State
State forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (prim :: * -> *) ty.
PrimMonad prim =>
MUArray ty (PrimState prim) -> prim (UArray ty)
A.unsafeFreeze MUArray Word8 RealWorld
newKey)

generateF64 :: State -> (Double, State)
generateF64 :: State -> (Double, State)
generateF64 (State UArray Word8
key) = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ forall (prim :: * -> *) a. PrimMonad prim => IO a -> prim a
unsafePrimFromIO forall a b. (a -> b) -> a -> b
$
    forall a b. Storable a => (Ptr a -> IO b) -> IO b
Foreign.Marshal.Alloc.alloca forall a b. (a -> b) -> a -> b
$ \Ptr Double
dst -> do
        MUArray Word8 RealWorld
newKey <- forall (prim :: * -> *) ty.
(PrimMonad prim, PrimType ty) =>
CountOf ty -> prim (MUArray ty (PrimState prim))
A.newPinned CountOf Word8
keySize
        forall (prim :: * -> *) ty a.
(PrimMonad prim, PrimType ty) =>
MUArray ty (PrimState prim) -> (Ptr ty -> prim a) -> prim a
A.withMutablePtr MUArray Word8 RealWorld
newKey forall a b. (a -> b) -> a -> b
$ \Ptr Word8
newKeyP ->
          forall ty (prim :: * -> *) a.
(PrimMonad prim, PrimType ty) =>
UArray ty -> (Ptr ty -> prim a) -> prim a
A.withPtr UArray Word8
key           forall a b. (a -> b) -> a -> b
$ \Ptr Word8
keyP  ->
            Ptr Word8 -> Ptr Double -> Ptr Word8 -> IO Word32
c_rngv1_generate_f64 Ptr Word8
newKeyP Ptr Double
dst Ptr Word8
keyP forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> forall (m :: * -> *) a. Monad m => a -> m a
return ()
        (,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Storable a => Ptr a -> IO a
peek Ptr Double
dst forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (UArray Word8 -> State
State forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (prim :: * -> *) ty.
PrimMonad prim =>
MUArray ty (PrimState prim) -> prim (UArray ty)
A.unsafeFreeze MUArray Word8 RealWorld
newKey)

-- return 0 on success, !0 for failure
foreign import ccall unsafe "foundation_rngV1_generate"
   c_rngv1_generate :: Ptr Word8 -- new key
                    -> Ptr Word8 -- destination
                    -> Ptr Word8 -- current key
                    -> CountOf Word8 -- number of bytes to generate
                    -> IO Word32

foreign import ccall unsafe "foundation_rngV1_generate_word64"
   c_rngv1_generate_word64 :: Ptr Word8  -- new key
                           -> Ptr Word64 -- destination
                           -> Ptr Word8  -- current key
                           -> IO Word32

foreign import ccall unsafe "foundation_rngV1_generate_f32"
   c_rngv1_generate_f32 :: Ptr Word8  -- new key
                        -> Ptr Float -- destination
                        -> Ptr Word8  -- current key
                        -> IO Word32

foreign import ccall unsafe "foundation_rngV1_generate_f64"
   c_rngv1_generate_f64 :: Ptr Word8  -- new key
                        -> Ptr Double -- destination
                        -> Ptr Word8  -- current key
                        -> IO Word32