{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

module Test.Cardano.Ledger.Shelley.UnitTests.InstantStakeTest (spec) where

import Cardano.Ledger.Coin (Coin (..))
import Cardano.Ledger.Compactible (CompactForm, fromCompact)
import Cardano.Ledger.Core
import Cardano.Ledger.Credential (Credential, StakeReference (..))
import Cardano.Ledger.Shelley.State
import Cardano.Ledger.TxIn (TxIn)
import Data.Default (def)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.MonoTuple (TupleN)
import qualified Data.Set as Set
import qualified Data.VMap as VMap
import Lens.Micro
import Test.Cardano.Ledger.Common
import Test.Cardano.Ledger.Core.Arbitrary ()
import Test.Cardano.Ledger.Core.KeyPair (mkAddr)
import Test.Cardano.Ledger.Shelley.Era
import Test.Cardano.Ledger.Shelley.ImpTest

sppIdL :: Lens' StakePoolParams (KeyHash StakePool)
sppIdL :: Lens' StakePoolParams (KeyHash StakePool)
sppIdL = (StakePoolParams -> KeyHash StakePool)
-> (StakePoolParams -> KeyHash StakePool -> StakePoolParams)
-> Lens' StakePoolParams (KeyHash StakePool)
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens StakePoolParams -> KeyHash StakePool
sppId (\StakePoolParams
x KeyHash StakePool
y -> StakePoolParams
x {sppId = y})

-- | Generate an arbitrary value and overwrite the specified value using the supplied lens.
arbitraryLens :: Arbitrary a => Lens' a b -> b -> Gen a
arbitraryLens :: forall a b. Arbitrary a => Lens' a b -> b -> Gen a
arbitraryLens Lens' a b
l b
b = ((b -> Identity b) -> a -> Identity a
Lens' a b
l ((b -> Identity b) -> a -> Identity a) -> b -> a -> a
forall s t a b. ASetter s t a b -> b -> s -> t
.~ b
b) (a -> a) -> Gen a -> Gen a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen a
forall a. Arbitrary a => Gen a
arbitrary

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

instantStakeIncludesRewards :: forall era. ShelleyEraImp era => Gen Property
instantStakeIncludesRewards :: forall era. ShelleyEraImp era => Gen Property
instantStakeIncludesRewards = do
  (pool1, pool2) <- forall a. Arbitrary a => Gen a
arbitrary @(TupleN 2 StakePoolParams)
  let
    poolId1 = StakePoolParams
pool1 StakePoolParams
-> Getting (KeyHash StakePool) StakePoolParams (KeyHash StakePool)
-> KeyHash StakePool
forall s a. s -> Getting a s a -> a
^. Getting (KeyHash StakePool) StakePoolParams (KeyHash StakePool)
Lens' StakePoolParams (KeyHash StakePool)
sppIdL
    poolId2 = StakePoolParams
pool2 StakePoolParams
-> Getting (KeyHash StakePool) StakePoolParams (KeyHash StakePool)
-> KeyHash StakePool
forall s a. s -> Getting a s a -> a
^. Getting (KeyHash StakePool) StakePoolParams (KeyHash StakePool)
Lens' StakePoolParams (KeyHash StakePool)
sppIdL
    poolParamsMap = [(KeyHash StakePool, StakePoolParams)]
-> Map (KeyHash StakePool) StakePoolParams
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(KeyHash StakePool
poolId1, StakePoolParams
pool1), (KeyHash StakePool
poolId2, StakePoolParams
pool2)]
  pState <- arbitraryLens psStakePoolsL $ mkStakePoolState mempty mempty <$> poolParamsMap
  let
    initCertState :: CertState era
    initCertState = CertState era
forall a. Default a => a
def CertState era -> (CertState era -> CertState era) -> CertState era
forall a b. a -> (a -> b) -> b
& (PState era -> Identity (PState era))
-> CertState era -> Identity (CertState era)
forall era. EraCertState era => Lens' (CertState era) (PState era)
Lens' (CertState era) (PState era)
certPStateL ((PState era -> Identity (PState era))
 -> CertState era -> Identity (CertState era))
-> PState era -> CertState era -> CertState era
forall s t a b. ASetter s t a b -> b -> s -> t
.~ PState era
pState
    registerAccount KeyHash StakePool
poolId CertState era
certState = do
      stakingCredential <- Gen (Credential Staking)
forall a. Arbitrary a => Gen a
arbitrary
      ptr <- arbitrary
      deposit <- arbitrary
      balance <- arbitrary
      let accounts :: Accounts era
          accounts =
            Map (Credential Staking) (CompactForm Coin)
-> Accounts era -> Accounts era
forall era.
EraAccounts era =>
Map (Credential Staking) (CompactForm Coin)
-> Accounts era -> Accounts era
addToBalanceAccounts (Credential Staking
-> CompactForm Coin -> Map (Credential Staking) (CompactForm Coin)
forall k a. k -> a -> Map k a
Map.singleton Credential Staking
stakingCredential CompactForm Coin
balance) (Accounts era -> Accounts era) -> Accounts era -> Accounts era
forall a b. (a -> b) -> a -> b
$
              Credential Staking
-> Maybe Ptr
-> CompactForm Coin
-> Maybe (KeyHash StakePool)
-> Maybe DRep
-> Accounts era
-> Accounts era
forall era.
(HasCallStack, EraTest era) =>
Credential Staking
-> Maybe Ptr
-> CompactForm Coin
-> Maybe (KeyHash StakePool)
-> Maybe DRep
-> Accounts era
-> Accounts era
registerTestAccount Credential Staking
stakingCredential (Ptr -> Maybe Ptr
forall a. a -> Maybe a
Just Ptr
ptr) CompactForm Coin
deposit (KeyHash StakePool -> Maybe (KeyHash StakePool)
forall a. a -> Maybe a
Just KeyHash StakePool
poolId) Maybe DRep
forall a. Maybe a
Nothing (Accounts era -> Accounts era) -> Accounts era -> Accounts era
forall a b. (a -> b) -> a -> b
$
                CertState era
certState CertState era
-> Getting (Accounts era) (CertState era) (Accounts era)
-> Accounts era
forall s a. s -> Getting a s a -> a
^. (DState era -> Const (Accounts era) (DState era))
-> CertState era -> Const (Accounts era) (CertState era)
forall era. EraCertState era => Lens' (CertState era) (DState era)
Lens' (CertState era) (DState era)
certDStateL ((DState era -> Const (Accounts era) (DState era))
 -> CertState era -> Const (Accounts era) (CertState era))
-> ((Accounts era -> Const (Accounts era) (Accounts era))
    -> DState era -> Const (Accounts era) (DState era))
-> Getting (Accounts era) (CertState era) (Accounts era)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Accounts era -> Const (Accounts era) (Accounts era))
-> DState era -> Const (Accounts era) (DState era)
forall era. Lens' (DState era) (Accounts era)
forall (t :: * -> *) era.
CanSetAccounts t =>
Lens' (t era) (Accounts era)
accountsL
          certState' =
            CertState era
certState
              CertState era -> (CertState era -> CertState era) -> CertState era
forall a b. a -> (a -> b) -> b
& (DState era -> Identity (DState era))
-> CertState era -> Identity (CertState era)
forall era. EraCertState era => Lens' (CertState era) (DState era)
Lens' (CertState era) (DState era)
certDStateL ((DState era -> Identity (DState era))
 -> CertState era -> Identity (CertState era))
-> ((Accounts era -> Identity (Accounts era))
    -> DState era -> Identity (DState era))
-> (Accounts era -> Identity (Accounts era))
-> CertState era
-> Identity (CertState era)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Accounts era -> Identity (Accounts era))
-> DState era -> Identity (DState era)
forall era. Lens' (DState era) (Accounts era)
forall (t :: * -> *) era.
CanSetAccounts t =>
Lens' (t era) (Accounts era)
accountsL ((Accounts era -> Identity (Accounts era))
 -> CertState era -> Identity (CertState era))
-> Accounts era -> CertState era -> CertState era
forall s t a b. ASetter s t a b -> b -> s -> t
.~ Accounts era
accounts
              CertState era -> (CertState era -> CertState era) -> CertState era
forall a b. a -> (a -> b) -> b
& (PState era -> Identity (PState era))
-> CertState era -> Identity (CertState era)
forall era. EraCertState era => Lens' (CertState era) (PState era)
Lens' (CertState era) (PState era)
certPStateL ((PState era -> Identity (PState era))
 -> CertState era -> Identity (CertState era))
-> ((Map (KeyHash StakePool) StakePoolState
     -> Identity (Map (KeyHash StakePool) StakePoolState))
    -> PState era -> Identity (PState era))
-> (Map (KeyHash StakePool) StakePoolState
    -> Identity (Map (KeyHash StakePool) StakePoolState))
-> CertState era
-> Identity (CertState era)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Map (KeyHash StakePool) StakePoolState
 -> Identity (Map (KeyHash StakePool) StakePoolState))
-> PState era -> Identity (PState era)
forall era (f :: * -> *).
Functor f =>
(Map (KeyHash StakePool) StakePoolState
 -> f (Map (KeyHash StakePool) StakePoolState))
-> PState era -> f (PState era)
psStakePoolsL
                ((Map (KeyHash StakePool) StakePoolState
  -> Identity (Map (KeyHash StakePool) StakePoolState))
 -> CertState era -> Identity (CertState era))
-> (Map (KeyHash StakePool) StakePoolState
    -> Map (KeyHash StakePool) StakePoolState)
-> CertState era
-> CertState era
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ (StakePoolState -> StakePoolState)
-> KeyHash StakePool
-> Map (KeyHash StakePool) StakePoolState
-> Map (KeyHash StakePool) StakePoolState
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
Map.adjust ((Set (Credential Staking) -> Identity (Set (Credential Staking)))
-> StakePoolState -> Identity StakePoolState
Lens' StakePoolState (Set (Credential Staking))
spsDelegatorsL ((Set (Credential Staking) -> Identity (Set (Credential Staking)))
 -> StakePoolState -> Identity StakePoolState)
-> (Set (Credential Staking) -> Set (Credential Staking))
-> StakePoolState
-> StakePoolState
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ Credential Staking
-> Set (Credential Staking) -> Set (Credential Staking)
forall a. Ord a => a -> Set a -> Set a
Set.insert Credential Staking
stakingCredential) KeyHash StakePool
poolId
      pure (stakingCredential, balance, certState')
  (tom, tomBalance, certState0) <- registerAccount poolId1 initCertState
  (john, johnBalance, certState1) <- registerAccount poolId2 certState0
  (ann, annBalance, certState2) <- registerAccount poolId2 certState1
  (ron, ronBalance, certState3) <- registerAccount poolId1 certState2
  (mary, maryBalance, certState) <- registerAccount poolId2 certState3
  (tomPay, johnPay, annPay, ronPay) <- arbitrary @(TupleN 4 (Credential Payment))

  let tomAddr = Credential Payment -> StakeReference -> Addr
forall p s.
(MakeCredential p Payment, MakeStakeReference s) =>
p -> s -> Addr
mkAddr Credential Payment
tomPay StakeReference
StakeRefNull
      johnAddr = Credential Payment -> Credential Staking -> Addr
forall p s.
(MakeCredential p Payment, MakeStakeReference s) =>
p -> s -> Addr
mkAddr Credential Payment
johnPay Credential Staking
john
      annAddr = Credential Payment -> Credential Staking -> Addr
forall p s.
(MakeCredential p Payment, MakeStakeReference s) =>
p -> s -> Addr
mkAddr Credential Payment
annPay Credential Staking
ann
      ronAddr = Credential Payment -> Credential Staking -> Addr
forall p s.
(MakeCredential p Payment, MakeStakeReference s) =>
p -> s -> Addr
mkAddr Credential Payment
ronPay Credential Staking
ron
  -- maryAddr is omitted on purpose. Mary will not have a UTxO entry

  -- Again mary is not included, because she will not have an UTxO entry, but tom will have 2
  (tomCoin1, tomCoin2, johnCoin, annCoin, ronCoin) <- arbitrary @(TupleN 5 (CompactForm Coin))

  (tomTxIn1, tomTxIn2, johnTxIn, annTxIn, ronTxIn) <- arbitrary @(TupleN 5 TxIn)

  let
    -- Each wallet (except mary) has one or more UTxO entries
    -- Since tom uses a StakeRefNull those entries will not be distributed
    utxo1 =
      forall era. Map TxIn (TxOut era) -> UTxO era
UTxO @era
        ( [(TxIn, TxOut era)] -> Map TxIn (TxOut era)
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList
            [ (TxIn
tomTxIn1, Addr -> Coin -> TxOut era
forall era. EraTxOut era => Addr -> Coin -> TxOut era
mkCoinTxOut Addr
tomAddr (Coin -> TxOut era) -> Coin -> TxOut era
forall a b. (a -> b) -> a -> b
$ CompactForm Coin -> Coin
forall a. Compactible a => CompactForm a -> a
fromCompact CompactForm Coin
tomCoin1) -- Not distrubuted, see tomAddr
            , (TxIn
tomTxIn2, Addr -> Coin -> TxOut era
forall era. EraTxOut era => Addr -> Coin -> TxOut era
mkCoinTxOut Addr
tomAddr (Coin -> TxOut era) -> Coin -> TxOut era
forall a b. (a -> b) -> a -> b
$ CompactForm Coin -> Coin
forall a. Compactible a => CompactForm a -> a
fromCompact CompactForm Coin
tomCoin2) -- Not distributed, see tomAddr
            , (TxIn
annTxIn, Addr -> Coin -> TxOut era
forall era. EraTxOut era => Addr -> Coin -> TxOut era
mkCoinTxOut Addr
annAddr (Coin -> TxOut era) -> Coin -> TxOut era
forall a b. (a -> b) -> a -> b
$ CompactForm Coin -> Coin
forall a. Compactible a => CompactForm a -> a
fromCompact CompactForm Coin
annCoin)
            , (TxIn
ronTxIn, Addr -> Coin -> TxOut era
forall era. EraTxOut era => Addr -> Coin -> TxOut era
mkCoinTxOut Addr
ronAddr (Coin -> TxOut era) -> Coin -> TxOut era
forall a b. (a -> b) -> a -> b
$ CompactForm Coin -> Coin
forall a. Compactible a => CompactForm a -> a
fromCompact CompactForm Coin
ronCoin)
            , (TxIn
johnTxIn, Addr -> Coin -> TxOut era
forall era. EraTxOut era => Addr -> Coin -> TxOut era
mkCoinTxOut Addr
johnAddr (Coin -> TxOut era) -> Coin -> TxOut era
forall a b. (a -> b) -> a -> b
$ CompactForm Coin -> Coin
forall a. Compactible a => CompactForm a -> a
fromCompact CompactForm Coin
johnCoin)
            -- Note Mary does not have a UTxO entry, but her rewards are still counted
            ]
        )
    instantStake = UTxO era -> InstantStake era -> InstantStake era
forall era.
EraStake era =>
UTxO era -> InstantStake era -> InstantStake era
addInstantStake UTxO era
utxo1 InstantStake era
forall a. Monoid a => a
mempty
  network <- arbitrary
  let snapShot =
        InstantStake era -> DState era -> PState era -> Network -> SnapShot
forall era.
EraStake era =>
InstantStake era -> DState era -> PState era -> Network -> SnapShot
snapShotFromInstantStake InstantStake era
instantStake (CertState era
certState CertState era
-> Getting (DState era) (CertState era) (DState era) -> DState era
forall s a. s -> Getting a s a -> a
^. Getting (DState era) (CertState era) (DState era)
forall era. EraCertState era => Lens' (CertState era) (DState era)
Lens' (CertState era) (DState era)
certDStateL) (CertState era
certState CertState era
-> Getting (PState era) (CertState era) (PState era) -> PState era
forall s a. s -> Getting a s a -> a
^. Getting (PState era) (CertState era) (PState era)
forall era. EraCertState era => Lens' (CertState era) (PState era)
Lens' (CertState era) (PState era)
certPStateL) Network
network
      computedStakeDistr = VMap VB VP (Credential Staking) (CompactForm Coin)
-> Map (Credential Staking) (CompactForm Coin)
forall (kv :: * -> *) k (vv :: * -> *) v.
(Vector kv k, Vector vv v) =>
VMap kv vv k v -> Map k v
VMap.toMap (Stake -> VMap VB VP (Credential Staking) (CompactForm Coin)
unStake (SnapShot -> Stake
ssStake SnapShot
snapShot))

      expectedStakeDistr :: Map (Credential Staking) (CompactForm Coin)
      expectedStakeDistr =
        [(Credential Staking, CompactForm Coin)]
-> Map (Credential Staking) (CompactForm Coin)
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(Credential Staking, CompactForm Coin)]
 -> Map (Credential Staking) (CompactForm Coin))
-> [(Credential Staking, CompactForm Coin)]
-> Map (Credential Staking) (CompactForm Coin)
forall a b. (a -> b) -> a -> b
$
          [ (Credential Staking, CompactForm Coin)
stake
          | stake :: (Credential Staking, CompactForm Coin)
stake@(Credential Staking
_, CompactForm Coin
balance) <-
              [ (Credential Staking
ann, CompactForm Coin
annBalance CompactForm Coin -> CompactForm Coin -> CompactForm Coin
forall a. Semigroup a => a -> a -> a
<> CompactForm Coin
annCoin)
              , (Credential Staking
ron, CompactForm Coin
ronBalance CompactForm Coin -> CompactForm Coin -> CompactForm Coin
forall a. Semigroup a => a -> a -> a
<> CompactForm Coin
ronCoin)
              , (Credential Staking
john, CompactForm Coin
johnBalance CompactForm Coin -> CompactForm Coin -> CompactForm Coin
forall a. Semigroup a => a -> a -> a
<> CompactForm Coin
johnCoin)
              , (Credential Staking
tom, CompactForm Coin
tomBalance)
              , (Credential Staking
mary, CompactForm Coin
maryBalance)
              ]
          , -- We need to filter out zero rewards from instant stake, since all coins are
          -- generated at random, stake for any one of the accounts can be zero
          CompactForm Coin
balance CompactForm Coin -> CompactForm Coin -> Bool
forall a. Eq a => a -> a -> Bool
/= CompactForm Coin
forall a. Monoid a => a
mempty
          ]

  pure (computedStakeDistr === expectedStakeDistr)

spec :: forall era. ShelleyEraImp era => Spec
spec :: forall era. ShelleyEraImp era => Spec
spec = String -> Gen Property -> Spec
forall prop.
(HasCallStack, Testable prop) =>
String -> prop -> Spec
prop String
"InstantStakeIncludesRewards" (forall era. ShelleyEraImp era => Gen Property
instantStakeIncludesRewards @era)