{-# 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 (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 (KeyHash 'StakePool)
ppIdL :: Lens' PoolParams (KeyHash 'StakePool)
ppIdL = forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens PoolParams -> KeyHash 'StakePool
ppId (\PoolParams
x KeyHash 'StakePool
y -> PoolParams
x {ppId :: KeyHash 'StakePool
ppId = KeyHash 'StakePool
y})

address :: Credential 'Payment -> Maybe (Credential 'Staking) -> Addr
address :: Credential 'Payment -> Maybe (Credential 'Staking) -> Addr
address Credential 'Payment
pc Maybe (Credential 'Staking)
Nothing = Network -> Credential 'Payment -> StakeReference -> Addr
Addr Network
Testnet Credential 'Payment
pc StakeReference
StakeRefNull
address Credential 'Payment
pc (Just Credential 'Staking
sc) = Network -> Credential 'Payment -> StakeReference -> Addr
Addr Network
Testnet Credential 'Payment
pc (Credential 'Staking -> StakeReference
StakeRefBase Credential 'Staking
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
tom, Credential 'Staking
john, Credential 'Staking
ann, Credential 'Staking
ron, Credential 'Staking
mary) <- forall a. Arbitrary a => Gen a
arbitrary @(TupleN 5 (Credential 'Staking))
  (Credential 'Payment
tomPay, Credential 'Payment
johnPay, Credential 'Payment
annPay, Credential 'Payment
ronPay) <- forall a. Arbitrary a => Gen a
arbitrary @(TupleN 4 (Credential 'Payment))

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

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

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

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

      delegations :: Map (Credential 'Staking) (KeyHash 'StakePool)
      delegations :: Map (Credential 'Staking) (KeyHash 'StakePool)
delegations =
        forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList
          [ (Credential 'Staking
tom, KeyHash 'StakePool
pool1) -- since every one is delegated
          , (Credential 'Staking
ann, KeyHash 'StakePool
pool2) -- no one's stake should be left out
          , (Credential 'Staking
ron, KeyHash 'StakePool
pool1)
          , (Credential 'Staking
john, KeyHash 'StakePool
pool2)
          , (Credential 'Staking
mary, KeyHash 'StakePool
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
tomTxIn1, TxIn
tomTxIn2, TxIn
johnTxIn, TxIn
annTxIn, TxIn
ronTxIn) <- forall a. Arbitrary a => Gen a
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 :: UTxO era
utxo1 =
      forall era. Map TxIn (TxOut era) -> UTxO era
UTxO @era
        ( forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList
            [ (TxIn
tomTxIn1, forall era. EraTxOut era => Addr -> Coin -> TxOut era
mkCoinTxOut Addr
tomAddr Coin
tomCoin1) -- Not distrubuted, see tomAddr
            , (TxIn
tomTxIn2, forall era. EraTxOut era => Addr -> Coin -> TxOut era
mkCoinTxOut Addr
tomAddr Coin
tomCoin2) -- Not distributed, see tomAddr
            , (TxIn
annTxIn, forall era. EraTxOut era => Addr -> Coin -> TxOut era
mkCoinTxOut Addr
annAddr Coin
annCoin)
            , (TxIn
ronTxIn, forall era. EraTxOut era => Addr -> Coin -> TxOut era
mkCoinTxOut Addr
ronAddr Coin
ronCoin)
            , (TxIn
johnTxIn, forall era. EraTxOut era => Addr -> Coin -> TxOut era
mkCoinTxOut Addr
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
incrementalStake = forall era.
EraTxOut era =>
PParams era
-> IncrementalStake -> UTxO era -> UTxO era -> IncrementalStake
updateStakeDistribution PParams era
pparams (Map (Credential 'Staking) (CompactForm Coin)
-> Map Ptr (CompactForm Coin) -> IncrementalStake
IStake forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty) forall a. Monoid a => a
mempty UTxO era
utxo1
    umap :: UMap
umap = Map (Credential 'Staking) RDPair
-> Map Ptr (Credential 'Staking)
-> Map (Credential 'Staking) (KeyHash 'StakePool)
-> Map (Credential 'Staking) DRep
-> UMap
UM.unify Map (Credential 'Staking) RDPair
rewards forall k a. Map k a
Map.empty Map (Credential 'Staking) (KeyHash 'StakePool)
delegations forall k a. Map k a
Map.empty
    poolparamMap :: Map (KeyHash 'StakePool) PoolParams
poolparamMap = forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(KeyHash 'StakePool
pool1, PoolParams
pool1Params), (KeyHash 'StakePool
pool2, PoolParams
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
dsUnifiedL UMap
umap
  PState era
pState <- forall a b. Arbitrary a => Lens' a b -> b -> Gen a
arbitraryLens forall era.
Lens' (PState era) (Map (KeyHash 'StakePool) PoolParams)
psStakePoolParamsL Map (KeyHash 'StakePool) PoolParams
poolparamMap
  let computedStakeDistr :: Map (Credential 'Staking) 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 (Stake -> VMap VB VP (Credential 'Staking) (CompactForm Coin)
unStake (SnapShot -> Stake
ssStake SnapShot
snap)))
        where
          snap :: SnapShot
snap =
            forall era.
EraPParams era =>
PParams era
-> IncrementalStake -> DState era -> PState era -> SnapShot
incrementalStakeDistr -- This computes the actual Incremental Stake
              PParams era
pparams
              IncrementalStake
incrementalStake
              DState era
dState
              PState era
pState

      expectedStakeDistr :: Map (Credential 'Staking) Coin
      expectedStakeDistr :: Map (Credential 'Staking) Coin
expectedStakeDistr =
        forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList -- Coin Part is (rdRewardCoin <> utxoCoin)
          [ (Credential 'Staking
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
ann, RDPair -> Coin
rdRewardCoin RDPair
annRD forall a. Semigroup a => a -> a -> a
<> Coin
annCoin)
          , (Credential 'Staking
ron, RDPair -> Coin
rdRewardCoin RDPair
ronRD forall a. Semigroup a => a -> a -> a
<> Coin
ronCoin)
          , (Credential 'Staking
john, RDPair -> Coin
rdRewardCoin RDPair
johnRD forall a. Semigroup a => a -> a -> a
<> Coin
johnCoin)
          , (Credential 'Staking
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 (Map (Credential 'Staking) Coin
computedStakeDistr forall a. (Eq a, Show a) => a -> a -> Property
=== Map (Credential 'Staking) 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)