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

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

import Cardano.Ledger.Address (Addr (..))
import Cardano.Ledger.BaseTypes (Network (..))
import Cardano.Ledger.Coin (Coin (..))
import Cardano.Ledger.Compactible (fromCompact)
import Cardano.Ledger.Core (Era (..), EraTxOut, emptyPParams, mkCoinTxOut)
import Cardano.Ledger.Credential (Credential, StakeReference (..))
import Cardano.Ledger.EpochBoundary (SnapShot (..), Stake (..))
import Cardano.Ledger.Keys (KeyHash, KeyRole (..))
import Cardano.Ledger.PoolParams (PoolParams (..))
import Cardano.Ledger.Shelley.LedgerState (
  IncrementalStake (..),
  dsUnifiedL,
  incrementalStakeDistr,
  psStakePoolParamsL,
  updateStakeDistribution,
 )
import Cardano.Ledger.TxIn (TxIn)
import Cardano.Ledger.UMap (RDPair (..), rdRewardCoin)
import qualified Cardano.Ledger.UMap as UM
import Cardano.Ledger.UTxO (UTxO (..))
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.MonoTuple (TupleN)
import qualified Data.VMap as VMap
import Lens.Micro
import Test.Cardano.Ledger.Common
import Test.Cardano.Ledger.Core.Arbitrary ()

ppIdL :: Lens' (PoolParams c) (KeyHash 'StakePool c)
ppIdL :: forall c. Lens' (PoolParams c) (KeyHash 'StakePool c)
ppIdL = forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens forall c. PoolParams c -> KeyHash 'StakePool c
ppId (\PoolParams c
x KeyHash 'StakePool c
y -> PoolParams c
x {ppId :: KeyHash 'StakePool c
ppId = KeyHash 'StakePool c
y})

address :: Credential 'Payment c -> Maybe (Credential 'Staking c) -> Addr c
address :: forall c.
Credential 'Payment c -> Maybe (Credential 'Staking c) -> Addr c
address Credential 'Payment c
pc Maybe (Credential 'Staking c)
Nothing = forall c.
Network -> PaymentCredential c -> StakeReference c -> Addr c
Addr Network
Testnet Credential 'Payment c
pc forall c. StakeReference c
StakeRefNull
address Credential 'Payment c
pc (Just Credential 'Staking c
sc) = forall c.
Network -> PaymentCredential c -> StakeReference c -> Addr c
Addr Network
Testnet Credential 'Payment c
pc (forall c. StakeCredential c -> StakeReference c
StakeRefBase Credential 'Staking c
sc)

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 = do a
a <- forall a. Arbitrary a => Gen a
arbitrary; forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a forall a b. a -> (a -> b) -> b
& Lens' a b
l forall s t a b. ASetter s t a b -> b -> s -> t
.~ b
b)

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

stakeDistrIncludesRewards :: forall era. EraTxOut era => Gen Property
stakeDistrIncludesRewards :: forall era. EraTxOut era => Gen Property
stakeDistrIncludesRewards = do
  (Credential 'Staking (EraCrypto era)
tom, Credential 'Staking (EraCrypto era)
john, Credential 'Staking (EraCrypto era)
ann, Credential 'Staking (EraCrypto era)
ron, Credential 'Staking (EraCrypto era)
mary) <- forall a. Arbitrary a => Gen a
arbitrary @(TupleN 5 (Credential 'Staking (EraCrypto era)))
  (Credential 'Payment (EraCrypto era)
tomPay, Credential 'Payment (EraCrypto era)
johnPay, Credential 'Payment (EraCrypto era)
annPay, Credential 'Payment (EraCrypto era)
ronPay) <- forall a. Arbitrary a => Gen a
arbitrary @(TupleN 4 (Credential 'Payment (EraCrypto era)))

  (KeyHash 'StakePool (EraCrypto era)
pool1, KeyHash 'StakePool (EraCrypto era)
pool2) <- forall a. Arbitrary a => Gen a
arbitrary @(TupleN 2 (KeyHash 'StakePool (EraCrypto era)))
  PoolParams (EraCrypto era)
pool1Params <- forall a b. Arbitrary a => Lens' a b -> b -> Gen a
arbitraryLens forall c. Lens' (PoolParams c) (KeyHash 'StakePool c)
ppIdL KeyHash 'StakePool (EraCrypto era)
pool1
  PoolParams (EraCrypto era)
pool2Params <- forall a b. Arbitrary a => Lens' a b -> b -> Gen a
arbitraryLens forall c. Lens' (PoolParams c) (KeyHash 'StakePool c)
ppIdL KeyHash 'StakePool (EraCrypto era)
pool2

  (RDPair
tomRD, RDPair
johnRD, RDPair
annRD, RDPair
ronRD, RDPair
maryRD) <- forall a. Arbitrary a => Gen a
arbitrary @(TupleN 5 RDPair)

  let tomAddr :: Addr (EraCrypto era)
tomAddr = forall c.
Credential 'Payment c -> Maybe (Credential 'Staking c) -> Addr c
address Credential 'Payment (EraCrypto era)
tomPay forall a. Maybe a
Nothing -- Nothing means tomAddr does not have a StakeReference
      johnAddr :: Addr (EraCrypto era)
johnAddr = forall c.
Credential 'Payment c -> Maybe (Credential 'Staking c) -> Addr c
address Credential 'Payment (EraCrypto era)
johnPay (forall a. a -> Maybe a
Just Credential 'Staking (EraCrypto era)
john)
      annAddr :: Addr (EraCrypto era)
annAddr = forall c.
Credential 'Payment c -> Maybe (Credential 'Staking c) -> Addr c
address Credential 'Payment (EraCrypto era)
annPay (forall a. a -> Maybe a
Just Credential 'Staking (EraCrypto era)
ann)
      ronAddr :: Addr (EraCrypto era)
ronAddr = forall c.
Credential 'Payment c -> Maybe (Credential 'Staking c) -> Addr c
address Credential 'Payment (EraCrypto era)
ronPay (forall a. a -> Maybe a
Just Credential 'Staking (EraCrypto era)
ron)
      -- maryAddr is omitted on purpose. Mary will not have a UTxO entry

      rewards :: Map (Credential 'Staking (EraCrypto era)) RDPair
      rewards :: Map (Credential 'Staking (EraCrypto era)) RDPair
rewards =
        forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList
          [ (Credential 'Staking (EraCrypto era)
tom, RDPair
tomRD)
          , (Credential 'Staking (EraCrypto era)
john, RDPair
johnRD)
          , (Credential 'Staking (EraCrypto era)
ann, RDPair
annRD)
          , (Credential 'Staking (EraCrypto era)
ron, RDPair
ronRD)
          , (Credential 'Staking (EraCrypto era)
mary, RDPair
maryRD)
          ]

      delegations :: Map (Credential 'Staking (EraCrypto era)) (KeyHash 'StakePool (EraCrypto era))
      delegations :: Map
  (Credential 'Staking (EraCrypto era))
  (KeyHash 'StakePool (EraCrypto era))
delegations =
        forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList
          [ (Credential 'Staking (EraCrypto era)
tom, KeyHash 'StakePool (EraCrypto era)
pool1) -- since every one is delegated
          , (Credential 'Staking (EraCrypto era)
ann, KeyHash 'StakePool (EraCrypto era)
pool2) -- no one's stake should be left out
          , (Credential 'Staking (EraCrypto era)
ron, KeyHash 'StakePool (EraCrypto era)
pool1)
          , (Credential 'Staking (EraCrypto era)
john, KeyHash 'StakePool (EraCrypto era)
pool2)
          , (Credential 'Staking (EraCrypto era)
mary, KeyHash 'StakePool (EraCrypto era)
pool2)
          ]

  -- Again mary is not included, because she will not have an UTxO entry, but tom will have 2
  (Coin
tomCoin1, Coin
tomCoin2, Coin
johnCoin, Coin
annCoin, Coin
ronCoin) <- forall a. Arbitrary a => Gen a
arbitrary @(TupleN 5 Coin)

  (TxIn (EraCrypto era)
tomTxIn1, TxIn (EraCrypto era)
tomTxIn2, TxIn (EraCrypto era)
johnTxIn, TxIn (EraCrypto era)
annTxIn, TxIn (EraCrypto era)
ronTxIn) <- forall a. Arbitrary a => Gen a
arbitrary @(TupleN 5 (TxIn (EraCrypto era)))

  let
    -- Each wallet (except mary) has one or more UTxO entries
    -- Since tom uses a StakeRefNull those entries will not be distributed
    utxo1 :: UTxO era
utxo1 =
      forall era. Map (TxIn (EraCrypto era)) (TxOut era) -> UTxO era
UTxO @era
        ( forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList
            [ (TxIn (EraCrypto era)
tomTxIn1, forall era.
EraTxOut era =>
Addr (EraCrypto era) -> Coin -> TxOut era
mkCoinTxOut Addr (EraCrypto era)
tomAddr Coin
tomCoin1) -- Not distrubuted, see tomAddr
            , (TxIn (EraCrypto era)
tomTxIn2, forall era.
EraTxOut era =>
Addr (EraCrypto era) -> Coin -> TxOut era
mkCoinTxOut Addr (EraCrypto era)
tomAddr Coin
tomCoin2) -- Not distributed, see tomAddr
            , (TxIn (EraCrypto era)
annTxIn, forall era.
EraTxOut era =>
Addr (EraCrypto era) -> Coin -> TxOut era
mkCoinTxOut Addr (EraCrypto era)
annAddr Coin
annCoin)
            , (TxIn (EraCrypto era)
ronTxIn, forall era.
EraTxOut era =>
Addr (EraCrypto era) -> Coin -> TxOut era
mkCoinTxOut Addr (EraCrypto era)
ronAddr Coin
ronCoin)
            , (TxIn (EraCrypto era)
johnTxIn, forall era.
EraTxOut era =>
Addr (EraCrypto era) -> Coin -> TxOut era
mkCoinTxOut Addr (EraCrypto era)
johnAddr Coin
johnCoin)
            -- Note Mary does not have a UTxO entry, but her rewards are still counted
            ]
        )

    pparams :: PParams era
pparams = forall era. EraPParams era => PParams era
emptyPParams @era

    incrementalStake :: IncrementalStake (EraCrypto era)
incrementalStake = forall era.
EraTxOut era =>
PParams era
-> IncrementalStake (EraCrypto era)
-> UTxO era
-> UTxO era
-> IncrementalStake (EraCrypto era)
updateStakeDistribution PParams era
pparams (forall c.
Map (Credential 'Staking c) (CompactForm Coin)
-> Map Ptr (CompactForm Coin) -> IncrementalStake c
IStake forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty) forall a. Monoid a => a
mempty UTxO era
utxo1
    umap :: UMap (EraCrypto era)
umap = forall c.
Map (Credential 'Staking c) RDPair
-> Map Ptr (Credential 'Staking c)
-> Map (Credential 'Staking c) (KeyHash 'StakePool c)
-> Map (Credential 'Staking c) (DRep c)
-> UMap c
UM.unify Map (Credential 'Staking (EraCrypto era)) RDPair
rewards forall k a. Map k a
Map.empty Map
  (Credential 'Staking (EraCrypto era))
  (KeyHash 'StakePool (EraCrypto era))
delegations forall k a. Map k a
Map.empty
    poolparamMap :: Map
  (KeyHash 'StakePool (EraCrypto era)) (PoolParams (EraCrypto era))
poolparamMap = forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(KeyHash 'StakePool (EraCrypto era)
pool1, PoolParams (EraCrypto era)
pool1Params), (KeyHash 'StakePool (EraCrypto era)
pool2, PoolParams (EraCrypto era)
pool2Params)]
  -- We can either use an emptyDstate with just the umap, like this
  -- dState = (emptyDState {dsUnified = umap}) :: DState era
  -- Or an arbitrary one, where we overwrite the umap, with the one we need.
  DState era
dState <- forall a b. Arbitrary a => Lens' a b -> b -> Gen a
arbitraryLens forall era. Lens' (DState era) (UMap (EraCrypto era))
dsUnifiedL UMap (EraCrypto era)
umap
  PState era
pState <- forall a b. Arbitrary a => Lens' a b -> b -> Gen a
arbitraryLens forall era.
Lens'
  (PState era)
  (Map
     (KeyHash 'StakePool (EraCrypto era)) (PoolParams (EraCrypto era)))
psStakePoolParamsL Map
  (KeyHash 'StakePool (EraCrypto era)) (PoolParams (EraCrypto era))
poolparamMap
  let computedStakeDistr :: Map (Credential 'Staking (EraCrypto era)) Coin
computedStakeDistr = forall a b k. (a -> b) -> Map k a -> Map k b
Map.map forall a. Compactible a => CompactForm a -> a
fromCompact (forall (kv :: * -> *) k (vv :: * -> *) v.
(Vector kv k, Vector vv v) =>
VMap kv vv k v -> Map k v
VMap.toMap (forall c.
Stake c -> VMap VB VP (Credential 'Staking c) (CompactForm Coin)
unStake (forall c. SnapShot c -> Stake c
ssStake SnapShot (EraCrypto era)
snap)))
        where
          snap :: SnapShot (EraCrypto era)
snap =
            forall era.
EraPParams era =>
PParams era
-> IncrementalStake (EraCrypto era)
-> DState era
-> PState era
-> SnapShot (EraCrypto era)
incrementalStakeDistr -- This computes the actual Incremental Stake
              PParams era
pparams
              IncrementalStake (EraCrypto era)
incrementalStake
              DState era
dState
              PState era
pState

      expectedStakeDistr :: Map (Credential 'Staking (EraCrypto era)) Coin
      expectedStakeDistr :: Map (Credential 'Staking (EraCrypto era)) Coin
expectedStakeDistr =
        forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList -- Coin Part is (rdRewardCoin <> utxoCoin)
          [ (Credential 'Staking (EraCrypto era)
tom, RDPair -> Coin
rdRewardCoin RDPair
tomRD forall a. Semigroup a => a -> a -> a
<> Integer -> Coin
Coin Integer
0) -- tom uxtxoCoin is zero because his address has StakeRefNull
          , (Credential 'Staking (EraCrypto era)
ann, RDPair -> Coin
rdRewardCoin RDPair
annRD forall a. Semigroup a => a -> a -> a
<> Coin
annCoin)
          , (Credential 'Staking (EraCrypto era)
ron, RDPair -> Coin
rdRewardCoin RDPair
ronRD forall a. Semigroup a => a -> a -> a
<> Coin
ronCoin)
          , (Credential 'Staking (EraCrypto era)
john, RDPair -> Coin
rdRewardCoin RDPair
johnRD forall a. Semigroup a => a -> a -> a
<> Coin
johnCoin)
          , (Credential 'Staking (EraCrypto era)
mary, RDPair -> Coin
rdRewardCoin RDPair
maryRD forall a. Semigroup a => a -> a -> a
<> Integer -> Coin
Coin Integer
0) -- mary uxtxoCoin is zero because she has no UtxO entry
          ]

  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ (Map (Credential 'Staking (EraCrypto era)) Coin
computedStakeDistr forall a. (Eq a, Show a) => a -> a -> Property
=== Map (Credential 'Staking (EraCrypto era)) Coin
expectedStakeDistr)

spec :: forall era. EraTxOut era => Spec
spec :: forall era. EraTxOut era => Spec
spec = forall prop.
(HasCallStack, Testable prop) =>
String -> prop -> Spec
prop String
"StakeDistrIncludesRewards" (forall era. EraTxOut era => Gen Property
stakeDistrIncludesRewards @era)