------------------------------------------------------------------------
-- |
-- Module           : Data.Parameterized.NonceGenerator
-- Description      : A counter in the ST monad.
-- Copyright        : (c) Galois, Inc 2014
-- Maintainer       : Joe Hendrix <jhendrix@galois.com>
-- Stability        : provisional
--
-- This module provides a simple generator of new indexes in the ST monad.
-- It is predictable and not intended for cryptographic purposes.
--
-- 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.  This is only OK because of the discipline
-- by which nonces should be used: they should only be generated from
-- a 'NonceGenerator' (i.e., should not be built directly), and nonces from
-- different generators must never be compared!  Arranging to compare
-- Nonces from different origins would allow users to build 'unsafeCoerce'
-- via the 'testEquality' function.
--
-- A somewhat safer API would be to brand the generated Nonces with the
-- state type variable of the NonceGenerator whence they came, and to only
-- provide NonceGenerators via a Rank-2 continuation-passing API, similar to
-- 'runST'. This would (via a meta-argument involving parametricity)
-- help to prevent nonces of different origin from being compared.
-- However, this would force us to push the 'ST' type brand into a significant
-- number of other structures and APIs.
--
-- Another alternative would be to use 'unsafePerformIO' magic to make
-- a global nonce generator, and make that the only way to generate nonces.
-- It is not clear that this is actually an improvement from a type safety
-- point of view, but an argument could be made.
--
-- For now, be careful using Nonces, and ensure that you do not mix
-- Nonces from different NonceGenerators.
------------------------------------------------------------------------
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE Unsafe #-}
module Data.Parameterized.Nonce.Unsafe
  ( NonceGenerator
  , newNonceGenerator
  , freshNonce
  , atLimit
  , Nonce
  , indexValue
  ) where

import Control.Monad.ST
import Data.Hashable
import Data.STRef
import Data.Word
import Unsafe.Coerce

import Data.Parameterized.Classes

-- | A simple type that for getting fresh indices in the 'ST' monad.
-- The type parameter @s@ is used for the 'ST' monad parameter.
newtype NonceGenerator s = NonceGenerator (STRef s Word64)

-- | Create a new counter.
newNonceGenerator :: ST s (NonceGenerator s)
newNonceGenerator = NonceGenerator `fmap` newSTRef (toEnum 0)

-- | An index generated by the counter.
newtype Nonce (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

instance TestEquality Nonce where
  testEquality x y | indexValue x == indexValue y = unsafeCoerce (Just Refl)
                   | otherwise = Nothing

instance OrdF Nonce where
  compareF x y =
    case compare (indexValue x) (indexValue y) of
      LT -> LTF
      EQ -> unsafeCoerce EQF
      GT -> GTF

instance HashableF Nonce where
  hashWithSaltF s (Nonce x) = hashWithSalt s x

instance ShowF Nonce

{-# INLINE freshNonce #-}
-- | Get a fresh index and increment the counter.
freshNonce :: NonceGenerator s -> ST s (Nonce tp)
freshNonce (NonceGenerator r) = do
  i <- readSTRef r
  writeSTRef r $! succ i
  return (Nonce i)

-- | Return true if counter has reached the limit, and can't be
-- incremented without risk of error.
atLimit :: NonceGenerator s -> ST s Bool
atLimit (NonceGenerator r) = do
  i <- readSTRef r
  return (i == maxBound)