{-# 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 = 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})

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 = (Lens' a b
l forall s t a b. ASetter s t a b -> b -> s -> t
.~ b
b) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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 <- 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 = forall p s.
(MakeCredential p 'Payment, MakeStakeReference s) =>
p -> s -> Addr
mkAddr Credential 'Payment
tomPay StakeReference
StakeRefNull
      johnAddr :: Addr
johnAddr = forall p s.
(MakeCredential p 'Payment, MakeStakeReference s) =>
p -> s -> Addr
mkAddr Credential 'Payment
johnPay Credential 'Staking
john
      annAddr :: Addr
annAddr = forall p s.
(MakeCredential p 'Payment, MakeStakeReference s) =>
p -> s -> Addr
mkAddr Credential 'Payment
annPay Credential 'Staking
ann
      ronAddr :: Addr
ronAddr = 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 =
        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
            ]
        )

    instantStake :: InstantStake era
instantStake = forall era.
EraStake era =>
UTxO era -> InstantStake era -> InstantStake era
addInstantStake UTxO era
utxo1 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 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 snapShot :: SnapShot
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 = 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
snapShot)))

      expectedStakeDistr :: Map (Credential 'Staking) Coin
      expectedStakeDistr :: Map (Credential 'Staking) Coin
expectedStakeDistr =
        forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList forall a b. (a -> b) -> a -> b
$
          [ (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)
          ]
            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 forall a. Eq a => a -> a -> Bool
/= Integer -> Coin
Coin Integer
0 -- We need to filter out zero rewards from instant stake
               ]

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