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

-- | 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.Keys (KeyHash, KeyRole (StakePool))
import Cardano.Ledger.PoolParams (PoolParams (ppVrf))
import Cardano.Ledger.Shelley.Rules (calculatePoolDistr)
import Cardano.Ledger.State (
  IndividualPoolStake (..),
  PoolDistr (..),
  SnapShot (..),
  Stake (..),
  sumAllStake,
 )
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
snap ->
                forall prop. Testable prop => TestName -> prop -> Property
counterexample
                  TestName
"BAD"
                  ((KeyHash 'StakePool -> Bool) -> SnapShot -> PoolDistr
oldCalculatePoolDistr (forall a b. a -> b -> a
const Bool
True) SnapShot
snap 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 Stake
stake VMap VB VB (Credential 'Staking) (KeyHash 'StakePool)
delegs VMap VB VB (KeyHash 'StakePool) PoolParams
poolParams) =
  let Coin Integer
totalc = Stake -> Coin
sumAllStake Stake
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) (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
d, (CompactForm Coin
compactCoin, Integer
c forall a. Integral a => a -> a -> Ratio a
% Integer
nonZeroTotal))
          | (Credential 'Staking
hk, CompactForm Coin
compactCoin) <- forall (kv :: * -> *) k (vv :: * -> *) v.
(Vector kv k, Vector vv v) =>
VMap kv vv k v -> [(k, v)]
VMap.toAscList (Stake -> VMap VB VP (Credential 'Staking) (CompactForm Coin)
unStake Stake
stake)
          , let Coin Integer
c = forall a. Compactible a => CompactForm a -> a
fromCompact CompactForm Coin
compactCoin
          , Just KeyHash 'StakePool
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
hk VMap VB VB (Credential 'Staking) (KeyHash 'StakePool)
delegs]
          , KeyHash 'StakePool -> Bool
includeHash KeyHash 'StakePool
d
          ]
   in Map (KeyHash 'StakePool) IndividualPoolStake
-> CompactForm Coin -> PoolDistr
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) 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
            (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 PoolParams -> VRFVerKeyHash 'StakePoolVRF
ppVrf VMap VB VB (KeyHash 'StakePool) PoolParams
poolParams))
        )
        (Word64 -> CompactForm Coin
CompactCoin forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
nonZeroTotal)