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

-- | Test that modifications to the calculatePoolDistr function
--   made when building the Tickf benchmarks behave the same as
--   the code that was replaced.
module Test.Cardano.Ledger.Tickf (oldCalculatePoolDistr, calcPoolDistOldEqualsNew) where

import Cardano.Ledger.BaseTypes (nonZeroOr, unNonZero)
import Cardano.Ledger.Coin (Coin (..), knownNonZeroCoin)
import Cardano.Ledger.Compactible (fromCompact)
import Cardano.Ledger.Keys (KeyHash, KeyRole (StakePool))
import Cardano.Ledger.Shelley.Rules (calculatePoolDistr)
import Cardano.Ledger.State (
  IndividualPoolStake (..),
  PoolDistr (..),
  SnapShot (..),
  Stake (..),
  StakePoolParams (sppVrf),
  sumAllStake,
 )
import qualified Data.Map.Strict as Map
import Data.Ratio ((%))
import qualified Data.VMap as VMap
import Test.Cardano.Ledger.Common
import Test.Cardano.Ledger.Shelley.Arbitrary ()
import Test.Cardano.Ledger.Shelley.Serialisation.EraIndepGenerators ()

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

calcPoolDistOldEqualsNew :: Spec
calcPoolDistOldEqualsNew :: Spec
calcPoolDistOldEqualsNew =
  String -> Spec -> Spec
forall a. HasCallStack => String -> SpecWith a -> SpecWith a
describe String
"calculatePoolDistr" (Spec -> Spec) -> Spec -> Spec
forall a b. (a -> b) -> a -> b
$ do
    String -> Property -> Spec
forall prop.
(HasCallStack, Testable prop) =>
String -> prop -> Spec
prop String
"old==new" (Property -> Spec) -> Property -> Spec
forall a b. (a -> b) -> a -> b
$
      Int -> (SnapShot -> Property) -> Property
forall prop. Testable prop => Int -> prop -> Property
withMaxSuccess Int
500 ((SnapShot -> Property) -> Property)
-> (SnapShot -> Property) -> Property
forall a b. (a -> b) -> a -> b
$ \SnapShot
snap ->
        String -> Property -> Property
forall prop. Testable prop => String -> prop -> Property
counterexample String
"BAD" (Property -> Property) -> Property -> Property
forall a b. (a -> b) -> a -> b
$
          (KeyHash StakePool -> Bool) -> SnapShot -> PoolDistr
oldCalculatePoolDistr (Bool -> KeyHash StakePool -> Bool
forall a b. a -> b -> a
const Bool
True) SnapShot
snap PoolDistr -> PoolDistr -> Property
forall a. (Eq a, Show a) => a -> a -> Property
=== HasCallStack => SnapShot -> PoolDistr
SnapShot -> PoolDistr
calculatePoolDistr SnapShot
snap

-- | The original version of calculatePoolDistr
oldCalculatePoolDistr :: (KeyHash StakePool -> Bool) -> SnapShot -> PoolDistr
oldCalculatePoolDistr :: (KeyHash StakePool -> Bool) -> SnapShot -> PoolDistr
oldCalculatePoolDistr KeyHash StakePool -> Bool
includeHash (SnapShot Stake
stake NonZero Coin
_ VMap VB VB (Credential Staking) (KeyHash StakePool)
delegs VMap VB VB (KeyHash StakePool) StakePoolParams
stakePoolParams VMap VB VB (KeyHash StakePool) StakePoolSnapShot
_) =
  let totalActiveStake :: Coin
totalActiveStake = Stake -> Coin
sumAllStake Stake
stake
      nonZeroTotalActiveStake :: NonZero Coin
nonZeroTotalActiveStake = Coin
totalActiveStake Coin -> NonZero Coin -> NonZero Coin
forall a. HasZero a => a -> NonZero a -> NonZero a
`nonZeroOr` forall (n :: Natural). (KnownNat n, 1 <= n) => NonZero Coin
knownNonZeroCoin @1
      withZeroStake :: Map (Credential Staking) (CompactForm Coin)
withZeroStake = 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 Stake
stake) Map (Credential Staking) (CompactForm Coin)
-> Map (Credential Staking) (CompactForm Coin)
-> Map (Credential Staking) (CompactForm Coin)
forall k a. Ord k => Map k a -> Map k a -> Map k a
`Map.union` (CompactForm Coin
forall a. Monoid a => a
mempty CompactForm Coin
-> Map (Credential Staking) (KeyHash StakePool)
-> Map (Credential Staking) (CompactForm Coin)
forall a b.
a -> Map (Credential Staking) b -> Map (Credential Staking) a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ VMap VB VB (Credential Staking) (KeyHash StakePool)
-> Map (Credential Staking) (KeyHash StakePool)
forall (kv :: * -> *) k (vv :: * -> *) v.
(Vector kv k, Vector vv v) =>
VMap kv vv k v -> Map k v
VMap.toMap VMap VB VB (Credential Staking) (KeyHash StakePool)
delegs)
      sd :: Map (KeyHash StakePool) (CompactForm Coin, Ratio Integer)
sd =
        ((CompactForm Coin, Ratio Integer)
 -> (CompactForm Coin, Ratio Integer)
 -> (CompactForm Coin, Ratio Integer))
-> [(KeyHash StakePool, (CompactForm Coin, Ratio Integer))]
-> Map (KeyHash StakePool) (CompactForm Coin, Ratio Integer)
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith (\(CompactForm Coin
cc, Ratio Integer
rat) (CompactForm Coin
cc', Ratio Integer
rat') -> (CompactForm Coin
cc CompactForm Coin -> CompactForm Coin -> CompactForm Coin
forall a. Semigroup a => a -> a -> a
<> CompactForm Coin
cc', Ratio Integer
rat Ratio Integer -> Ratio Integer -> Ratio Integer
forall a. Num a => a -> a -> a
+ Ratio Integer
rat')) ([(KeyHash StakePool, (CompactForm Coin, Ratio Integer))]
 -> Map (KeyHash StakePool) (CompactForm Coin, Ratio Integer))
-> [(KeyHash StakePool, (CompactForm Coin, Ratio Integer))]
-> Map (KeyHash StakePool) (CompactForm Coin, Ratio Integer)
forall a b. (a -> b) -> a -> b
$
          [ (KeyHash StakePool
d, (CompactForm Coin
compactCoin, Integer
c Integer -> Integer -> Ratio Integer
forall a. Integral a => a -> a -> Ratio a
% Coin -> Integer
unCoin (NonZero Coin -> Coin
forall a. NonZero a -> a
unNonZero NonZero Coin
nonZeroTotalActiveStake)))
          | (Credential Staking
hk, CompactForm Coin
compactCoin) <- Map (Credential Staking) (CompactForm Coin)
-> [(Credential Staking, CompactForm Coin)]
forall k a. Map k a -> [(k, a)]
Map.toAscList Map (Credential Staking) (CompactForm Coin)
withZeroStake
          , let Coin Integer
c = CompactForm Coin -> Coin
forall a. Compactible a => CompactForm a -> a
fromCompact CompactForm Coin
compactCoin
          , Just KeyHash StakePool
d <- [Credential Staking
-> VMap VB VB (Credential Staking) (KeyHash StakePool)
-> Maybe (KeyHash StakePool)
forall k (kv :: * -> *) (vv :: * -> *) v.
(Ord k, Vector kv k, Vector vv v) =>
k -> VMap kv vv k v -> Maybe v
VMap.lookup Credential Staking
hk VMap VB VB (Credential Staking) (KeyHash StakePool)
delegs]
          , KeyHash StakePool -> Bool
includeHash KeyHash StakePool
d
          ]
   in Map (KeyHash StakePool) IndividualPoolStake
-> NonZero Coin -> PoolDistr
PoolDistr
        ( ((CompactForm Coin, Ratio Integer)
 -> VRFVerKeyHash StakePoolVRF -> IndividualPoolStake)
-> Map (KeyHash StakePool) (CompactForm Coin, Ratio Integer)
-> Map (KeyHash StakePool) (VRFVerKeyHash StakePoolVRF)
-> Map (KeyHash StakePool) IndividualPoolStake
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
Map.intersectionWith
            (\(CompactForm Coin
cc, Ratio Integer
rat) VRFVerKeyHash StakePoolVRF
vrf -> Ratio Integer
-> CompactForm Coin
-> VRFVerKeyHash StakePoolVRF
-> IndividualPoolStake
IndividualPoolStake Ratio Integer
rat CompactForm Coin
cc VRFVerKeyHash StakePoolVRF
vrf)
            Map (KeyHash StakePool) (CompactForm Coin, Ratio Integer)
sd
            (VMap VB VB (KeyHash StakePool) (VRFVerKeyHash StakePoolVRF)
-> Map (KeyHash StakePool) (VRFVerKeyHash StakePoolVRF)
forall (kv :: * -> *) k (vv :: * -> *) v.
(Vector kv k, Vector vv v) =>
VMap kv vv k v -> Map k v
VMap.toMap ((StakePoolParams -> VRFVerKeyHash StakePoolVRF)
-> VMap VB VB (KeyHash StakePool) StakePoolParams
-> VMap VB VB (KeyHash StakePool) (VRFVerKeyHash StakePoolVRF)
forall (vv :: * -> *) a b (kv :: * -> *) k.
(Vector vv a, Vector vv b) =>
(a -> b) -> VMap kv vv k a -> VMap kv vv k b
VMap.map StakePoolParams -> VRFVerKeyHash StakePoolVRF
sppVrf VMap VB VB (KeyHash StakePool) StakePoolParams
stakePoolParams))
        )
        NonZero Coin
nonZeroTotalActiveStake