content: {-|
  Description : Index generator in the ST monad.
  Copyright   : (c) Galois, Inc 2014-2019
  Maintainer  : Joe Hendrix <jhendrix@galois.com>
  
  This module provides a simple generator of new indexes in the 'ST' monad.
  It is predictable and not intended for cryptographic purposes.
  
  This module also provides a global nonce generator that will generate
  2^64 nonces before repeating.
  
  NOTE: The 'TestEquality' and 'OrdF' instances for the 'Nonce' type simply
  compare the generated nonce values and then assert to the compiler
  (via 'unsafeCoerce') that the types ascribed to the nonces are equal
  if their values are equal.
  -}
  {-# LANGUAGE CPP #-}
  {-# LANGUAGE ExistentialQuantification #-}
  {-# LANGUAGE EmptyDataDecls #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE KindSignatures #-}
  {-# LANGUAGE PolyKinds #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE RoleAnnotations #-}
  {-# LANGUAGE Trustworthy #-}
  #if __GLASGOW_HASKELL__ < 806
  {-# LANGUAGE TypeInType #-}
  #endif
  module Data.Parameterized.Nonce
    ( -- * NonceGenerator
      NonceGenerator
    , freshNonce
    , countNoncesGenerated
    , Nonce
    , indexValue
      -- * Accessing a nonce generator
    , newSTNonceGenerator
    , newIONonceGenerator
    , withIONonceGenerator
    , withSTNonceGenerator
    , runSTNonceGenerator
      -- * Global nonce generator
    , withGlobalSTNonceGenerator
    , GlobalNonceGenerator
    , globalNonceGenerator
    ) where
  
  import Control.Monad.ST
  import Data.Hashable
  import Data.Kind
  import Data.IORef
  import Data.STRef
  import Data.Word
  import Unsafe.Coerce
  import System.IO.Unsafe (unsafePerformIO)
  
  import Data.Parameterized.Axiom
  import Data.Parameterized.Classes
  import Data.Parameterized.Some
  
  -- | Provides a monadic action for getting fresh typed names.
  --
  -- The first type parameter @m@ is the monad used for generating names, and
  -- the second parameter @s@ is used for the counter.
  data NonceGenerator (m :: Type -> Type) (s :: Type) where
    STNG :: !(STRef t Word64) -> NonceGenerator (ST t) s
    IONG :: !(IORef Word64) -> NonceGenerator IO s
  
  freshNonce :: forall m s k (tp :: k) . NonceGenerator m s -> m (Nonce s tp)
  freshNonce (IONG r) =
    atomicModifyIORef' r $ \n -> (n+1, Nonce n)
  freshNonce (STNG r) = do
    i <- readSTRef r
    writeSTRef r $! i+1
    return $ Nonce i
    -- (Weirdly, there's no atomicModifySTRef'.  Yes, only the IO monad
    -- does concurrency, but the ST monad is part of the IO monad via
    -- stToIO, so there's no guarantee that ST code won't be run in
    -- multiple threads.)
  
  {-# INLINE freshNonce #-}
    -- Inlining is particularly necessary since there's no @Monad m@
    -- constraint on 'freshNonce', so SPECIALIZE doesn't work on it.  In
    -- this case, though, we get specialization for free from inlining.
    -- For instance, a @NonceGenerator IO s@ must be an @IONG@, so the
    -- simplifier eliminates the STNG branch.
  
  -- | The number of nonces generated so far by this generator.  Only
  -- really useful for profiling.
  countNoncesGenerated :: NonceGenerator m s -> m Integer
  countNoncesGenerated (IONG r) = toInteger <$> readIORef r
  countNoncesGenerated (STNG r) = toInteger <$> readSTRef r
  
  -- | Create a new nonce generator in the 'ST' monad.
  newSTNonceGenerator :: ST t (Some (NonceGenerator (ST t)))
  newSTNonceGenerator = Some . STNG <$> newSTRef (toEnum 0)
  
  -- | This combines `runST` and `newSTNonceGenerator` to create a nonce
  -- generator that shares the same phantom type parameter as the @ST@ monad.
  --
  -- This can be used to reduce the number of type parameters when we know a
  -- ST computation only needs a single `NonceGenerator`.
  runSTNonceGenerator :: (forall s . NonceGenerator (ST s) s -> ST s a)
                      -> a
  runSTNonceGenerator f = runST $ f . STNG =<< newSTRef 0
  
  -- | Create a new nonce generator in the 'IO' monad.
  newIONonceGenerator :: IO (Some (NonceGenerator IO))
  newIONonceGenerator = Some . IONG <$> newIORef (toEnum 0)
  
  -- | Run an 'ST' computation with a new nonce generator in the 'ST' monad.
  withSTNonceGenerator :: (forall s . NonceGenerator (ST t) s -> ST t r) -> ST t r
  withSTNonceGenerator f = do
    Some r <- newSTNonceGenerator
    f r
  
  -- | Run an 'IO' computation with a new nonce generator in the 'IO' monad.
  withIONonceGenerator :: (forall s . NonceGenerator IO s -> IO r) -> IO r
  withIONonceGenerator f = do
    Some r <- newIONonceGenerator
    f r
  
  -- | An index generated by the counter.
  newtype Nonce (s :: Type) (tp :: k) = Nonce { indexValue :: Word64 }
    deriving (Eq, Ord, Hashable, Show)
  
  --  Force the type role of Nonce to be nominal: this prevents Data.Coerce.coerce
  --  from casting the types of nonces, which it would otherwise be able to do
  --  because tp is a phantom type parameter.  This partially helps to protect
  --  the nonce abstraction.
  type role Nonce nominal nominal
  
  instance EqF (Nonce s) where
    eqF x y = isJust (testEquality x y)
  
  instance TestEquality (Nonce s) where
    testEquality x y | indexValue x == indexValue y = Just unsafeAxiom
                     | otherwise = Nothing
  
  instance OrdF (Nonce s) where
    compareF x y =
      case compare (indexValue x) (indexValue y) of
        LT -> LTF
        EQ -> unsafeCoerce EQF
        GT -> GTF
  
  instance HashableF (Nonce s) where
    hashWithSaltF s (Nonce x) = hashWithSalt s x
  
  instance ShowF (Nonce s)
  
  ------------------------------------------------------------------------
  -- * GlobalNonceGenerator
  
  data GlobalNonceGenerator
  
  globalNonceIORef :: IORef Word64
  globalNonceIORef = unsafePerformIO (newIORef 0)
  {-# NOINLINE globalNonceIORef #-}
  
  -- | A nonce generator that uses a globally-defined counter.
  globalNonceGenerator :: NonceGenerator IO GlobalNonceGenerator
  globalNonceGenerator = IONG globalNonceIORef
  
  -- | Create a new counter.
  withGlobalSTNonceGenerator :: (forall t . NonceGenerator (ST t) t -> ST t r) -> r
  withGlobalSTNonceGenerator f = runST $ do
    r <- newSTRef (toEnum 0)
    f $! STNG r
  