{-# LANGUAGE TypeApplications #-}

module Test.Cardano.Ledger.Shelley.Examples (
  CHAINExample (..),
  testCHAINExample,
)
where

import Cardano.Ledger.Block (Block)
import Cardano.Ledger.Shelley ()
import Cardano.Ledger.Shelley.LedgerState (nesPdL)
import Cardano.Ledger.Shelley.Scripts ()
import Cardano.Ledger.State (individualTotalPoolStakeL, poolDistrDistrL, poolDistrTotalL)
import Cardano.Protocol.TPraos.BHeader (BHeader)
import Control.State.Transition.Extended hiding (Assertion)
import Data.List.NonEmpty (NonEmpty)
import GHC.Stack
import Lens.Micro
import Test.Cardano.Ledger.Shelley.ConcreteCryptoTypes (C, MockCrypto)
import Test.Cardano.Ledger.Shelley.Rules.Chain (CHAIN, ChainState, chainStateNesL, totalAda)
import Test.Cardano.Ledger.Shelley.TreeDiff (expectExprEqual)
import Test.Cardano.Ledger.Shelley.Utils (applySTSTest, maxLLSupply, runShelleyBase)
import Test.Control.State.Transition.Trace (checkTrace, (.-), (.->>))
import Test.Tasty.HUnit (Assertion, (@?=))

data CHAINExample era = CHAINExample
  { forall era. CHAINExample era -> ChainState era
startState :: ChainState era
  -- ^ State to start testing with
  , forall era. CHAINExample era -> Block (BHeader MockCrypto) era
newBlock :: Block (BHeader MockCrypto) era
  -- ^ Block to run chain state transition system on
  , forall era.
CHAINExample era
-> Either
     (NonEmpty (PredicateFailure (CHAIN era))) (ChainState era)
intendedResult :: Either (NonEmpty (PredicateFailure (CHAIN era))) (ChainState era)
  -- ^ type of fatal error, if failure expected and final chain state if success expected
  }

-- | Runs example, applies chain state transition system rule (STS),
--   and checks that trace ends with expected state or expected error.
testCHAINExample :: HasCallStack => CHAINExample C -> Assertion
testCHAINExample :: HasCallStack => CHAINExample C -> Assertion
testCHAINExample (CHAINExample ChainState C
initSt Block (BHeader MockCrypto) C
block (Right ChainState C
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 @(CHAIN C) forall a. ShelleyBase a -> a
runShelleyBase () forall a b. (a -> b) -> a -> b
$
      ( forall (f :: * -> *) a. Applicative f => a -> f a
pure ChainState C
initSt forall (m :: * -> *) st sig err.
(MonadIO m, MonadReader (st -> sig -> Either err st) m, Show err,
 HasCallStack) =>
m st -> sig -> m st
.- Block (BHeader MockCrypto) C
block
          forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> forall era. Lens' (ChainState era) (NewEpochState era)
chainStateNesL forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall era. Lens' (NewEpochState era) PoolDistr
nesPdL forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lens' PoolDistr (CompactForm Coin)
poolDistrTotalL forall s t a b. ASetter s t a b -> b -> s -> t
.~ forall a. Monoid a => a
mempty
          forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> forall era. Lens' (ChainState era) (NewEpochState era)
chainStateNesL forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall era. Lens' (NewEpochState era) PoolDistr
nesPdL forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lens' PoolDistr (Map (KeyHash 'StakePool) IndividualPoolStake)
poolDistrDistrL forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ (forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> Lens' IndividualPoolStake (CompactForm Coin)
individualTotalPoolStakeL forall s t a b. ASetter s t a b -> b -> s -> t
.~ forall a. Monoid a => a
mempty)
      )
        forall (m :: * -> *) st.
(MonadIO m, Eq st, ToExpr st, HasCallStack) =>
m st -> st -> m st
.->> ( ChainState C
expectedSt
                forall a b. a -> (a -> b) -> b
& forall era. Lens' (ChainState era) (NewEpochState era)
chainStateNesL forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall era. Lens' (NewEpochState era) PoolDistr
nesPdL forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lens' PoolDistr (CompactForm Coin)
poolDistrTotalL forall s t a b. ASetter s t a b -> b -> s -> t
.~ forall a. Monoid a => a
mempty
                forall a b. a -> (a -> b) -> b
& forall era. Lens' (ChainState era) (NewEpochState era)
chainStateNesL forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall era. Lens' (NewEpochState era) PoolDistr
nesPdL forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lens' PoolDistr (Map (KeyHash 'StakePool) IndividualPoolStake)
poolDistrDistrL forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ (forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> Lens' IndividualPoolStake (CompactForm Coin)
individualTotalPoolStakeL forall s t a b. ASetter s t a b -> b -> s -> t
.~ forall a. Monoid a => a
mempty)
             )
    )
    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall a. (Eq a, ToExpr a) => a -> a -> Assertion
expectExprEqual (forall era. (EraTxOut era, EraGov era) => ChainState era -> Coin
totalAda ChainState C
expectedSt) Coin
maxLLSupply
testCHAINExample (CHAINExample ChainState C
initSt Block (BHeader MockCrypto) C
block predicateFailure :: Either (NonEmpty (PredicateFailure (CHAIN C))) (ChainState C)
predicateFailure@(Left NonEmpty (PredicateFailure (CHAIN C))
_)) = do
  let st :: Either (NonEmpty (TestChainPredicateFailure C)) (ChainState C)
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 @(CHAIN C) (forall sts. (Environment sts, State sts, Signal sts) -> TRC sts
TRC ((), ChainState C
initSt, Block (BHeader MockCrypto) C
block))
  Either (NonEmpty (TestChainPredicateFailure C)) (ChainState C)
st forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Either (NonEmpty (PredicateFailure (CHAIN C))) (ChainState C)
predicateFailure