{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Test.Cardano.Ledger.Shelley.Utils (
  mkSeedFromWords,
  mkCertifiedVRF,
  epochFromSlotNo,
  evolveKESUntil,
  slotFromEpoch,
  epochSize,
  mkHash,
  mkKeyPair,
  mkKeyPair',
  mkGenKey,
  mkKESKeyPair,
  mkVRFKeyPair,
  runShelleyBase,
  maxKESIterations,
  slotsPerKESIteration,
  testSTS,
  maxLLSupply,
  applySTSTest,
  GenesisKeyPair,
  getBlockNonce,
  ChainProperty,
  RawSeed (..),
  Split (..),
  module CoreUtils,
)
where

import Cardano.Crypto.DSIGN.Class (DSIGNAlgorithm (..))
import Cardano.Crypto.Hash (
  Blake2b_256,
  Hash,
  HashAlgorithm,
  hashToBytes,
 )
import Cardano.Crypto.KES (
  KESAlgorithm (..),
  deriveVerKeyKES,
  genKeyKES,
 )
import Cardano.Crypto.Seed (Seed, mkSeedFromBytes)
import Cardano.Crypto.VRF (
  CertifiedVRF,
  SignKeyVRF,
  VRFAlgorithm (..),
  certifiedOutput,
  deriveVerKeyVRF,
  evalCertified,
  genKeyVRF,
 )
import qualified Cardano.Crypto.VRF as VRF
import Cardano.Ledger.BaseTypes (
  Globals (..),
  Nonce,
  ShelleyBase,
  epochInfoPure,
  mkNonceFromOutputVRF,
 )
import Cardano.Ledger.Binary (EncCBOR (..), hashWithEncoder, shelleyProtVer)
import Cardano.Ledger.Block (Block, bheader)
import Cardano.Ledger.Coin (Coin (..))
import Cardano.Ledger.Crypto (Crypto (DSIGN))
import Cardano.Ledger.Shelley.API (ApplyBlock, KeyRole (..), VKey (..))
import Cardano.Ledger.Shelley.Core
import Cardano.Ledger.Slot (EpochNo, EpochSize (..), SlotNo)
import Cardano.Protocol.TPraos.API (GetLedgerView)
import Cardano.Protocol.TPraos.BHeader (BHBody (..), BHeader, bhbody)
import Cardano.Slotting.EpochInfo (
  epochInfoEpoch,
  epochInfoFirst,
  epochInfoSize,
 )
import Control.Monad.Reader.Class (asks)
import Control.Monad.Trans.Reader (runReaderT)
import Control.State.Transition.Extended hiding (Assertion)
import Data.Coerce (Coercible, coerce)
import Data.Functor.Identity (runIdentity)
import Data.List.NonEmpty (NonEmpty)
import Data.Typeable (Proxy (Proxy))
import Data.Word (Word64)
import Test.Cardano.Ledger.Core.KeyPair (KeyPair, pattern KeyPair)
import Test.Cardano.Ledger.Core.Utils as CoreUtils
import Test.Cardano.Ledger.Shelley.Arbitrary (RawSeed (..))
import Test.Cardano.Ledger.Shelley.ConcreteCryptoTypes (Mock)
import Test.Cardano.Ledger.TreeDiff (ToExpr)
import Test.Cardano.Protocol.TPraos.Create (KESKeyPair (..), VRFKeyPair (..), evolveKESUntil)
import Test.Control.State.Transition.Trace (
  applySTSTest,
  checkTrace,
  (.-),
  (.->>),
 )
import Test.Tasty.HUnit (
  Assertion,
  (@?=),
 )

type ChainProperty era =
  ( Mock (EraCrypto era)
  , ApplyBlock era
  , GetLedgerView era
  , EraTx era
  )

-- ================================================

class Split v where
  vsplit :: v -> Integer -> ([v], Coin)

-- ===============================================================================
-- Generating random transactions requires splitting Values into multiple Values
-- with the same underlying amount of Coin. This property is crucial to generating
-- transactions which have the preservation of ADA property. (vsplit n v) breaks
-- v into n different values, and one remainder Coin, where the sum of the Coin
-- in the original value, and the sum of the underlying Coin in the list plus the
-- remainder coin are equal.
-- Given:    let (vs,coin) = split n value
-- Then:     (coin value) == sum(map coin vs) <+> coin

-- We introduce a new class Split which supplies this operation.
-- As new kinds of values become instances of the Val class, and we want to generate
-- transactions over these values, we will have to add additional instances here.

instance Split Coin where
  vsplit :: Coin -> Integer -> ([Coin], Coin)
vsplit (Coin Integer
n) Integer
0 = ([], Integer -> Coin
Coin Integer
n)
  vsplit (Coin Integer
n) Integer
m
    | Integer
m forall a. Ord a => a -> a -> Bool
<= Integer
0 = forall a. HasCallStack => [Char] -> a
error [Char]
"must split coins into positive parts"
    | Bool
otherwise = (forall a. Int -> [a] -> [a]
take (forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
m) (forall a. a -> [a]
repeat (Integer -> Coin
Coin (Integer
n forall a. Integral a => a -> a -> a
`div` Integer
m))), Integer -> Coin
Coin (Integer
n forall a. Integral a => a -> a -> a
`rem` Integer
m))

type GenesisKeyPair c = KeyPair 'Genesis c

instance EncCBOR RawSeed where
  encCBOR :: RawSeed -> Encoding
encCBOR (RawSeed Word64
w1 Word64
w2 Word64
w3 Word64
w4 Word64
w5) = forall a. EncCBOR a => a -> Encoding
encCBOR (Word64
w1, Word64
w2, Word64
w3, Word64
w4, Word64
w5)
  encodedSizeExpr :: (forall t. EncCBOR t => Proxy t -> Size) -> Proxy RawSeed -> Size
encodedSizeExpr forall t. EncCBOR t => Proxy t -> Size
size Proxy RawSeed
_ = Size
1 forall a. Num a => a -> a -> a
+ forall t. EncCBOR t => Proxy t -> Size
size (forall {k} (t :: k). Proxy t
Proxy :: Proxy Word64) forall a. Num a => a -> a -> a
* Size
5

-- | Construct a seed from a bunch of Word64s
--
--   We multiply these words by some extra stuff to make sure they contain
--   enough bits for our seed.
mkSeedFromWords ::
  RawSeed ->
  Seed
mkSeedFromWords :: RawSeed -> Seed
mkSeedFromWords RawSeed
stuff =
  ByteString -> Seed
mkSeedFromBytes forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall h a. Hash h a -> ByteString
hashToBytes forall a b. (a -> b) -> a -> b
$ forall h a.
HashAlgorithm h =>
Version -> (a -> Encoding) -> a -> Hash h a
hashWithEncoder @Blake2b_256 Version
shelleyProtVer forall a. EncCBOR a => a -> Encoding
encCBOR RawSeed
stuff

-- | For testing purposes, generate a deterministic genesis key pair given a seed.
mkGenKey ::
  DSIGNAlgorithm (DSIGN c) =>
  RawSeed ->
  (SignKeyDSIGN (DSIGN c), VKey kd c)
mkGenKey :: forall c (kd :: KeyRole).
DSIGNAlgorithm (DSIGN c) =>
RawSeed -> (SignKeyDSIGN (DSIGN c), VKey kd c)
mkGenKey RawSeed
seed =
  let sk :: SignKeyDSIGN (DSIGN c)
sk = forall v. DSIGNAlgorithm v => Seed -> SignKeyDSIGN v
genKeyDSIGN forall a b. (a -> b) -> a -> b
$ RawSeed -> Seed
mkSeedFromWords RawSeed
seed
   in (SignKeyDSIGN (DSIGN c)
sk, forall (kd :: KeyRole) c. VerKeyDSIGN (DSIGN c) -> VKey kd c
VKey forall a b. (a -> b) -> a -> b
$ forall v. DSIGNAlgorithm v => SignKeyDSIGN v -> VerKeyDSIGN v
deriveVerKeyDSIGN SignKeyDSIGN (DSIGN c)
sk)

-- | For testing purposes, generate a deterministic key pair given a seed.
mkKeyPair ::
  forall c kd.
  DSIGNAlgorithm (DSIGN c) =>
  RawSeed ->
  (SignKeyDSIGN (DSIGN c), VKey kd c)
mkKeyPair :: forall c (kd :: KeyRole).
DSIGNAlgorithm (DSIGN c) =>
RawSeed -> (SignKeyDSIGN (DSIGN c), VKey kd c)
mkKeyPair RawSeed
seed =
  let sk :: SignKeyDSIGN (DSIGN c)
sk = forall v. DSIGNAlgorithm v => Seed -> SignKeyDSIGN v
genKeyDSIGN forall a b. (a -> b) -> a -> b
$ RawSeed -> Seed
mkSeedFromWords RawSeed
seed
   in (SignKeyDSIGN (DSIGN c)
sk, forall (kd :: KeyRole) c. VerKeyDSIGN (DSIGN c) -> VKey kd c
VKey forall a b. (a -> b) -> a -> b
$ forall v. DSIGNAlgorithm v => SignKeyDSIGN v -> VerKeyDSIGN v
deriveVerKeyDSIGN SignKeyDSIGN (DSIGN c)
sk)

-- | For testing purposes, generate a deterministic key pair given a seed.
mkKeyPair' ::
  DSIGNAlgorithm (DSIGN c) =>
  RawSeed ->
  KeyPair kd c
mkKeyPair' :: forall c (kd :: KeyRole).
DSIGNAlgorithm (DSIGN c) =>
RawSeed -> KeyPair kd c
mkKeyPair' RawSeed
seed = forall (kd :: KeyRole) c.
VKey kd c -> SignKeyDSIGN (DSIGN c) -> KeyPair kd c
KeyPair VKey kd c
vk SignKeyDSIGN (DSIGN c)
sk
  where
    (SignKeyDSIGN (DSIGN c)
sk, VKey kd c
vk) = forall c (kd :: KeyRole).
DSIGNAlgorithm (DSIGN c) =>
RawSeed -> (SignKeyDSIGN (DSIGN c), VKey kd c)
mkKeyPair RawSeed
seed

-- | For testing purposes, generate a deterministic VRF key pair given a seed.
mkVRFKeyPair :: Crypto c => RawSeed -> VRFKeyPair c
mkVRFKeyPair :: forall c. Crypto c => RawSeed -> VRFKeyPair c
mkVRFKeyPair RawSeed
seed =
  let sk :: SignKeyVRF (VRF c)
sk = forall v. VRFAlgorithm v => Seed -> SignKeyVRF v
genKeyVRF forall a b. (a -> b) -> a -> b
$ RawSeed -> Seed
mkSeedFromWords RawSeed
seed
   in VRFKeyPair
        { vrfSignKey :: SignKeyVRF (VRF c)
vrfSignKey = SignKeyVRF (VRF c)
sk
        , vrfVerKey :: VerKeyVRF c
vrfVerKey = forall v. VRFAlgorithm v => SignKeyVRF v -> VerKeyVRF v
deriveVerKeyVRF SignKeyVRF (VRF c)
sk
        }

-- | For testing purposes, create a VRF value
mkCertifiedVRF ::
  ( VRF.Signable v a
  , VRFAlgorithm v
  , ContextVRF v ~ ()
  , Coercible b (CertifiedVRF v a)
  ) =>
  a ->
  SignKeyVRF v ->
  b
mkCertifiedVRF :: forall v a b.
(Signable v a, VRFAlgorithm v, ContextVRF v ~ (),
 Coercible b (CertifiedVRF v a)) =>
a -> SignKeyVRF v -> b
mkCertifiedVRF a
a SignKeyVRF v
sk =
  coerce :: forall a b. Coercible a b => a -> b
coerce forall a b. (a -> b) -> a -> b
$ forall v a.
(VRFAlgorithm v, Signable v a) =>
ContextVRF v -> a -> SignKeyVRF v -> CertifiedVRF v a
evalCertified () a
a SignKeyVRF v
sk

-- | For testing purposes, generate a deterministic KES key pair given a seed.
mkKESKeyPair :: Crypto c => RawSeed -> KESKeyPair c
mkKESKeyPair :: forall c. Crypto c => RawSeed -> KESKeyPair c
mkKESKeyPair RawSeed
seed =
  let sk :: SignKeyKES (KES c)
sk = forall v. KESAlgorithm v => Seed -> SignKeyKES v
genKeyKES forall a b. (a -> b) -> a -> b
$ RawSeed -> Seed
mkSeedFromWords RawSeed
seed
   in KESKeyPair
        { kesSignKey :: SignKeyKES (KES c)
kesSignKey = SignKeyKES (KES c)
sk
        , kesVerKey :: VerKeyKES c
kesVerKey = forall v. KESAlgorithm v => SignKeyKES v -> VerKeyKES v
deriveVerKeyKES SignKeyKES (KES c)
sk
        }

runShelleyBase :: ShelleyBase a -> a
runShelleyBase :: forall a. ShelleyBase a -> a
runShelleyBase ShelleyBase a
act = forall a. Identity a -> a
runIdentity forall a b. (a -> b) -> a -> b
$ forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ShelleyBase a
act Globals
testGlobals

epochFromSlotNo :: SlotNo -> EpochNo
epochFromSlotNo :: SlotNo -> EpochNo
epochFromSlotNo = forall a. Identity a -> a
runIdentity forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *).
HasCallStack =>
EpochInfo m -> SlotNo -> m EpochNo
epochInfoEpoch (Globals -> EpochInfo Identity
epochInfoPure Globals
testGlobals)

slotFromEpoch :: EpochNo -> SlotNo
slotFromEpoch :: EpochNo -> SlotNo
slotFromEpoch = forall a. Identity a -> a
runIdentity forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *).
HasCallStack =>
EpochInfo m -> EpochNo -> m SlotNo
epochInfoFirst (Globals -> EpochInfo Identity
epochInfoPure Globals
testGlobals)

epochSize :: EpochNo -> EpochSize
epochSize :: EpochNo -> EpochSize
epochSize = forall a. Identity a -> a
runIdentity forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *).
HasCallStack =>
EpochInfo m -> EpochNo -> m EpochSize
epochInfoSize (Globals -> EpochInfo Identity
epochInfoPure Globals
testGlobals)

maxKESIterations :: Word64
maxKESIterations :: Word64
maxKESIterations = forall a. ShelleyBase a -> a
runShelleyBase (forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Globals -> Word64
maxKESEvo)

slotsPerKESIteration :: Word64
slotsPerKESIteration :: Word64
slotsPerKESIteration = forall a. ShelleyBase a -> a
runShelleyBase (forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Globals -> Word64
slotsPerKESPeriod)

maxLLSupply :: Coin
maxLLSupply :: Coin
maxLLSupply = Integer -> Coin
Coin forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ forall a. ShelleyBase a -> a
runShelleyBase (forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Globals -> Word64
maxLovelaceSupply)

testSTS ::
  forall s.
  (BaseM s ~ ShelleyBase, STS s, Eq (State s), Show (State s), ToExpr (State s)) =>
  Environment s ->
  State s ->
  Signal s ->
  Either (NonEmpty (PredicateFailure s)) (State s) ->
  Assertion
testSTS :: forall s.
(BaseM s ~ ReaderT Globals Identity, STS s, Eq (State s),
 Show (State s), ToExpr (State s)) =>
Environment s
-> State s
-> Signal s
-> Either (NonEmpty (PredicateFailure s)) (State s)
-> Assertion
testSTS Environment s
env State s
initSt Signal s
signal (Right State s
expectedSt) = do
  forall s (m :: * -> *).
(STS s, BaseM s ~ m) =>
(forall a. m a -> a)
-> Environment s
-> ReaderT
     (State s
      -> Signal s -> Either (NonEmpty (PredicateFailure s)) (State s))
     IO
     (State s)
-> Assertion
checkTrace @s forall a. ShelleyBase a -> a
runShelleyBase Environment s
env forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure State s
initSt forall (m :: * -> *) st sig err.
(MonadIO m, MonadReader (st -> sig -> Either err st) m, Show err,
 HasCallStack) =>
m st -> sig -> m st
.- Signal s
signal forall (m :: * -> *) st.
(MonadIO m, Eq st, ToExpr st, HasCallStack) =>
m st -> st -> m st
.->> State s
expectedSt
testSTS Environment s
env State s
initSt Signal s
sig predicateFailure :: Either (NonEmpty (PredicateFailure s)) (State s)
predicateFailure@(Left NonEmpty (PredicateFailure s)
_) = do
  let st :: Either (NonEmpty (PredicateFailure s)) (State s)
st = forall a. ShelleyBase a -> a
runShelleyBase forall a b. (a -> b) -> a -> b
$ forall s (m :: * -> *) (rtype :: RuleType).
(STS s, RuleTypeRep rtype, m ~ BaseM s) =>
RuleContext rtype s
-> m (Either (NonEmpty (PredicateFailure s)) (State s))
applySTSTest @s (forall sts. (Environment sts, State sts, Signal sts) -> TRC sts
TRC (Environment s
env, State s
initSt, Signal s
sig))
  Either (NonEmpty (PredicateFailure s)) (State s)
st forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Either (NonEmpty (PredicateFailure s)) (State s)
predicateFailure

mkHash :: forall a h. HashAlgorithm h => Int -> Hash h a
mkHash :: forall a h. HashAlgorithm h => Int -> Hash h a
mkHash Int
i = coerce :: forall a b. Coercible a b => a -> b
coerce (forall h a.
HashAlgorithm h =>
Version -> (a -> Encoding) -> a -> Hash h a
hashWithEncoder @h Version
shelleyProtVer forall a. EncCBOR a => a -> Encoding
encCBOR Int
i)

getBlockNonce :: forall era. Era era => Block (BHeader (EraCrypto era)) era -> Nonce
getBlockNonce :: forall era. Era era => Block (BHeader (EraCrypto era)) era -> Nonce
getBlockNonce =
  forall v. OutputVRF v -> Nonce
mkNonceFromOutputVRF forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall v a. CertifiedVRF v a -> OutputVRF v
certifiedOutput forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall c. BHBody c -> CertifiedVRF c Nonce
bheaderEta forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall c. Crypto c => BHeader c -> BHBody c
bhbody forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall h era. Block h era -> h
bheader