{-# 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.Shelley.Rules (calculatePoolDistr)
import Cardano.Ledger.State (
  IndividualPoolStake (..),
  PoolDistr (..),
  PoolParams (ppVrf),
  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"
    [ TestName -> Property -> TestTree
forall a. Testable a => TestName -> a -> TestTree
testProperty
        TestName
"old==new"
        ( Int -> (SnapShot -> Property) -> Property
forall prop. Testable prop => Int -> prop -> Property
withMaxSuccess
            Int
500
            ( \SnapShot
snap ->
                TestName -> Property -> Property
forall prop. Testable prop => TestName -> prop -> Property
counterexample
                  TestName
"BAD"
                  ((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 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 Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 then Integer
1 else Integer
totalc
      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
% Integer
nonZeroTotal))
          | (Credential 'Staking
hk, CompactForm Coin
compactCoin) <- VMap VB VP (Credential 'Staking) (CompactForm Coin)
-> [(Credential 'Staking, CompactForm Coin)]
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 = 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
-> CompactForm 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 ((PoolParams -> VRFVerKeyHash 'StakePoolVRF)
-> VMap VB VB (KeyHash 'StakePool) PoolParams
-> 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 PoolParams -> VRFVerKeyHash 'StakePoolVRF
ppVrf VMap VB VB (KeyHash 'StakePool) PoolParams
poolParams))
        )
        (Word64 -> CompactForm Coin
CompactCoin (Word64 -> CompactForm Coin) -> Word64 -> CompactForm Coin
forall a b. (a -> b) -> a -> b
$ Integer -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
nonZeroTotal)