{-|
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 ExistentialQuantification #-}
{-# LANGUAGE EmptyDataDecls #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE TypeInType #-}
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) =
  forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Word64
r forall a b. (a -> b) -> a -> b
$ \Word64
n -> (Word64
nforall a. Num a => a -> a -> a
+Word64
1, forall k s (tp :: k). Word64 -> Nonce s tp
Nonce Word64
n)
freshNonce (STNG STRef t Word64
r) = do
  Word64
i <- forall s a. STRef s a -> ST s a
readSTRef STRef t Word64
r
  forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef t Word64
r forall a b. (a -> b) -> a -> b
$! Word64
iforall a. Num a => a -> a -> a
+Word64
1
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall k s (tp :: k). Word64 -> Nonce s tp
Nonce Word64
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) = forall a. Integral a => a -> Integer
toInteger forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. IORef a -> IO a
readIORef IORef Word64
r
countNoncesGenerated (STNG STRef t Word64
r) = forall a. Integral a => a -> Integer
toInteger forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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 = forall k (f :: k -> *) (x :: k). f x -> Some f
Some forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall t s. STRef t Word64 -> NonceGenerator (ST t) s
STNG forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a s. a -> ST s (STRef s a)
newSTRef (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 a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ forall s. NonceGenerator (ST s) s -> ST s a
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall t s. STRef t Word64 -> NonceGenerator (ST t) s
STNG forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< 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 = forall k (f :: k -> *) (x :: k). f x -> Some f
Some forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s. IORef Word64 -> NonceGenerator IO s
IONG forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. a -> IO (IORef a)
newIORef (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 NonceGenerator (ST t) x
r <- forall t. ST t (Some (NonceGenerator (ST t)))
newSTNonceGenerator
  forall s. NonceGenerator (ST t) s -> ST t r
f NonceGenerator (ST t) x
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 NonceGenerator IO x
r <- IO (Some (NonceGenerator IO))
newIONonceGenerator
  forall s. NonceGenerator IO s -> IO r
f NonceGenerator IO x
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
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
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
Eq, 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
min :: Nonce s tp -> Nonce s tp -> Nonce s tp
$cmin :: forall s k (tp :: k). Nonce s tp -> Nonce s tp -> Nonce s tp
max :: Nonce s tp -> Nonce s tp -> Nonce s tp
$cmax :: forall s k (tp :: k). Nonce s tp -> Nonce s tp -> Nonce s tp
>= :: 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
$c< :: forall s k (tp :: k). Nonce s tp -> Nonce s tp -> Bool
compare :: Nonce s tp -> Nonce s tp -> Ordering
$ccompare :: forall s k (tp :: k). Nonce s tp -> Nonce s tp -> Ordering
Ord, 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
hash :: Nonce s tp -> Int
$chash :: forall s k (tp :: k). Nonce s tp -> Int
hashWithSalt :: Int -> Nonce s tp -> Int
$chashWithSalt :: forall s k (tp :: k). Int -> Nonce s tp -> Int
Hashable, Int -> Nonce s tp -> ShowS
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
showList :: [Nonce s tp] -> ShowS
$cshowList :: forall s k (tp :: k). [Nonce s tp] -> ShowS
show :: Nonce s tp -> String
$cshow :: forall s k (tp :: k). Nonce s tp -> String
showsPrec :: Int -> Nonce s tp -> ShowS
$cshowsPrec :: forall s k (tp :: k). Int -> 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 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 | forall k s (tp :: k). Nonce s tp -> Word64
indexValue Nonce s a
x forall a. Eq a => a -> a -> Bool
== forall k s (tp :: k). Nonce s tp -> Word64
indexValue Nonce s b
y = forall a. a -> Maybe a
Just forall {k} (a :: k) (b :: k). a :~: b
unsafeAxiom
                   | Bool
otherwise = 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 forall a. Ord a => a -> a -> Ordering
compare (forall k s (tp :: k). Nonce s tp -> Word64
indexValue Nonce s x
x) (forall k s (tp :: k). Nonce s tp -> Word64
indexValue Nonce s y
y) of
      Ordering
LT -> forall {k} (x :: k) (y :: k). OrderingF x y
LTF
      Ordering
EQ -> forall a b. a -> b
unsafeCoerce forall {k} (x :: k). OrderingF x x
EQF
      Ordering
GT -> 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) = 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 = forall a. IO a -> a
unsafePerformIO (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 = 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 a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
  STRef s Word64
r <- forall a s. a -> ST s (STRef s a)
newSTRef (forall a. Enum a => Int -> a
toEnum Int
0)
  forall t. NonceGenerator (ST t) t -> ST t r
f forall a b. (a -> b) -> a -> b
$! forall t s. STRef t Word64 -> NonceGenerator (ST t) s
STNG STRef s Word64
r