-- |
-- 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, ScrubbedBytes)
import qualified Crypto.Internal.ByteArray as B
import           Foreign.Storable (pokeElemOff)

import qualified Crypto.Cipher.ChaCha as C

instance DRG ChaChaDRG where
    randomBytesGenerate = generate

-- | ChaCha Deterministic Random Generator
newtype ChaChaDRG = ChaChaDRG C.StateSimple
    deriving (NFData)

-- | Initialize a new ChaCha context with the number of rounds,
-- the key and the nonce associated.
initialize :: ByteArray seed
           => seed        -- ^ 40 bytes of seed
           -> ChaChaDRG   -- ^ the initial ChaCha state
initialize seed = ChaChaDRG $ C.initializeSimple 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 (a,b,c,d,e) = initialize (B.allocAndFreeze 40 fill :: ScrubbedBytes)
  where fill s = mapM_ (uncurry (pokeElemOff s)) [(0,a), (1,b), (2,c), (3,d), (4,e)]

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