{-# 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 (unNonZero)
import Cardano.Ledger.Coin (Coin (..))
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 (..),
  StakePoolSnapShot (spssVrf),
  StakeWithDelegation (..),
  sumAllActiveStake,
  unActiveStake,
 )
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
=== 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 ActiveStake
activeStake NonZero Coin
_ VMap VB VB (KeyHash StakePool) StakePoolSnapShot
stakePoolsSnapShot) =
  let nonZeroTotalActiveStake :: NonZero Coin
nonZeroTotalActiveStake = ActiveStake -> NonZero Coin
sumAllActiveStake ActiveStake
activeStake
      activeStakeMap :: Map (Credential Staking) StakeWithDelegation
activeStakeMap = VMap VB VB (Credential Staking) StakeWithDelegation
-> Map (Credential Staking) StakeWithDelegation
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) StakeWithDelegation
 -> Map (Credential Staking) StakeWithDelegation)
-> VMap VB VB (Credential Staking) StakeWithDelegation
-> Map (Credential Staking) StakeWithDelegation
forall a b. (a -> b) -> a -> b
$ ActiveStake -> VMap VB VB (Credential Staking) StakeWithDelegation
unActiveStake ActiveStake
activeStake
      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)))
          | StakeWithDelegation NonZero (CompactForm Coin)
nzc KeyHash StakePool
d <- Map (Credential Staking) StakeWithDelegation
-> [StakeWithDelegation]
forall k a. Map k a -> [a]
Map.elems Map (Credential Staking) StakeWithDelegation
activeStakeMap
          , let compactCoin :: CompactForm Coin
compactCoin = NonZero (CompactForm Coin) -> CompactForm Coin
forall a. NonZero a -> a
unNonZero NonZero (CompactForm Coin)
nzc
                Coin Integer
c = CompactForm Coin -> Coin
forall a. Compactible a => CompactForm a -> a
fromCompact CompactForm Coin
compactCoin
          , 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 ((StakePoolSnapShot -> VRFVerKeyHash StakePoolVRF)
-> VMap VB VB (KeyHash StakePool) StakePoolSnapShot
-> VMap VB VB (KeyHash StakePool) (VRFVerKeyHash StakePoolVRF)
forall (kv :: * -> *) k (vv :: * -> *) a b.
(Vector kv k, Vector vv a, Vector vv b) =>
(a -> b) -> VMap kv vv k a -> VMap kv vv k b
VMap.map StakePoolSnapShot -> VRFVerKeyHash StakePoolVRF
spssVrf VMap VB VB (KeyHash StakePool) StakePoolSnapShot
stakePoolsSnapShot))
        )
        NonZero Coin
nonZeroTotalActiveStake