{-# 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 qualified Cardano.Ledger.Shelley.Rules as Shelley
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
Shelley.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 VS (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 VS (Credential Staking) StakeWithDelegation
 -> Map (Credential Staking) StakeWithDelegation)
-> VMap VB VS (Credential Staking) StakeWithDelegation
-> Map (Credential Staking) StakeWithDelegation
forall a b. (a -> b) -> a -> b
$ ActiveStake -> VMap VB VS (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