-- |
-- Module      : Crypto.Random.ChaChaDRG
-- License     : BSD-style
-- Maintainer  : Vincent Hanquez <vincent@snarc.org>
-- Stability   : stable
-- Portability : good
--
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
module Crypto.Random.ChaChaDRG
    ( ChaChaDRG
    , initialize
    , initializeWords
    ) where

import           Crypto.Random.Types
import           Crypto.Internal.Imports
import           Crypto.Internal.ByteArray (ByteArray, ByteArrayAccess, ScrubbedBytes)
import qualified Crypto.Internal.ByteArray as B
import           Foreign.Storable (pokeElemOff)

import qualified Crypto.Cipher.ChaCha as C

instance DRG ChaChaDRG where
    randomBytesGenerate :: Int -> ChaChaDRG -> (byteArray, ChaChaDRG)
randomBytesGenerate = Int -> ChaChaDRG -> (byteArray, ChaChaDRG)
forall byteArray.
ByteArray byteArray =>
Int -> ChaChaDRG -> (byteArray, ChaChaDRG)
generate

-- | ChaCha Deterministic Random Generator
newtype ChaChaDRG = ChaChaDRG C.StateSimple
    deriving (ChaChaDRG -> ()
(ChaChaDRG -> ()) -> NFData ChaChaDRG
forall a. (a -> ()) -> NFData a
rnf :: ChaChaDRG -> ()
$crnf :: ChaChaDRG -> ()
NFData)

-- | Initialize a new ChaCha context with the number of rounds,
-- the key and the nonce associated.
initialize :: ByteArrayAccess seed
           => seed        -- ^ 40 bytes of seed
           -> ChaChaDRG   -- ^ the initial ChaCha state
initialize :: seed -> ChaChaDRG
initialize seed
seed = StateSimple -> ChaChaDRG
ChaChaDRG (StateSimple -> ChaChaDRG) -> StateSimple -> ChaChaDRG
forall a b. (a -> b) -> a -> b
$ seed -> StateSimple
forall seed. ByteArrayAccess seed => seed -> StateSimple
C.initializeSimple seed
seed

-- | Initialize a new ChaCha context from 5-tuple of words64.
-- This interface is useful when creating a RNG out of tests generators (e.g. QuickCheck).
initializeWords :: (Word64, Word64, Word64, Word64, Word64) -> ChaChaDRG
initializeWords :: (Word64, Word64, Word64, Word64, Word64) -> ChaChaDRG
initializeWords (Word64
a,Word64
b,Word64
c,Word64
d,Word64
e) = ScrubbedBytes -> ChaChaDRG
forall seed. ByteArrayAccess seed => seed -> ChaChaDRG
initialize (Int -> (Ptr Word64 -> IO ()) -> ScrubbedBytes
forall a p. ByteArray a => Int -> (Ptr p -> IO ()) -> a
B.allocAndFreeze Int
40 Ptr Word64 -> IO ()
fill :: ScrubbedBytes)
  where fill :: Ptr Word64 -> IO ()
fill Ptr Word64
s = ((Int, Word64) -> IO ()) -> [(Int, Word64)] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ((Int -> Word64 -> IO ()) -> (Int, Word64) -> IO ()
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (Ptr Word64 -> Int -> Word64 -> IO ()
forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff Ptr Word64
s)) [(Int
0,Word64
a), (Int
1,Word64
b), (Int
2,Word64
c), (Int
3,Word64
d), (Int
4,Word64
e)]

generate :: ByteArray output => Int -> ChaChaDRG -> (output, ChaChaDRG)
generate :: Int -> ChaChaDRG -> (output, ChaChaDRG)
generate Int
nbBytes st :: ChaChaDRG
st@(ChaChaDRG StateSimple
prevSt)
    | Int
nbBytes Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = (output
forall a. ByteArray a => a
B.empty, ChaChaDRG
st)
    | Bool
otherwise    = let (output
output, StateSimple
newSt) = StateSimple -> Int -> (output, StateSimple)
forall ba. ByteArray ba => StateSimple -> Int -> (ba, StateSimple)
C.generateSimple StateSimple
prevSt Int
nbBytes in (output
output, StateSimple -> ChaChaDRG
ChaChaDRG StateSimple
newSt)