Store number of occurrences of VRF key hashes in CertState

in order to avoid pools retiring after the HardFork from removing the
record of a known (but duplicate) VRF key hash

Author: teodanciu

eras/conway/impl/src/Cardano/Ledger/Conway/Rules/HardFork.hs

@@ -4,6 +4,7 @@
 {-# LANGUAGE EmptyDataDeriving #-}
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeApplications #-}
@@ -16,7 +17,7 @@ module Cardano.Ledger.Conway.Rules.HardFork (
   ConwayHardForkEvent (..),
 ) where
 
-import Cardano.Ledger.BaseTypes (ProtVer (..), ShelleyBase, natVersion)
+import Cardano.Ledger.BaseTypes
 import Cardano.Ledger.Conway.Core
 import Cardano.Ledger.Conway.Era (ConwayEra, ConwayHARDFORK)
 import Cardano.Ledger.Conway.State
@@ -36,9 +37,12 @@ import Control.State.Transition (
   tellEvent,
   transitionRules,
  )
+import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
+import Data.Maybe (fromMaybe)
 import qualified Data.Set as Set
 import Data.Void (Void)
+import Data.Word (Word64)
 import GHC.Generics (Generic)
 import Lens.Micro
 
@@ -102,8 +106,19 @@ updateDRepDelegations certState =
 
 populateVRFKeyHashes :: PState era -> PState era
 populateVRFKeyHashes pState =
-  let allVRFKeyHashes =
-        spsVrf
-          <$> Map.elems (pState ^. psStakePoolsL)
-            <> Map.elems (pState ^. psFutureStakePoolsL)
-   in pState & psVRFKeyHashesL .~ Set.fromList allVRFKeyHashes
+  pState
+    & psVRFKeyHashesL
+      %~ accumulateVRFKeyHashes (pState ^. psStakePoolsL)
+        . accumulateVRFKeyHashes (pState ^. psFutureStakePoolsL)
+  where
+    accumulateVRFKeyHashes ::
+      Map (KeyHash 'StakePool) StakePoolState ->
+      Map (VRFVerKeyHash 'StakePoolVRF) (NonZero Word64) ->
+      Map (VRFVerKeyHash 'StakePoolVRF) (NonZero Word64)
+    accumulateVRFKeyHashes spsMap acc =
+      Map.foldr' (\sps -> addVRFKeyHashOccurrence (sps ^. spsVrfL)) acc spsMap
+    addVRFKeyHashOccurrence vrfKeyHash =
+      Map.insertWith combine vrfKeyHash (knownNonZeroBounded @1)
+      where
+        -- Saturates at maxBound: if (+1) would overflow to 0, keep existing value
+        combine _ oldVal = fromMaybe oldVal $ mapNonZero (+ 1) oldVal

eras/conway/impl/testlib/Test/Cardano/Ledger/Conway/Imp/HardForkSpec.hs

@@ -15,7 +15,7 @@ import Cardano.Ledger.Conway.PParams
 import Cardano.Ledger.Conway.State
 import Cardano.Ledger.Shelley.LedgerState
 import Cardano.Ledger.Shelley.Rules (ShelleyPoolPredFailure (..))
-import qualified Data.Set as Set
+import qualified Data.Map.Strict as Map
 import Lens.Micro
 import Test.Cardano.Ledger.Conway.ImpTest
 import Test.Cardano.Ledger.Core.Rational
@@ -52,7 +52,7 @@ spec = do
 
       expectVRFs [] -- VRF keyhashes in PState is not yet populated
       enactHardForkV11
-      expectVRFs [vrf2, vrf3, vrf5]
+      expectVRFs [(vrf2, 2), (vrf3, 1), (vrf5, 1)]
 
   it "Retiring a stake pool with a duplicate VRF Keyhash after v11 HardFork" $ do
     whenMajorVersion @10 $ do
@@ -61,21 +61,31 @@ spec = do
       registerStakePool kh1 vrf
       kh2 <- freshKeyHash
       registerStakePool kh2 vrf
+      kh3 <- freshKeyHash
+      registerStakePool kh3 vrf
 
       enactHardForkV11
-      expectVRFs [vrf]
+      expectVRFs [(vrf, 3)]
       -- retire one of the pools after the hard fork
       retireStakePool kh1 (EpochInterval 1)
+      retireStakePool kh2 (EpochInterval 1)
       passEpoch
       -- the vrf keyhash should still be present, since another pool is registered with it
-      expectVRFs [vrf]
+      expectVRFs [(vrf, 1)]
 
       -- registration of the same vrf should be disallowed
-      kh3 <- freshKeyHash
-      registerStakePoolTx kh3 vrf >>= \tx ->
+      kh4 <- freshKeyHash
+      registerStakePoolTx kh4 vrf >>= \tx ->
         submitFailingTx
           tx
-          [injectFailure $ VRFKeyHashAlreadyRegistered kh3 vrf]
+          [injectFailure $ VRFKeyHashAlreadyRegistered kh4 vrf]
+
+      retireStakePool kh3 (EpochInterval 1)
+      passEpoch
+      expectVRFs []
+
+      registerStakePool kh4 vrf
+      expectVRFs [(vrf, 1)]
   where
     enactHardForkV11 = do
       modifyPParams $ \pp ->
@@ -102,5 +112,7 @@ spec = do
         mkBasicTx mkBasicTxBody
           & bodyTxL . certsTxBodyL .~ [RetirePoolTxCert kh retirement]
     expectVRFs vrfs =
-      psVRFKeyHashes <$> getPState `shouldReturn` Set.fromList vrfs
+      psVRFKeyHashes
+        <$> getPState
+          `shouldReturn` Map.fromList [(k, unsafeNonZero v) | (k, v) <- vrfs]
     getPState = getsNES @era $ nesEsL . esLStateL . lsCertStateL . certPStateL

eras/shelley/impl/src/Cardano/Ledger/Shelley/Rules/Pool.hs

@@ -34,6 +34,7 @@ import Cardano.Ledger.BaseTypes (
   ShelleyBase,
   addEpochInterval,
   invalidKey,
+  knownNonZeroBounded,
   networkId,
  )
 import Cardano.Ledger.Binary (
@@ -68,7 +69,6 @@ import Control.State.Transition (
 import qualified Data.ByteString as BS
 import Data.Kind (Type)
 import qualified Data.Map as Map
-import qualified Data.Set as Set
 import Data.Word (Word8)
 import GHC.Generics (Generic)
 import Lens.Micro
@@ -253,9 +253,9 @@ poolDelegationTransition = do
         -- register new, Pool-Reg
         Nothing -> do
           when (hardforkConwayDisallowDuplicatedVRFKeys pv) $ do
-            Set.notMember ppVrf psVRFKeyHashes ?! VRFKeyHashAlreadyRegistered ppId ppVrf
+            Map.notMember ppVrf psVRFKeyHashes ?! VRFKeyHashAlreadyRegistered ppId ppVrf
           let updateVRFKeyHash
-                | hardforkConwayDisallowDuplicatedVRFKeys pv = Set.insert ppVrf
+                | hardforkConwayDisallowDuplicatedVRFKeys pv = Map.insert ppVrf (knownNonZeroBounded @1)
                 | otherwise = id
           tellEvent $ RegisterPool ppId
           pure $
@@ -266,17 +266,18 @@ poolDelegationTransition = do
         Just stakePoolState -> do
           when (hardforkConwayDisallowDuplicatedVRFKeys pv) $ do
             ppVrf == stakePoolState ^. spsVrfL
-              || Set.notMember ppVrf psVRFKeyHashes ?! VRFKeyHashAlreadyRegistered ppId ppVrf
+              || Map.notMember ppVrf psVRFKeyHashes ?! VRFKeyHashAlreadyRegistered ppId ppVrf
           let updateFutureVRFKeyHash
                 | hardforkConwayDisallowDuplicatedVRFKeys pv =
-                    -- If a pool re-registers with a fresh VRF, we have to add it to the list,
+                    -- If a pool re-registers with a fresh VRF, we have to record it in the map,
                     -- but also remove the previous VRFHashKey potentially stored in previous re-registration within the same epoch,
                     -- which we retrieve from futureStakePools.
                     case Map.lookup ppId psFutureStakePools of
-                      Nothing -> Set.insert ppVrf
+                      Nothing -> Map.insert ppVrf (knownNonZeroBounded @1)
                       Just futureStakePoolState
                         | futureStakePoolState ^. spsVrfL /= ppVrf ->
-                            Set.insert ppVrf . Set.delete (futureStakePoolState ^. spsVrfL)
+                            (Map.insert ppVrf (knownNonZeroBounded @1))
+                              . Map.delete (futureStakePoolState ^. spsVrfL)
                         | otherwise -> id
                 | otherwise = id
           tellEvent $ ReregisterPool ppId

eras/shelley/impl/src/Cardano/Ledger/Shelley/Rules/PoolReap.hs

@@ -23,7 +23,7 @@ module Cardano.Ledger.Shelley.Rules.PoolReap (
 ) where
 
 import Cardano.Ledger.Address
-import Cardano.Ledger.BaseTypes (ShelleyBase)
+import Cardano.Ledger.BaseTypes
 import Cardano.Ledger.Coin (Coin, CompactForm)
 import Cardano.Ledger.Compactible (fromCompact)
 import Cardano.Ledger.Core
@@ -35,7 +35,6 @@ import Cardano.Ledger.Shelley.LedgerState (
   utxosGovStateL,
  )
 import Cardano.Ledger.Shelley.LedgerState.Types (potEqualsObligation)
-import Cardano.Ledger.Slot (EpochNo (..))
 import Cardano.Ledger.State
 import Cardano.Ledger.Val ((<+>), (<->))
 import Control.DeepSeq (NFData)
@@ -50,11 +49,14 @@ import Control.State.Transition (
  )
 import Data.Default (Default, def)
 import Data.Foldable (fold)
+import Data.Foldable as F (foldl')
 import qualified Data.Map.Merge.Strict as Map
+import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
 import Data.Set (Set)
 import qualified Data.Set as Set
 import Data.Void (Void)
+import Data.Word (Word64)
 import GHC.Generics (Generic)
 import Lens.Micro
 
@@ -135,7 +137,7 @@ poolReapTransition = do
     ps0 = cs0 ^. certPStateL
     -- find the set of VRF key hashes that are no longer relevant, since they have been overwritten
     -- via pool re-registration
-    danglingVrfKeyHashes =
+    danglingVRFKeyHashes =
       Set.fromList $
         Map.elems $
           Map.merge
@@ -163,7 +165,8 @@ poolReapTransition = do
     retiringPools :: Map.Map (KeyHash 'StakePool) StakePoolState
     retiringPools = Map.restrictKeys (psStakePools ps) retired
     -- collect all accounts for stake pools that will retire
-    retiredVRFs = foldMap (Set.singleton . spsVrf) retiringPools
+    retiredVRFKeyHashes = spsVrf <$> Map.elems retiringPools
+
     -- collect all of the potential refunds
     accountRefunds :: Map.Map (Credential 'Staking) (CompactForm Coin)
     accountRefunds =
@@ -209,8 +212,19 @@ poolReapTransition = do
           & certPStateL . psStakePoolsL %~ (`Map.withoutKeys` retired)
           & certPStateL . psRetiringL %~ (`Map.withoutKeys` retired)
           & certPStateL . psVRFKeyHashesL
-            %~ ((`Set.difference` retiredVRFs) . (`Set.difference` danglingVrfKeyHashes))
+            %~ ( removeVRFKeyHashOccurrences (retiredVRFKeyHashes)
+                   . (`Map.withoutKeys` danglingVRFKeyHashes)
+               )
       )
+  where
+    removeVRFKeyHashOccurrences ::
+      [VRFVerKeyHash 'StakePoolVRF] ->
+      Map (VRFVerKeyHash 'StakePoolVRF) (NonZero Word64) ->
+      Map (VRFVerKeyHash 'StakePoolVRF) (NonZero Word64)
+    removeVRFKeyHashOccurrences vrfs vrfsMap = F.foldl' (flip removeVRFKeyHashOccurrence) vrfsMap vrfs
+    removeVRFKeyHashOccurrence =
+      -- Removes the key from the map if the value drops to 0
+      Map.update (mapNonZero (\n -> n - 1))
 
 renderPoolReapViolation ::
   ( EraGov era

eras/shelley/impl/testlib/Test/Cardano/Ledger/Shelley/Imp/PoolSpec.hs

@@ -344,7 +344,7 @@ spec = describe "POOL" $ do
         $ Map.member poolKh retiring == isRetiring
     expectVRFs vrfs = do
       whenMajorVersionAtLeast @11 $
-        psVRFKeyHashes <$> getPState `shouldReturn` vrfs
+        Map.keysSet . psVRFKeyHashes <$> getPState `shouldReturn` vrfs
     poolParams kh vrf = do
       pps <- registerRewardAccount >>= freshPoolParams kh
       pure $ pps & ppVrfL .~ vrf

eras/shelley/test-suite/test/Test/Cardano/Ledger/Shelley/Examples/Combinators.hs

@@ -354,7 +354,7 @@ reapPool pool cs = cs {chainNes = nes'}
       ps
         { psRetiring = Map.delete poolId (psRetiring ps)
         , psStakePools = Map.delete poolId (psStakePools ps)
-        , psVRFKeyHashes = Set.delete (ppVrf pool) (psVRFKeyHashes ps)
+        , psVRFKeyHashes = Map.delete (ppVrf pool) (psVRFKeyHashes ps)
         }
     pp = es ^. curPParamsEpochStateL
     ds = dps ^. certDStateL

eras/shelley/test-suite/test/Test/Cardano/Ledger/Shelley/Serialisation/Golden/Encoding.hs

@@ -1085,7 +1085,7 @@ tests =
           expectedHex =
             mconcat
               [ "8700a1581ce0a714319812c3f773ba04ec5d6b3ffcd5aad85006805b047b0825410aa1581ca64647"
-              , "4b8f5431261506b6c273d307c7569a4eb6c96b42dd4a29520a0384821927101903e882828480a0a0"
+              , "4b8f5431261506b6c273d307c7569a4eb6c96b42dd4a29520a0384821927101903e8828284a0a0a0"
               , "a08482a0a0a0a084a0a0000086a15822ee155ace9c40292074cb6aff8c9ccdd273c81648ff1149ef"
               , "36bcea6ebb8a3e250000583d003900cb9358529df4729c3246a2a033cb9821abbfd16de488800590"
               , "4abc410d6a577e9441ad8ed9663931906e4d43ece8f82c712b1d0235affb06000a1903e80185a0a0"

libs/cardano-ledger-core/src/Cardano/Ledger/State/CertState.hs

@@ -48,6 +48,7 @@ import Cardano.Ledger.BaseTypes (
   Anchor (..),
   AnchorData,
   KeyValuePairs (..),
+  NonZero,
   StrictMaybe,
   ToKeyValuePairs (..),
  )
@@ -84,8 +85,8 @@ import qualified Data.Foldable as F
 import Data.Kind (Type)
 import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
-import Data.Set (Set)
 import qualified Data.Set as Set
+import Data.Word (Word64)
 import GHC.Generics (Generic)
 import Lens.Micro
 import NoThunks.Class (NoThunks (..))
@@ -221,7 +222,7 @@ lookupRewardDState DState {dsAccounts} cred = do
 
 -- | The state used by the POOL rule, which tracks stake pool information.
 data PState era = PState
-  { psVRFKeyHashes :: !(Set (VRFVerKeyHash 'StakePoolVRF))
+  { psVRFKeyHashes :: !(Map (VRFVerKeyHash 'StakePoolVRF) (NonZero Word64))
   -- ^ VRF key hashes that have been registered via PoolParams
   , psStakePools :: !(Map (KeyHash 'StakePool) StakePoolState)
   -- ^ The state of current stake pools.
@@ -386,7 +387,7 @@ instance Default (Accounts era) => Default (DState era) where
   def = DState def Map.empty (GenDelegs Map.empty) def
 
 instance Default (PState era) where
-  def = PState Set.empty Map.empty Map.empty Map.empty
+  def = PState Map.empty Map.empty Map.empty Map.empty
 
 -- | A composite of all the Deposits the system is obligated to eventually pay back.
 data Obligations = Obligations
@@ -443,5 +444,5 @@ psFutureStakePoolsL = lens psFutureStakePools (\ps u -> ps {psFutureStakePools =
 psRetiringL :: Lens' (PState era) (Map (KeyHash 'StakePool) EpochNo)
 psRetiringL = lens psRetiring (\ps u -> ps {psRetiring = u})
 
-psVRFKeyHashesL :: Lens' (PState era) (Set (VRFVerKeyHash 'StakePoolVRF))
+psVRFKeyHashesL :: Lens' (PState era) (Map (VRFVerKeyHash 'StakePoolVRF) (NonZero Word64))
 psVRFKeyHashesL = lens psVRFKeyHashes (\ps u -> ps {psVRFKeyHashes = u})

libs/cardano-ledger-test/src/Test/Cardano/Ledger/Constrained/Conway/Instances/Basic.hs

@@ -418,3 +418,10 @@ class
   ppToSubset :: PParams era -> SimplePParams era
   updateToPPU :: SimplePPUpdate -> PParamsUpdate era
   ppuToUpdate :: PParamsUpdate era -> SimplePPUpdate
+
+instance HasSimpleRep (NonZero Word64) where
+  type SimpleRep (NonZero Word64) = Word64
+  toSimpleRep = unNonZero
+  fromSimpleRep = unsafeNonZero
+
+instance HasSpec (NonZero Word64)

libs/cardano-ledger-test/src/Test/Cardano/Ledger/Generic/GenState.hs

@@ -669,7 +669,7 @@ initialLedgerState gstate = LedgerState utxostate dpstate
         instantaneousRewardsZero
     pstate =
       PState
-        Set.empty
+        Map.empty
         (mkStakePoolState poolDeposit <$> pools)
         Map.empty
         Map.empty