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

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

import Cardano.Ledger.Coin (Coin (..))
import Cardano.Ledger.Compactible (fromCompact)
import Cardano.Ledger.Core (EraTxOut, mkCoinTxOut)
import Cardano.Ledger.Credential (Credential, StakeReference (..))
import Cardano.Ledger.Keys (KeyHash, KeyRole (..))
import Cardano.Ledger.PoolParams (PoolParams (..))
import Cardano.Ledger.State hiding (delegations, rewards)
import Cardano.Ledger.TxIn (TxIn)
import Cardano.Ledger.UMap (RDPair (..), rdRewardCoin)
import qualified Cardano.Ledger.UMap as UM
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 ()
import Test.Cardano.Ledger.Core.KeyPair (mkAddr)

ppIdL :: Lens' PoolParams (KeyHash 'StakePool)
ppIdL :: Lens' PoolParams (KeyHash 'StakePool)
ppIdL = (PoolParams -> KeyHash 'StakePool)
-> (PoolParams -> KeyHash 'StakePool -> PoolParams)
-> Lens' PoolParams (KeyHash 'StakePool)
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 = y})

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. (EraTxOut era, EraStake era) => Gen Property
instantStakeIncludesRewards :: forall era. (EraTxOut era, EraStake era) => Gen Property
instantStakeIncludesRewards = 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 <- Lens' PoolParams (KeyHash 'StakePool)
-> KeyHash 'StakePool -> Gen PoolParams
forall a b. Arbitrary a => Lens' a b -> b -> Gen a
arbitraryLens (KeyHash 'StakePool -> f (KeyHash 'StakePool))
-> PoolParams -> f PoolParams
Lens' PoolParams (KeyHash 'StakePool)
ppIdL KeyHash 'StakePool
pool1
  PoolParams
pool2Params <- Lens' PoolParams (KeyHash 'StakePool)
-> KeyHash 'StakePool -> Gen PoolParams
forall a b. Arbitrary a => Lens' a b -> b -> Gen a
arbitraryLens (KeyHash 'StakePool -> f (KeyHash 'StakePool))
-> PoolParams -> f PoolParams
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 -> StakeReference -> Addr
forall p s.
(MakeCredential p 'Payment, MakeStakeReference s) =>
p -> s -> Addr
mkAddr Credential 'Payment
tomPay StakeReference
StakeRefNull
      johnAddr :: Addr
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 :: Addr
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 :: Addr
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

      rewards :: Map (Credential 'Staking) RDPair
      rewards :: Map (Credential 'Staking) RDPair
rewards =
        [(Credential 'Staking, RDPair)] -> Map (Credential 'Staking) RDPair
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 =
        [(Credential 'Staking, KeyHash 'StakePool)]
-> Map (Credential 'Staking) (KeyHash 'StakePool)
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
        ( [(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
tomCoin1) -- Not distrubuted, see tomAddr
            , (TxIn
tomTxIn2, Addr -> Coin -> TxOut era
forall era. EraTxOut era => Addr -> Coin -> TxOut era
mkCoinTxOut Addr
tomAddr Coin
tomCoin2) -- Not distributed, see tomAddr
            , (TxIn
annTxIn, Addr -> Coin -> TxOut era
forall era. EraTxOut era => Addr -> Coin -> TxOut era
mkCoinTxOut Addr
annAddr Coin
annCoin)
            , (TxIn
ronTxIn, Addr -> Coin -> TxOut era
forall era. EraTxOut era => Addr -> Coin -> TxOut era
mkCoinTxOut Addr
ronAddr Coin
ronCoin)
            , (TxIn
johnTxIn, Addr -> Coin -> TxOut era
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
            ]
        )

    instantStake :: InstantStake era
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
    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 Map Ptr (Credential 'Staking)
forall k a. Map k a
Map.empty Map (Credential 'Staking) (KeyHash 'StakePool)
delegations Map (Credential 'Staking) DRep
forall k a. Map k a
Map.empty
    poolparamMap :: Map (KeyHash 'StakePool) PoolParams
poolparamMap = [(KeyHash 'StakePool, PoolParams)]
-> Map (KeyHash 'StakePool) PoolParams
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 <- Lens' (DState era) UMap -> UMap -> Gen (DState era)
forall a b. Arbitrary a => Lens' a b -> b -> Gen a
arbitraryLens (UMap -> f UMap) -> DState era -> f (DState era)
forall era (f :: * -> *).
Functor f =>
(UMap -> f UMap) -> DState era -> f (DState era)
Lens' (DState era) UMap
dsUnifiedL UMap
umap
  PState era
pState <- Lens' (PState era) (Map (KeyHash 'StakePool) PoolParams)
-> Map (KeyHash 'StakePool) PoolParams -> Gen (PState era)
forall a b. Arbitrary a => Lens' a b -> b -> Gen a
arbitraryLens (Map (KeyHash 'StakePool) PoolParams
 -> f (Map (KeyHash 'StakePool) PoolParams))
-> PState era -> f (PState era)
forall era (f :: * -> *).
Functor f =>
(Map (KeyHash 'StakePool) PoolParams
 -> f (Map (KeyHash 'StakePool) PoolParams))
-> PState era -> f (PState era)
Lens' (PState era) (Map (KeyHash 'StakePool) PoolParams)
psStakePoolParamsL Map (KeyHash 'StakePool) PoolParams
poolparamMap
  let snapShot :: SnapShot
snapShot = InstantStake era -> DState era -> PState era -> SnapShot
forall era.
EraStake era =>
InstantStake era -> DState era -> PState era -> SnapShot
snapShotFromInstantStake InstantStake era
instantStake DState era
dState PState era
pState
      computedStakeDistr :: Map (Credential 'Staking) Coin
computedStakeDistr = (CompactForm Coin -> Coin)
-> Map (Credential 'Staking) (CompactForm Coin)
-> Map (Credential 'Staking) Coin
forall a b k. (a -> b) -> Map k a -> Map k b
Map.map CompactForm Coin -> Coin
forall a. Compactible a => CompactForm a -> a
fromCompact (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) Coin
      expectedStakeDistr :: Map (Credential 'Staking) Coin
expectedStakeDistr =
        [(Credential 'Staking, Coin)] -> Map (Credential 'Staking) Coin
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(Credential 'Staking, Coin)] -> Map (Credential 'Staking) Coin)
-> [(Credential 'Staking, Coin)] -> Map (Credential 'Staking) Coin
forall a b. (a -> b) -> a -> b
$
          [ (Credential 'Staking
ann, RDPair -> Coin
rdRewardCoin RDPair
annRD Coin -> Coin -> Coin
forall a. Semigroup a => a -> a -> a
<> Coin
annCoin)
          , (Credential 'Staking
ron, RDPair -> Coin
rdRewardCoin RDPair
ronRD Coin -> Coin -> Coin
forall a. Semigroup a => a -> a -> a
<> Coin
ronCoin)
          , (Credential 'Staking
john, RDPair -> Coin
rdRewardCoin RDPair
johnRD Coin -> Coin -> Coin
forall a. Semigroup a => a -> a -> a
<> Coin
johnCoin)
          ]
            [(Credential 'Staking, Coin)]
-> [(Credential 'Staking, Coin)] -> [(Credential 'Staking, Coin)]
forall a. [a] -> [a] -> [a]
++ [ (Credential 'Staking, Coin)
stake
               | stake :: (Credential 'Staking, Coin)
stake@(Credential 'Staking
_, Coin
reward) <-
                   [ (Credential 'Staking
tom, RDPair -> Coin
rdRewardCoin RDPair
tomRD) -- tom uxtxoCoin cab be zero because his address has StakeRefNull
                   , (Credential 'Staking
mary, RDPair -> Coin
rdRewardCoin RDPair
maryRD) -- mary uxtxoCoin can be zero because she has no UtxO entry
                   ]
               , Coin
reward Coin -> Coin -> Bool
forall a. Eq a => a -> a -> Bool
/= Integer -> Coin
Coin Integer
0 -- We need to filter out zero rewards from instant stake
               ]

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

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