{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}

module Test.Cardano.Ledger.Shelley.Rules.IncrementalStake (
  incrStakeComputationTest,
  incrStakeComparisonTest,
  stakeDistr,
  aggregateUtxoCoinByCredential,
) where

import Test.Cardano.Ledger.Shelley.Rules.TestChain (
  TestingLedger,
  forAllChainTrace,
  ledgerTraceFromBlock,
  longTraceLen,
  traceLen,
 )

import Cardano.Ledger.Address (Addr (..))
import Cardano.Ledger.Coin
import Cardano.Ledger.Compactible (fromCompact)
import Cardano.Ledger.Core
import Cardano.Ledger.Credential (Credential (..), Ptr, StakeReference (StakeRefBase, StakeRefPtr))
import Cardano.Ledger.EpochBoundary (SnapShot (..), Stake (..))
import Cardano.Ledger.Shelley.Core
import Cardano.Ledger.Shelley.LedgerState (
  CertState (..),
  DState (..),
  EpochState (..),
  IncrementalStake (..),
  LedgerState (..),
  NewEpochState (..),
  PState (..),
  UTxOState (..),
  credMap,
  curPParamsEpochStateL,
  incrementalStakeDistr,
  ptrsMap,
 )
import qualified Cardano.Ledger.UMap as UM
import Cardano.Ledger.UTxO (UTxO (..), coinBalance)
import Control.SetAlgebra (dom, eval, (▷), (◁))
import Data.Foldable (fold)
import Data.Map (Map)
import qualified Data.Map.Strict as Map
import Data.Proxy
import qualified Data.VMap as VMap
import Lens.Micro hiding (ix)
import Test.Cardano.Ledger.Shelley.Constants (defaultConstants)
import Test.Cardano.Ledger.Shelley.Generator.Core (GenEnv)
import Test.Cardano.Ledger.Shelley.Generator.EraGen (EraGen (..))
import Test.Cardano.Ledger.Shelley.Generator.ShelleyEraGen ()
import Test.Cardano.Ledger.Shelley.Rules.Chain (CHAIN, ChainState (..))
import Test.Cardano.Ledger.Shelley.Utils (
  ChainProperty,
 )
import Test.Cardano.Ledger.TerseTools (tersemapdiffs)
import Test.Control.State.Transition.Trace (
  SourceSignalTarget (..),
  sourceSignalTargets,
 )
import qualified Test.Control.State.Transition.Trace.Generator.QuickCheck as QC
import Test.QuickCheck (
  Property,
  conjoin,
  counterexample,
  (===),
 )
import Test.Tasty (TestTree)
import Test.Tasty.QuickCheck (testProperty)

incrStakeComputationTest ::
  forall era ledger.
  ( EraGen era
  , TestingLedger era ledger
  , ChainProperty era
  , QC.HasTrace (CHAIN era) (GenEnv era)
  ) =>
  TestTree
incrStakeComputationTest :: forall era ledger.
(EraGen era, TestingLedger era ledger, ChainProperty era,
 HasTrace (CHAIN era) (GenEnv era)) =>
TestTree
incrStakeComputationTest =
  forall a. Testable a => TestName -> a -> TestTree
testProperty TestName
"incremental stake calc" forall a b. (a -> b) -> a -> b
$
    forall era prop.
(Testable prop, EraGen era, HasTrace (CHAIN era) (GenEnv era),
 EraGov era) =>
Word64 -> Constants -> (Trace (CHAIN era) -> prop) -> Property
forAllChainTrace @era Word64
longTraceLen Constants
defaultConstants forall a b. (a -> b) -> a -> b
$ \Trace (CHAIN era)
tr -> do
      let ssts :: [SourceSignalTarget (CHAIN era)]
ssts = forall a. Trace a -> [SourceSignalTarget a]
sourceSignalTargets Trace (CHAIN era)
tr

      forall prop. Testable prop => [prop] -> Property
conjoin forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall a b. (a -> b) -> a -> b
$
        [ -- preservation properties
          forall a b. (a -> b) -> [a] -> [b]
map (forall era ledger.
(EraSegWits era, ChainProperty era, TestingLedger era ledger) =>
SourceSignalTarget (CHAIN era) -> Property
incrStakeComp @era @ledger) [SourceSignalTarget (CHAIN era)]
ssts
        ]

incrStakeComp ::
  forall era ledger.
  (EraSegWits era, ChainProperty era, TestingLedger era ledger) =>
  SourceSignalTarget (CHAIN era) ->
  Property
incrStakeComp :: forall era ledger.
(EraSegWits era, ChainProperty era, TestingLedger era ledger) =>
SourceSignalTarget (CHAIN era) -> Property
incrStakeComp SourceSignalTarget {source :: forall a. SourceSignalTarget a -> State a
source = State (CHAIN era)
chainSt, signal :: forall a. SourceSignalTarget a -> Signal a
signal = Signal (CHAIN era)
block} =
  forall prop. Testable prop => [prop] -> Property
conjoin forall a b. (a -> b) -> a -> b
$
    forall a b. (a -> b) -> [a] -> [b]
map SourceSignalTarget ledger -> Property
checkIncrStakeComp forall a b. (a -> b) -> a -> b
$
      forall a. Trace a -> [SourceSignalTarget a]
sourceSignalTargets Trace ledger
ledgerTr
  where
    (ChainState era
_, Trace ledger
ledgerTr) = forall era ledger.
(ChainProperty era, EraSegWits era, TestingLedger era ledger) =>
ChainState era
-> Block (BHeader MockCrypto) era -> (ChainState era, Trace ledger)
ledgerTraceFromBlock @era @ledger State (CHAIN era)
chainSt Signal (CHAIN era)
block
    checkIncrStakeComp :: SourceSignalTarget ledger -> Property
    checkIncrStakeComp :: SourceSignalTarget ledger -> Property
checkIncrStakeComp
      SourceSignalTarget
        { source :: forall a. SourceSignalTarget a -> State a
source = LedgerState UTxOState {utxosUtxo :: forall era. UTxOState era -> UTxO era
utxosUtxo = UTxO era
u, utxosStakeDistr :: forall era. UTxOState era -> IncrementalStake
utxosStakeDistr = IncrementalStake
sd} CertState era
dp
        , signal :: forall a. SourceSignalTarget a -> Signal a
signal = Signal ledger
tx
        , target :: forall a. SourceSignalTarget a -> State a
target = LedgerState UTxOState {utxosUtxo :: forall era. UTxOState era -> UTxO era
utxosUtxo = UTxO era
u', utxosStakeDistr :: forall era. UTxOState era -> IncrementalStake
utxosStakeDistr = IncrementalStake
sd'} CertState era
dp'
        } =
        forall prop. Testable prop => TestName -> prop -> Property
counterexample
          ( forall a. Monoid a => [a] -> a
mconcat
              ( [ TestName
"\nDetails:\n"
                , TestName
"\ntx\n"
                , forall a. Show a => a -> TestName
show Signal ledger
tx
                , TestName
"\nsize original utxo\n"
                , forall a. Show a => a -> TestName
show (forall k a. Map k a -> Int
Map.size forall a b. (a -> b) -> a -> b
$ forall era. UTxO era -> Map TxIn (TxOut era)
unUTxO UTxO era
u)
                , TestName
"\noriginal utxo\n"
                , forall a. Show a => a -> TestName
show UTxO era
u
                , TestName
"\noriginal sd\n"
                , forall a. Show a => a -> TestName
show IncrementalStake
sd
                , TestName
"\nfinal utxo\n"
                , forall a. Show a => a -> TestName
show UTxO era
u'
                , TestName
"\nfinal sd\n"
                , forall a. Show a => a -> TestName
show IncrementalStake
sd'
                , TestName
"\noriginal ptrs\n"
                , forall a. Show a => a -> TestName
show Map Ptr (Credential 'Staking)
ptrs
                , TestName
"\nfinal ptrs\n"
                , forall a. Show a => a -> TestName
show Map Ptr (Credential 'Staking)
ptrs'
                ]
              )
          )
          forall a b. (a -> b) -> a -> b
$ Coin
utxoBal forall a. (Eq a, Show a) => a -> a -> Property
=== forall a. Compactible a => CompactForm a -> a
fromCompact CompactForm Coin
incrStakeBal
        where
          utxoBal :: Coin
utxoBal = forall era. EraTxOut era => UTxO era -> Coin
coinBalance UTxO era
u'
          incrStakeBal :: CompactForm Coin
incrStakeBal = forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold (IncrementalStake -> Map (Credential 'Staking) (CompactForm Coin)
credMap IncrementalStake
sd') forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold (IncrementalStake -> Map Ptr (CompactForm Coin)
ptrMap IncrementalStake
sd')
          ptrs :: Map Ptr (Credential 'Staking)
ptrs = forall era. DState era -> Map Ptr (Credential 'Staking)
ptrsMap forall a b. (a -> b) -> a -> b
$ forall era. CertState era -> DState era
certDState CertState era
dp
          ptrs' :: Map Ptr (Credential 'Staking)
ptrs' = forall era. DState era -> Map Ptr (Credential 'Staking)
ptrsMap forall a b. (a -> b) -> a -> b
$ forall era. CertState era -> DState era
certDState CertState era
dp'

incrStakeComparisonTest ::
  forall era.
  ( EraGen era
  , QC.HasTrace (CHAIN era) (GenEnv era)
  , EraGov era
  ) =>
  Proxy era ->
  TestTree
incrStakeComparisonTest :: forall era.
(EraGen era, HasTrace (CHAIN era) (GenEnv era), EraGov era) =>
Proxy era -> TestTree
incrStakeComparisonTest Proxy era
Proxy =
  forall a. Testable a => TestName -> a -> TestTree
testProperty TestName
"Incremental stake distribution at epoch boundaries agrees" forall a b. (a -> b) -> a -> b
$
    forall era prop.
(Testable prop, EraGen era, HasTrace (CHAIN era) (GenEnv era),
 EraGov era) =>
Word64 -> Constants -> (Trace (CHAIN era) -> prop) -> Property
forAllChainTrace Word64
traceLen Constants
defaultConstants forall a b. (a -> b) -> a -> b
$ \Trace (CHAIN era)
tr ->
      forall prop. Testable prop => [prop] -> Property
conjoin forall a b. (a -> b) -> a -> b
$
        forall a b. (a -> b) -> [a] -> [b]
map (\(SourceSignalTarget State (CHAIN era)
_ State (CHAIN era)
target Signal (CHAIN era)
_) -> forall era.
(EraTxOut era, EraGov era) =>
EpochState era -> Property
checkIncrementalStake @era ((forall era. NewEpochState era -> EpochState era
nesEs forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall era. ChainState era -> NewEpochState era
chainNes) State (CHAIN era)
target)) forall a b. (a -> b) -> a -> b
$
          forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {a} {era}.
(State a ~ ChainState era) =>
SourceSignalTarget a -> Bool
sameEpoch) (forall a. Trace a -> [SourceSignalTarget a]
sourceSignalTargets Trace (CHAIN era)
tr)
  where
    sameEpoch :: SourceSignalTarget a -> Bool
sameEpoch SourceSignalTarget {State a
source :: State a
source :: forall a. SourceSignalTarget a -> State a
source, State a
target :: State a
target :: forall a. SourceSignalTarget a -> State a
target} = forall {era}. ChainState era -> EpochNo
epoch State a
source forall a. Eq a => a -> a -> Bool
== forall {era}. ChainState era -> EpochNo
epoch State a
target
    epoch :: ChainState era -> EpochNo
epoch = forall era. NewEpochState era -> EpochNo
nesEL forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall era. ChainState era -> NewEpochState era
chainNes

checkIncrementalStake ::
  forall era.
  (EraTxOut era, EraGov era) =>
  EpochState era ->
  Property
checkIncrementalStake :: forall era.
(EraTxOut era, EraGov era) =>
EpochState era -> Property
checkIncrementalStake EpochState era
es =
  let
    (LedgerState (UTxOState UTxO era
utxo Coin
_ Coin
_ GovState era
_ IncrementalStake
incStake Coin
_) (CertState VState era
_vstate PState era
pstate DState era
dstate)) = forall era. EpochState era -> LedgerState era
esLState EpochState era
es
    stake :: SnapShot
stake = forall era.
EraTxOut era =>
UTxO era -> DState era -> PState era -> SnapShot
stakeDistr @era UTxO era
utxo DState era
dstate PState era
pstate
    istake :: SnapShot
istake = forall era.
EraPParams era =>
PParams era
-> IncrementalStake -> DState era -> PState era -> SnapShot
incrementalStakeDistr (EpochState era
es forall s a. s -> Getting a s a -> a
^. forall era. EraGov era => Lens' (EpochState era) (PParams era)
curPParamsEpochStateL) IncrementalStake
incStake DState era
dstate PState era
pstate
   in
    forall prop. Testable prop => TestName -> prop -> Property
counterexample
      ( TestName
"\nIncremental stake distribution does not match old style stake distribution"
          forall a. [a] -> [a] -> [a]
++ TestName -> Stake -> Stake -> TestName
tersediffincremental TestName
"differences: Old vs Incremental" (SnapShot -> Stake
ssStake SnapShot
stake) (SnapShot -> Stake
ssStake SnapShot
istake)
      )
      (SnapShot
stake forall a. (Eq a, Show a) => a -> a -> Property
=== SnapShot
istake)

tersediffincremental :: String -> Stake -> Stake -> String
tersediffincremental :: TestName -> Stake -> Stake -> TestName
tersediffincremental TestName
message (Stake VMap VB VP (Credential 'Staking) (CompactForm Coin)
a) (Stake VMap VB VP (Credential 'Staking) (CompactForm Coin)
c) =
  forall a b.
(Terse a, Terse b, Ord a, Eq b) =>
TestName -> Map a b -> Map a b -> TestName
tersemapdiffs (TestName
message forall a. [a] -> [a] -> [a]
++ TestName
" " forall a. [a] -> [a] -> [a]
++ TestName
"hashes") (VMap VB VP (Credential 'Staking) (CompactForm Coin)
-> Map (Credential 'Staking) Coin
mp VMap VB VP (Credential 'Staking) (CompactForm Coin)
a) (VMap VB VP (Credential 'Staking) (CompactForm Coin)
-> Map (Credential 'Staking) Coin
mp VMap VB VP (Credential 'Staking) (CompactForm Coin)
c)
  where
    mp :: VMap VB VP (Credential 'Staking) (CompactForm Coin)
-> Map (Credential 'Staking) Coin
mp = forall a b k. (a -> b) -> Map k a -> Map k b
Map.map forall a. Compactible a => CompactForm a -> a
fromCompact forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (kv :: * -> *) k (vv :: * -> *) v.
(Vector kv k, Vector vv v) =>
VMap kv vv k v -> Map k v
VMap.toMap

-- | Compute the current Stake Distribution. This was called at the Epoch boundary in the Snap Rule.
--   Now it is called in the tests to see that its incremental analog 'incrementalStakeDistr' agrees.
stakeDistr ::
  forall era.
  EraTxOut era =>
  UTxO era ->
  DState era ->
  PState era ->
  SnapShot
stakeDistr :: forall era.
EraTxOut era =>
UTxO era -> DState era -> PState era -> SnapShot
stakeDistr UTxO era
u DState era
ds PState era
ps =
  Stake
-> VMap VB VB (Credential 'Staking) (KeyHash 'StakePool)
-> VMap VB VB (KeyHash 'StakePool) PoolParams
-> SnapShot
SnapShot
    (VMap VB VP (Credential 'Staking) (CompactForm Coin) -> Stake
Stake forall a b. (a -> b) -> a -> b
$ forall (kv :: * -> *) k (vv :: * -> *) v.
(Vector kv k, Vector vv v) =>
Map k v -> VMap kv vv k v
VMap.fromMap (HasCallStack => Coin -> CompactForm Coin
UM.compactCoinOrError forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall s t. Embed s t => Exp t -> s
eval (forall k s (f :: * -> * -> *) v.
(Ord k, HasExp s (f k v)) =>
s -> Exp (Sett k ())
dom Map (Credential 'Staking) (KeyHash 'StakePool)
activeDelegs forall k s1 s2 (f :: * -> * -> *) v.
(Ord k, HasExp s1 (Sett k ()), HasExp s2 (f k v)) =>
s1 -> s2 -> Exp (f k v)
◁ Map (Credential 'Staking) Coin
stakeRelation)))
    (forall (kv :: * -> *) k (vv :: * -> *) v.
(Vector kv k, Vector vv v) =>
Map k v -> VMap kv vv k v
VMap.fromMap Map (Credential 'Staking) (KeyHash 'StakePool)
delegs)
    (forall (kv :: * -> *) k (vv :: * -> *) v.
(Vector kv k, Vector vv v) =>
Map k v -> VMap kv vv k v
VMap.fromMap Map (KeyHash 'StakePool) PoolParams
poolParams)
  where
    rewards' :: Map.Map (Credential 'Staking) Coin
    rewards' :: Map (Credential 'Staking) Coin
rewards' = UMap -> Map (Credential 'Staking) Coin
UM.rewardMap (forall era. DState era -> UMap
dsUnified DState era
ds)
    delegs :: Map.Map (Credential 'Staking) (KeyHash 'StakePool)
    delegs :: Map (Credential 'Staking) (KeyHash 'StakePool)
delegs = UMap -> Map (Credential 'Staking) (KeyHash 'StakePool)
UM.sPoolMap (forall era. DState era -> UMap
dsUnified DState era
ds)
    ptrs' :: Map Ptr (Credential 'Staking)
ptrs' = forall era. DState era -> Map Ptr (Credential 'Staking)
ptrsMap DState era
ds
    PState {psStakePoolParams :: forall era. PState era -> Map (KeyHash 'StakePool) PoolParams
psStakePoolParams = Map (KeyHash 'StakePool) PoolParams
poolParams} = PState era
ps
    stakeRelation :: Map (Credential 'Staking) Coin
    stakeRelation :: Map (Credential 'Staking) Coin
stakeRelation = forall era.
EraTxOut era =>
Map Ptr (Credential 'Staking)
-> UTxO era
-> Map (Credential 'Staking) Coin
-> Map (Credential 'Staking) Coin
aggregateUtxoCoinByCredential Map Ptr (Credential 'Staking)
ptrs' UTxO era
u Map (Credential 'Staking) Coin
rewards'
    activeDelegs :: Map.Map (Credential 'Staking) (KeyHash 'StakePool)
    activeDelegs :: Map (Credential 'Staking) (KeyHash 'StakePool)
activeDelegs = forall s t. Embed s t => Exp t -> s
eval ((forall k s (f :: * -> * -> *) v.
(Ord k, HasExp s (f k v)) =>
s -> Exp (Sett k ())
dom Map (Credential 'Staking) Coin
rewards' forall k s1 s2 (f :: * -> * -> *) v.
(Ord k, HasExp s1 (Sett k ()), HasExp s2 (f k v)) =>
s1 -> s2 -> Exp (f k v)
◁ Map (Credential 'Staking) (KeyHash 'StakePool)
delegs) forall k (g :: * -> * -> *) v s1 (f :: * -> * -> *) s2.
(Ord k, Iter g, Ord v, HasExp s1 (f k v), HasExp s2 (g v ())) =>
s1 -> s2 -> Exp (f k v)
▷ forall k s (f :: * -> * -> *) v.
(Ord k, HasExp s (f k v)) =>
s -> Exp (Sett k ())
dom Map (KeyHash 'StakePool) PoolParams
poolParams)

-- | Sum up all the Coin for each staking Credential. This function has an
--   incremental analog. See 'incrementalAggregateUtxoCoinByCredential'
aggregateUtxoCoinByCredential ::
  forall era.
  EraTxOut era =>
  Map Ptr (Credential 'Staking) ->
  UTxO era ->
  Map (Credential 'Staking) Coin ->
  Map (Credential 'Staking) Coin
aggregateUtxoCoinByCredential :: forall era.
EraTxOut era =>
Map Ptr (Credential 'Staking)
-> UTxO era
-> Map (Credential 'Staking) Coin
-> Map (Credential 'Staking) Coin
aggregateUtxoCoinByCredential Map Ptr (Credential 'Staking)
ptrs (UTxO Map TxIn (TxOut era)
u) Map (Credential 'Staking) Coin
initial =
  forall a b k. (a -> b -> a) -> a -> Map k b -> a
Map.foldl' Map (Credential 'Staking) Coin
-> TxOut era -> Map (Credential 'Staking) Coin
accum Map (Credential 'Staking) Coin
initial Map TxIn (TxOut era)
u
  where
    accum :: Map (Credential 'Staking) Coin
-> TxOut era -> Map (Credential 'Staking) Coin
accum Map (Credential 'Staking) Coin
ans TxOut era
out =
      let c :: Coin
c = TxOut era
out forall s a. s -> Getting a s a -> a
^. forall era. (HasCallStack, EraTxOut era) => Lens' (TxOut era) Coin
coinTxOutL
       in case TxOut era
out forall s a. s -> Getting a s a -> a
^. forall era. EraTxOut era => Lens' (TxOut era) Addr
addrTxOutL of
            Addr Network
_ PaymentCredential
_ (StakeRefPtr Ptr
p)
              | Just Credential 'Staking
cred <- forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Ptr
p Map Ptr (Credential 'Staking)
ptrs -> forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith forall a. Semigroup a => a -> a -> a
(<>) Credential 'Staking
cred Coin
c Map (Credential 'Staking) Coin
ans
            Addr Network
_ PaymentCredential
_ (StakeRefBase Credential 'Staking
hk) -> forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith forall a. Semigroup a => a -> a -> a
(<>) Credential 'Staking
hk Coin
c Map (Credential 'Staking) Coin
ans
            Addr
_other -> Map (Credential 'Staking) Coin
ans