{-|
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 :: 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
  Word64
i <- STRef t Word64 -> ST t Word64
forall s a. STRef s a -> ST s a
readSTRef STRef t Word64
r
  STRef t Word64 -> Word64 -> ST t ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef t Word64
r (Word64 -> ST t ()) -> Word64 -> ST t ()
forall a b. (a -> b) -> a -> b
$! Word64
iWord64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+Word64
1
  Nonce s tp -> m (Nonce s tp)
forall (m :: * -> *) a. Monad m => a -> m a
return (Nonce s tp -> m (Nonce s tp)) -> Nonce s tp -> m (Nonce s tp)
forall a b. (a -> b) -> a -> b
$ Word64 -> Nonce s tp
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 :: NonceGenerator m s -> m Integer
countNoncesGenerated (IONG IORef Word64
r) = Word64 -> Integer
forall a. Integral a => a -> Integer
toInteger (Word64 -> Integer) -> IO Word64 -> IO 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) -> ST t Word64 -> ST t 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 :: ST t (Some (NonceGenerator (ST t)))
newSTNonceGenerator = NonceGenerator (ST t) Any -> Some (NonceGenerator (ST t))
forall k (f :: k -> *) (x :: k). f x -> Some f
Some (NonceGenerator (ST t) Any -> Some (NonceGenerator (ST t)))
-> (STRef t Word64 -> NonceGenerator (ST t) Any)
-> STRef t Word64
-> Some (NonceGenerator (ST t))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STRef t Word64 -> NonceGenerator (ST t) Any
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 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 Any -> Some (NonceGenerator IO)
forall k (f :: k -> *) (x :: k). f x -> Some f
Some (NonceGenerator IO Any -> Some (NonceGenerator IO))
-> (IORef Word64 -> NonceGenerator IO Any)
-> IORef Word64
-> Some (NonceGenerator IO)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IORef Word64 -> NonceGenerator IO Any
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 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 <- ST t (Some (NonceGenerator (ST t)))
forall t. ST t (Some (NonceGenerator (ST t)))
newSTNonceGenerator
  NonceGenerator (ST t) x -> ST t r
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 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
  NonceGenerator IO x -> IO r
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 { 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
/= :: 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, 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
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
$cp1Ord :: forall s k (tp :: k). Eq (Nonce s tp)
Ord, Int -> Nonce s tp -> Int
Nonce s tp -> Int
(Int -> Nonce s tp -> Int)
-> (Nonce s tp -> Int) -> Hashable (Nonce s tp)
forall a. (Int -> a -> Int) -> (a -> Int) -> Hashable a
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
[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
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 :: Nonce s a -> Nonce s b -> Maybe (a :~: b)
testEquality Nonce s a
x Nonce s b
y | Nonce s a -> Word64
forall s k (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 s k (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 :: 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 s k (tp :: k). Nonce s tp -> Word64
indexValue Nonce s x
x) (Nonce s y -> Word64
forall s k (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 Any Any -> OrderingF x y
forall a b. a -> b
unsafeCoerce OrderingF Any Any
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 :: 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 t. NonceGenerator (ST t) t -> ST t r) -> r
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
  STRef s Word64
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)
  NonceGenerator (ST s) s -> ST s r
forall t. NonceGenerator (ST t) t -> ST t r
f (NonceGenerator (ST s) s -> ST s r)
-> NonceGenerator (ST s) s -> ST s r
forall a b. (a -> b) -> a -> b
$! STRef s Word64 -> NonceGenerator (ST s) s
forall t s. STRef t Word64 -> NonceGenerator (ST t) s
STNG STRef s Word64
r