    {-|
    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 :: forall (m :: * -> *) s k (tp :: k).
    NonceGenerator m s -> m (Nonce s tp)
    freshNonce (IONG IORef Word64
    r) =
      IORef Word64 -> (Word64 -> (Word64, Nonce s tp)) -> IO (Nonce s tp)
    forall a b. IORef a -> (a -> (a, b)) -> IO b
    atomicModifyIORef' IORef Word64
    r ((Word64 -> (Word64, Nonce s tp)) -> IO (Nonce s tp))
    -> (Word64 -> (Word64, Nonce s tp)) -> IO (Nonce s tp)
    forall a b. (a -> b) -> a -> b
    $ \Word64
    n -> (Word64
    nWord64 -> Word64 -> Word64
    forall a. Num a => a -> a -> a
    +Word64
    1, Word64 -> Nonce s tp
    forall k s (tp :: k). Word64 -> Nonce s tp
    Nonce Word64
    n)
    freshNonce (STNG STRef t Word64
    r) = do
      i <- STRef t Word64 -> ST t Word64
    forall s a. STRef s a -> ST s a
    readSTRef STRef t Word64
    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 :: forall (m :: * -> *) s. NonceGenerator m s -> m Integer
    countNoncesGenerated (IONG IORef Word64
    r) = Word64 -> Integer
    forall a. Integral a => a -> Integer
    toInteger (Word64 -> Integer) -> m Word64 -> m Integer
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> IORef Word64 -> IO Word64
    forall a. IORef a -> IO a
    readIORef IORef Word64
    r
    countNoncesGenerated (STNG STRef t Word64
    r) = Word64 -> Integer
    forall a. Integral a => a -> Integer
    toInteger (Word64 -> Integer) -> m Word64 -> m Integer
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> STRef t Word64 -> ST t Word64
    forall s a. STRef s a -> ST s a
    readSTRef STRef t Word64
    r

    -- | Create a new nonce generator in the 'ST' monad.
    newSTNonceGenerator :: ST t (Some (NonceGenerator (ST t)))
    newSTNonceGenerator :: forall t. ST t (Some (NonceGenerator (ST t)))
    newSTNonceGenerator = NonceGenerator (ST t) (ZonkAny 1) -> Some (NonceGenerator (ST t))
    forall k (f :: k -> *) (x :: k). f x -> Some f
    Some (NonceGenerator (ST t) (ZonkAny 1) -> Some (NonceGenerator (ST t)))
    -> (STRef t Word64 -> NonceGenerator (ST t) (ZonkAny 1))
    -> STRef t Word64
    -> Some (NonceGenerator (ST t))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . STRef t Word64 -> NonceGenerator (ST t) (ZonkAny 1)
    forall t s. STRef t Word64 -> NonceGenerator (ST t) s
    STNG (STRef t Word64 -> Some (NonceGenerator (ST t)))
    -> ST t (STRef t Word64) -> ST t (Some (NonceGenerator (ST t)))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Word64 -> ST t (STRef t Word64)
    forall a s. a -> ST s (STRef s a)
    newSTRef (Int -> Word64
    forall a. Enum a => Int -> a
    toEnum Int
    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 :: forall a. (forall s. NonceGenerator (ST s) s -> ST s a) -> a
    runSTNonceGenerator forall s. NonceGenerator (ST s) s -> ST s a
    f = (forall s. ST s a) -> a
    forall a. (forall s. ST s a) -> a
    runST ((forall s. ST s a) -> a) -> (forall s. ST s a) -> a
    forall a b. (a -> b) -> a -> b
    $ NonceGenerator (ST s) s -> ST s a
    forall s. NonceGenerator (ST s) s -> ST s a
    f (NonceGenerator (ST s) s -> ST s a)
    -> (STRef s Word64 -> NonceGenerator (ST s) s)
    -> STRef s Word64
    -> ST s a
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . STRef s Word64 -> NonceGenerator (ST s) s
    forall t s. STRef t Word64 -> NonceGenerator (ST t) s
    STNG (STRef s Word64 -> ST s a) -> ST s (STRef s Word64) -> ST s a
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< Word64 -> ST s (STRef s Word64)
    forall a s. a -> ST s (STRef s a)
    newSTRef Word64
    0

    -- | Create a new nonce generator in the 'IO' monad.
    newIONonceGenerator :: IO (Some (NonceGenerator IO))
    newIONonceGenerator :: IO (Some (NonceGenerator IO))
    newIONonceGenerator = NonceGenerator IO (ZonkAny 0) -> Some (NonceGenerator IO)
    forall k (f :: k -> *) (x :: k). f x -> Some f
    Some (NonceGenerator IO (ZonkAny 0) -> Some (NonceGenerator IO))
    -> (IORef Word64 -> NonceGenerator IO (ZonkAny 0))
    -> IORef Word64
    -> Some (NonceGenerator IO)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . IORef Word64 -> NonceGenerator IO (ZonkAny 0)
    forall s. IORef Word64 -> NonceGenerator IO s
    IONG (IORef Word64 -> Some (NonceGenerator IO))
    -> IO (IORef Word64) -> IO (Some (NonceGenerator IO))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Word64 -> IO (IORef Word64)
    forall a. a -> IO (IORef a)
    newIORef (Int -> Word64
    forall a. Enum a => Int -> a
    toEnum Int
    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 :: forall t r. (forall s. NonceGenerator (ST t) s -> ST t r) -> ST t r
    withSTNonceGenerator forall s. NonceGenerator (ST t) s -> ST t r
    f = do
      Some r <- ST t (Some (NonceGenerator (ST t)))
    forall t. ST t (Some (NonceGenerator (ST t)))
    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 :: forall r. (forall s. NonceGenerator IO s -> IO r) -> IO r
    withIONonceGenerator forall s. NonceGenerator IO s -> IO r
    f = do
      Some r <- IO (Some (NonceGenerator IO))
    newIONonceGenerator
      f r

    -- | An index generated by the counter.
    newtype Nonce (s :: Type) (tp :: k) = Nonce { forall k s (tp :: k). Nonce s tp -> Word64
    indexValue :: Word64 }
      deriving (Nonce s tp -> Nonce s tp -> Bool
    (Nonce s tp -> Nonce s tp -> Bool)
    -> (Nonce s tp -> Nonce s tp -> Bool) -> Eq (Nonce s tp)
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    forall s k (tp :: k). Nonce s tp -> Nonce s tp -> Bool
    $c== :: forall s k (tp :: k). Nonce s tp -> Nonce s tp -> Bool
    == :: Nonce s tp -> Nonce s tp -> Bool
    $c/= :: forall s k (tp :: k). Nonce s tp -> Nonce s tp -> Bool
    /= :: Nonce s tp -> Nonce s tp -> Bool
    Eq, Eq (Nonce s tp)
    Eq (Nonce s tp) =>
    (Nonce s tp -> Nonce s tp -> Ordering)
    -> (Nonce s tp -> Nonce s tp -> Bool)
    -> (Nonce s tp -> Nonce s tp -> Bool)
    -> (Nonce s tp -> Nonce s tp -> Bool)
    -> (Nonce s tp -> Nonce s tp -> Bool)
    -> (Nonce s tp -> Nonce s tp -> Nonce s tp)
    -> (Nonce s tp -> Nonce s tp -> Nonce s tp)
    -> Ord (Nonce s tp)
    Nonce s tp -> Nonce s tp -> Bool
    Nonce s tp -> Nonce s tp -> Ordering
    Nonce s tp -> Nonce s tp -> Nonce s tp
    forall a.
    Eq a =>
    (a -> a -> Ordering)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> a)
    -> (a -> a -> a)
    -> Ord a
    forall s k (tp :: k). Eq (Nonce s tp)
    forall s k (tp :: k). Nonce s tp -> Nonce s tp -> Bool
    forall s k (tp :: k). Nonce s tp -> Nonce s tp -> Ordering
    forall s k (tp :: k). Nonce s tp -> Nonce s tp -> Nonce s tp
    $ccompare :: forall s k (tp :: k). Nonce s tp -> Nonce s tp -> Ordering
    compare :: Nonce s tp -> Nonce s tp -> Ordering
    $c< :: forall s k (tp :: k). Nonce s tp -> Nonce s tp -> Bool
    < :: Nonce s tp -> Nonce s tp -> Bool
    $c<= :: forall s k (tp :: k). Nonce s tp -> Nonce s tp -> Bool
    <= :: Nonce s tp -> Nonce s tp -> Bool
    $c> :: forall s k (tp :: k). Nonce s tp -> Nonce s tp -> Bool
    > :: Nonce s tp -> Nonce s tp -> Bool
    $c>= :: forall s k (tp :: k). Nonce s tp -> Nonce s tp -> Bool
    >= :: Nonce s tp -> Nonce s tp -> Bool
    $cmax :: forall s k (tp :: k). Nonce s tp -> Nonce s tp -> Nonce s tp
    max :: Nonce s tp -> Nonce s tp -> Nonce s tp
    $cmin :: forall s k (tp :: k). Nonce s tp -> Nonce s tp -> Nonce s tp
    min :: Nonce s tp -> Nonce s tp -> Nonce s tp
    Ord, Eq (Nonce s tp)
    Eq (Nonce s tp) =>
    (Int -> Nonce s tp -> Int)
    -> (Nonce s tp -> Int) -> Hashable (Nonce s tp)
    Int -> Nonce s tp -> Int
    Nonce s tp -> Int
    forall a. Eq a => (Int -> a -> Int) -> (a -> Int) -> Hashable a
    forall s k (tp :: k). Eq (Nonce s tp)
    forall s k (tp :: k). Int -> Nonce s tp -> Int
    forall s k (tp :: k). Nonce s tp -> Int
    $chashWithSalt :: forall s k (tp :: k). Int -> Nonce s tp -> Int
    hashWithSalt :: Int -> Nonce s tp -> Int
    $chash :: forall s k (tp :: k). Nonce s tp -> Int
    hash :: Nonce s tp -> Int
    Hashable, Int -> Nonce s tp -> ShowS
    [Nonce s tp] -> ShowS
    Nonce s tp -> String
    (Int -> Nonce s tp -> ShowS)
    -> (Nonce s tp -> String)
    -> ([Nonce s tp] -> ShowS)
    -> Show (Nonce s tp)
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    forall s k (tp :: k). Int -> Nonce s tp -> ShowS
    forall s k (tp :: k). [Nonce s tp] -> ShowS
    forall s k (tp :: k). Nonce s tp -> String
    $cshowsPrec :: forall s k (tp :: k). Int -> Nonce s tp -> ShowS
    showsPrec :: Int -> Nonce s tp -> ShowS
    $cshow :: forall s k (tp :: k). Nonce s tp -> String
    show :: Nonce s tp -> String
    $cshowList :: forall s k (tp :: k). [Nonce s tp] -> ShowS
    showList :: [Nonce s tp] -> ShowS
    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 :: forall (a :: k). Nonce s a -> Nonce s a -> Bool
    eqF Nonce s a
    x Nonce s a
    y = Maybe (a :~: a) -> Bool
    forall a. Maybe a -> Bool
    isJust (Nonce s a -> Nonce s a -> Maybe (a :~: a)
    forall (a :: k) (b :: k). Nonce s a -> Nonce s b -> Maybe (a :~: b)
    forall {k} (f :: k -> *) (a :: k) (b :: k).
    TestEquality f =>
    f a -> f b -> Maybe (a :~: b)
    testEquality Nonce s a
    x Nonce s a
    y)

    instance TestEquality (Nonce s) where
      testEquality :: forall (a :: k) (b :: k). Nonce s a -> Nonce s b -> Maybe (a :~: b)
    testEquality Nonce s a
    x Nonce s b
    y | Nonce s a -> Word64
    forall k s (tp :: k). Nonce s tp -> Word64
    indexValue Nonce s a
    x Word64 -> Word64 -> Bool
    forall a. Eq a => a -> a -> Bool
    == Nonce s b -> Word64
    forall k s (tp :: k). Nonce s tp -> Word64
    indexValue Nonce s b
    y = (a :~: b) -> Maybe (a :~: b)
    forall a. a -> Maybe a
    Just a :~: b
    forall {k} (a :: k) (b :: k). a :~: b
    unsafeAxiom
                       | Bool
    otherwise = Maybe (a :~: b)
    forall a. Maybe a
    Nothing

    instance OrdF (Nonce s) where
      compareF :: forall (x :: k) (y :: k). Nonce s x -> Nonce s y -> OrderingF x y
    compareF Nonce s x
    x Nonce s y
    y =
        case Word64 -> Word64 -> Ordering
    forall a. Ord a => a -> a -> Ordering
    compare (Nonce s x -> Word64
    forall k s (tp :: k). Nonce s tp -> Word64
    indexValue Nonce s x
    x) (Nonce s y -> Word64
    forall k s (tp :: k). Nonce s tp -> Word64
    indexValue Nonce s y
    y) of
          Ordering
    LT -> OrderingF x y
    forall {k} (x :: k) (y :: k). OrderingF x y
    LTF
          Ordering
    EQ -> OrderingF (ZonkAny 3) (ZonkAny 3) -> OrderingF x y
    forall a b. a -> b
    unsafeCoerce OrderingF (ZonkAny 3) (ZonkAny 3)
    forall {k} (x :: k). OrderingF x x
    EQF
          Ordering
    GT -> OrderingF x y
    forall {k} (x :: k) (y :: k). OrderingF x y
    GTF

    instance HashableF (Nonce s) where
      hashWithSaltF :: forall (tp :: k). Int -> Nonce s tp -> Int
    hashWithSaltF Int
    s (Nonce Word64
    x) = Int -> Word64 -> Int
    forall a. Hashable a => Int -> a -> Int
    hashWithSalt Int
    s Word64
    x

    instance ShowF (Nonce s)

    ------------------------------------------------------------------------
    -- * GlobalNonceGenerator

    data GlobalNonceGenerator

    globalNonceIORef :: IORef Word64
    globalNonceIORef :: IORef Word64
    globalNonceIORef = IO (IORef Word64) -> IORef Word64
    forall a. IO a -> a
    unsafePerformIO (Word64 -> IO (IORef Word64)
    forall a. a -> IO (IORef a)
    newIORef Word64
    0)
    {-# NOINLINE globalNonceIORef #-}

    -- | A nonce generator that uses a globally-defined counter.
    globalNonceGenerator :: NonceGenerator IO GlobalNonceGenerator
    globalNonceGenerator :: NonceGenerator IO GlobalNonceGenerator
    globalNonceGenerator = IORef Word64 -> NonceGenerator IO GlobalNonceGenerator
    forall s. IORef Word64 -> NonceGenerator IO s
    IONG IORef Word64
    globalNonceIORef

    -- | Create a new counter.
    withGlobalSTNonceGenerator :: (forall t . NonceGenerator (ST t) t -> ST t r) -> r
    withGlobalSTNonceGenerator :: forall a. (forall s. NonceGenerator (ST s) s -> ST s a) -> a
    withGlobalSTNonceGenerator forall t. NonceGenerator (ST t) t -> ST t r
    f = (forall s. ST s r) -> r
    forall a. (forall s. ST s a) -> a
    runST ((forall s. ST s r) -> r) -> (forall s. ST s r) -> r
    forall a b. (a -> b) -> a -> b
    $ do
      r <- Word64 -> ST s (STRef s Word64)
    forall a s. a -> ST s (STRef s a)
    newSTRef (Int -> Word64
    forall a. Enum a => Int -> a
    toEnum Int
    0)
      f $! STNG r