{-# 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.Coin (Coin (Coin), CompactForm (CompactCoin))
import Cardano.Ledger.Compactible (fromCompact)
import Cardano.Ledger.Crypto (StandardCrypto)
import Cardano.Ledger.EpochBoundary (SnapShot (..), Stake (..), sumAllStake)
import Cardano.Ledger.Keys (KeyHash, KeyRole (StakePool))
import Cardano.Ledger.PoolDistr (IndividualPoolStake (..), PoolDistr (..))
import Cardano.Ledger.PoolParams (PoolParams (ppVrf))
import Cardano.Ledger.Shelley.Rules (calculatePoolDistr)
import qualified Data.Map.Strict as Map
import Data.Ratio ((%))
import qualified Data.VMap as VMap
import Test.Cardano.Ledger.Shelley.Arbitrary ()
import Test.Cardano.Ledger.Shelley.Serialisation.EraIndepGenerators ()
import Test.Tasty
import Test.Tasty.QuickCheck

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

calcPoolDistOldEqualsNew :: TestTree
calcPoolDistOldEqualsNew :: TestTree
calcPoolDistOldEqualsNew =
  TestName -> [TestTree] -> TestTree
testGroup
    TestName
"calculatePoolDistr"
    [ forall a. Testable a => TestName -> a -> TestTree
testProperty
        TestName
"old==new"
        ( forall prop. Testable prop => Int -> prop -> Property
withMaxSuccess
            Int
500
            ( \SnapShot StandardCrypto
snap ->
                forall prop. Testable prop => TestName -> prop -> Property
counterexample
                  TestName
"BAD"
                  (forall c.
(KeyHash 'StakePool c -> Bool) -> SnapShot c -> PoolDistr c
oldCalculatePoolDistr @StandardCrypto (forall a b. a -> b -> a
const Bool
True) SnapShot StandardCrypto
snap forall a. (Eq a, Show a) => a -> a -> Property
=== forall c. SnapShot c -> PoolDistr c
calculatePoolDistr SnapShot StandardCrypto
snap)
            )
        )
    ]

-- | The original version of calculatePoolDistr
oldCalculatePoolDistr :: forall c. (KeyHash 'StakePool c -> Bool) -> SnapShot c -> PoolDistr c
oldCalculatePoolDistr :: forall c.
(KeyHash 'StakePool c -> Bool) -> SnapShot c -> PoolDistr c
oldCalculatePoolDistr KeyHash 'StakePool c -> Bool
includeHash (SnapShot Stake c
stake VMap VB VB (Credential 'Staking c) (KeyHash 'StakePool c)
delegs VMap VB VB (KeyHash 'StakePool c) (PoolParams c)
poolParams) =
  let Coin Integer
totalc = forall c. Stake c -> Coin
sumAllStake Stake c
stake
      -- totalc could be zero (in particular when shrinking)
      nonZeroTotal :: Integer
nonZeroTotal = if Integer
totalc forall a. Eq a => a -> a -> Bool
== Integer
0 then Integer
1 else Integer
totalc
      sd :: Map (KeyHash 'StakePool c) (CompactForm Coin, Ratio Integer)
sd =
        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 forall a. Semigroup a => a -> a -> a
<> CompactForm Coin
cc', Ratio Integer
rat forall a. Num a => a -> a -> a
+ Ratio Integer
rat')) forall a b. (a -> b) -> a -> b
$
          [ (KeyHash 'StakePool c
d, (CompactForm Coin
compactCoin, Integer
c forall a. Integral a => a -> a -> Ratio a
% Integer
nonZeroTotal))
          | (Credential 'Staking c
hk, CompactForm Coin
compactCoin) <- forall (kv :: * -> *) k (vv :: * -> *) v.
(Vector kv k, Vector vv v) =>
VMap kv vv k v -> [(k, v)]
VMap.toAscList (forall c.
Stake c -> VMap VB VP (Credential 'Staking c) (CompactForm Coin)
unStake Stake c
stake)
          , let Coin Integer
c = forall a. Compactible a => CompactForm a -> a
fromCompact CompactForm Coin
compactCoin
          , Just KeyHash 'StakePool c
d <- [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 c
hk VMap VB VB (Credential 'Staking c) (KeyHash 'StakePool c)
delegs]
          , KeyHash 'StakePool c -> Bool
includeHash KeyHash 'StakePool c
d
          ]
   in forall c.
Map (KeyHash 'StakePool c) (IndividualPoolStake c)
-> CompactForm Coin -> PoolDistr c
PoolDistr
        ( 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) Hash (HASH c) (VerKeyVRF (VRF c))
vrf -> forall c.
Ratio Integer
-> CompactForm Coin
-> Hash c (VerKeyVRF c)
-> IndividualPoolStake c
IndividualPoolStake Ratio Integer
rat CompactForm Coin
cc Hash (HASH c) (VerKeyVRF (VRF c))
vrf)
            Map (KeyHash 'StakePool c) (CompactForm Coin, Ratio Integer)
sd
            (forall (kv :: * -> *) k (vv :: * -> *) v.
(Vector kv k, Vector vv v) =>
VMap kv vv k v -> Map k v
VMap.toMap (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 forall c. PoolParams c -> Hash c (VerKeyVRF c)
ppVrf VMap VB VB (KeyHash 'StakePool c) (PoolParams c)
poolParams))
        )
        (Word64 -> CompactForm Coin
CompactCoin forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
nonZeroTotal)