Refactor PState update logic to handle stake undelegation/redelegation

uniformly across eras

Author: teodanciu

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

@@ -67,6 +67,7 @@ import Control.State.Transition (
   State,
   TRC (TRC),
   TransitionRule,
+  failBecause,
   failOnJust,
   judgmentContext,
   transitionRules,
@@ -204,7 +205,7 @@ conwayDelegTransition = do
     checkStakeKeyIsRegistered stakeCred = do
       let mAccountState = lookupAccountState stakeCred accounts
       isJust mAccountState ?! StakeKeyNotRegisteredDELEG stakeCred
-      pure $ mAccountState >>= accountStateDelegatee
+      pure mAccountState
     checkStakeDelegateeRegistered =
       let checkPoolRegistered targetPool =
             targetPool `Map.member` pools ?! DelegateeStakePoolNotRegisteredDELEG targetPool
@@ -253,18 +254,23 @@ conwayDelegTransition = do
             guard (balanceCompact /= mempty)
             Just $ fromCompact balanceCompact
       failOnJust checkInvalidRefund id
-      isJust mAccountState ?! StakeKeyNotRegisteredDELEG stakeCred
       failOnJust checkStakeKeyHasZeroRewardBalance StakeKeyHasNonZeroRewardAccountBalanceDELEG
-      pure $
-        certState
-          & certDStateL . accountsL .~ newAccounts
-          & certVStateL %~ unDelegDRep stakeCred mCurDelegatee
-          & certPStateL %~ unDelegStakePool stakeCred mCurDelegatee Nothing
+      case mAccountState of
+        Nothing ->
+          do
+            failBecause $ StakeKeyNotRegisteredDELEG stakeCred
+            pure certState
+        Just accountState ->
+          pure $
+            certState
+              & certDStateL . accountsL .~ newAccounts
+              & certVStateL %~ unDelegDRep stakeCred mCurDelegatee
+              & certPStateL %~ unDelegReDelegStakePool stakeCred accountState Nothing
     ConwayDelegCert stakeCred delegatee -> do
-      mCurDelegatee <- checkStakeKeyIsRegistered stakeCred
+      mAccountState <- checkStakeKeyIsRegistered stakeCred
       checkStakeDelegateeRegistered delegatee
       pure $
-        processDelegationInternal (pvMajor pv < natVersion @10) stakeCred mCurDelegatee delegatee certState
+        processDelegationInternal (pvMajor pv < natVersion @10) stakeCred mAccountState delegatee certState
     ConwayRegDelegCert stakeCred delegatee deposit -> do
       checkDepositAgainstPParams deposit
       checkStakeKeyNotRegistered stakeCred
@@ -287,9 +293,8 @@ processDelegation ::
   CertState era
 processDelegation stakeCred newDelegatee !certState = certState'
   where
-    !certState' = processDelegationInternal False stakeCred mCurDelegatee newDelegatee certState
+    !certState' = processDelegationInternal False stakeCred mAccountState newDelegatee certState
     mAccountState = Map.lookup stakeCred (certState ^. certDStateL . accountsL . accountsMapL)
-    mCurDelegatee = mAccountState >>= accountStateDelegatee
 
 -- | Same as `processDelegation`, except it expects the current delegation supplied as an
 -- argument, because in ledger rules we already have it readily available.
@@ -299,23 +304,27 @@ processDelegationInternal ::
   Bool ->
   -- | Delegator
   Credential 'Staking ->
-  -- | Current delegatee for the above stake credential that needs to be cleaned up.
-  Maybe Delegatee ->
+  -- | Account state for the above stake credential
+  Maybe (AccountState era) ->
   -- | New delegatee
   Delegatee ->
   CertState era ->
   CertState era
-processDelegationInternal preserveIncorrectDelegation stakeCred mCurDelegatee newDelegatee =
+processDelegationInternal preserveIncorrectDelegation stakeCred mAccountState newDelegatee =
   case newDelegatee of
     DelegStake sPool -> delegStake sPool
     DelegVote dRep -> delegVote dRep
     DelegStakeVote sPool dRep -> delegVote dRep . delegStake sPool
   where
+    mCurDelegatee = mAccountState >>= accountStateDelegatee
     delegStake stakePool cState =
       cState
         & certDStateL . accountsL
           %~ adjustAccountState (stakePoolDelegationAccountStateL ?~ stakePool) stakeCred
-        & certPStateL %~ adjustPState stakePool
+        & maybe
+          (certPStateL . psStakePoolsL %~ Map.adjust (spsDelegatorsL %~ Set.insert stakeCred) stakePool)
+          (\accountState -> certPStateL %~ unDelegReDelegStakePool stakeCred accountState (Just stakePool))
+          mAccountState
     delegVote dRep cState =
       let cState' =
             cState
@@ -331,26 +340,6 @@ processDelegationInternal preserveIncorrectDelegation stakeCred mCurDelegatee ne
                   let dRepState' = dRepState {drepDelegs = Set.insert stakeCred (drepDelegs dRepState)}
                    in cState' & certVStateL . vsDRepsL .~ Map.insert targetDRep dRepState' dReps
             _ -> cState'
-    adjustPState newPool =
-      (psStakePoolsL %~ Map.adjust (spsDelegatorsL %~ Set.insert stakeCred) newPool)
-        . unDelegStakePool stakeCred mCurDelegatee (Just newPool)
-
-unDelegStakePool ::
-  Credential 'Staking ->
-  Maybe Delegatee ->
-  Maybe (KeyHash 'StakePool) ->
-  PState era ->
-  PState era
-unDelegStakePool stakeCred mCurDelegatee mNewPool =
-  maybe
-    id
-    (\oldPool -> psStakePoolsL %~ Map.adjust (spsDelegatorsL %~ Set.delete stakeCred) oldPool)
-    (mCurDelegatee >>= stakePoolToUnDeleg)
-  where
-    stakePoolToUnDeleg = \case
-      DelegStake oldPool | Just oldPool /= mNewPool -> Just oldPool
-      DelegStakeVote oldPool _ | Just oldPool /= mNewPool -> Just oldPool
-      _ -> Nothing
 
 unDelegDRep ::
   Credential 'Staking ->

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

@@ -282,7 +282,7 @@ delegationTransition = do
             certState
               & certDStateL . accountsL .~ accounts
               & certPStateL
-                %~ unDelegStakePool cred (accountState ^. stakePoolDelegationAccountStateL) Nothing
+                %~ unDelegReDelegStakePool cred accountState Nothing
     DelegStakeTxCert cred stakePool -> do
       -- note that pattern match is used instead of cwitness and dpool, as in the spec
       -- (hk ∈ dom (rewards ds))
@@ -294,10 +294,7 @@ delegationTransition = do
           pure $
             certState
               & certDStateL . accountsL %~ adjustAccountState (stakePoolDelegationAccountStateL ?~ stakePool) cred
-              & certPStateL %~ \ps ->
-                ps
-                  & unDelegStakePool cred (accountState ^. stakePoolDelegationAccountStateL) (Just stakePool)
-                  & psStakePoolsL %~ Map.adjust (spsDelegatorsL %~ Set.insert cred) stakePool
+              & certPStateL %~ unDelegReDelegStakePool cred accountState (Just stakePool)
     GenesisDelegTxCert gkh vkh vrf -> do
       sp <- liftSTS $ asks stabilityWindow
       -- note that pattern match is used instead of genesisDeleg, as in the spec
@@ -420,19 +417,3 @@ updateReservesAndTreasury targetPot combinedMap available certState = do
     case targetPot of
       ReservesMIR -> certState & certDStateL . dsIRewardsL . iRReservesL .~ combinedMap
       TreasuryMIR -> certState & certDStateL . dsIRewardsL . iRTreasuryL .~ combinedMap
-
-unDelegStakePool ::
-  Credential 'Staking ->
-  Maybe (KeyHash 'StakePool) ->
-  Maybe (KeyHash 'StakePool) ->
-  PState era ->
-  PState era
-unDelegStakePool stakeCred mCurStakePool mNewPool =
-  maybe
-    id
-    (\oldPool -> psStakePoolsL %~ Map.adjust (spsDelegatorsL %~ Set.delete stakeCred) oldPool)
-    (mCurStakePool >>= stakePoolToUnDeleg)
-  where
-    stakePoolToUnDeleg oldPool
-      | Just oldPool /= mNewPool = Just oldPool
-      | otherwise = Nothing

libs/cardano-ledger-core/CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## 1.19.0.0
 
+* Add `unDelegReDelegStakePool` to `CertState` module
 * Add `iRReservesL`, `iRTreasuryL`, `iRDeltaReservesL`, `iRDeltaTreasuryL`
 * Add `spsDelegators` field to `StakePool`
 * Add `spsDelegatorsL`

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

@@ -34,6 +34,7 @@ module Cardano.Ledger.State.CertState (
   lookupRewardDState,
   Obligations (..),
   sumObligation,
+  unDelegReDelegStakePool,
   -- Lenses
   iRReservesL,
   dsIRewardsL,
@@ -80,7 +81,7 @@ import Cardano.Ledger.DRep (DRep (..), DRepState (..))
 import Cardano.Ledger.Hashes (GenDelegPair (..), GenDelegs (..))
 import Cardano.Ledger.Slot (EpochNo (..), SlotNo (..))
 import Cardano.Ledger.State.Account
-import Cardano.Ledger.State.StakePool (StakePoolState (..))
+import Cardano.Ledger.State.StakePool (StakePoolState (..), spsDelegatorsL)
 import Control.DeepSeq (NFData (..))
 import Control.Monad.Trans
 import Data.Aeson (ToJSON (..), object, (.=))
@@ -89,6 +90,7 @@ import qualified Data.Foldable as F
 import Data.Kind (Type)
 import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
+import Data.Maybe (fromMaybe)
 import qualified Data.Set as Set
 import Data.Word (Word64)
 import GHC.Generics (Generic)
@@ -273,6 +275,30 @@ instance ToKeyValuePairs (PState era) where
     , "retiring" .= psRetiring
     ]
 
+-- | Reverses stake pool delegation.
+-- To be called when a stake credential is unregistered or its delegation target changes.
+-- If the new delegation matches the previous one, this is a noop.
+unDelegReDelegStakePool ::
+  EraAccounts era =>
+  Credential 'Staking ->
+  -- | Account that is losing its current delegation and/or acquiring a new one
+  AccountState era ->
+  -- | Optional new delegation target. Use 'Nothing' when the stake credential unregisters.
+  Maybe (KeyHash 'StakePool) ->
+  PState era ->
+  PState era
+unDelegReDelegStakePool stakeCred accountState mNewStakePool =
+  fromMaybe (psStakePoolsL %~ addNewDelegation) $ do
+    curStakePool <- accountState ^. stakePoolDelegationAccountStateL
+    pure $
+      -- no need to update the set of delegations if the delegation is unchanged
+      if Just curStakePool == mNewStakePool
+        then id
+        else
+          psStakePoolsL %~ addNewDelegation . Map.adjust (spsDelegatorsL %~ Set.delete stakeCred) curStakePool
+  where
+    addNewDelegation = maybe id (Map.adjust (spsDelegatorsL %~ Set.insert stakeCred)) mNewStakePool
+
 data CommitteeAuthorization
   = -- | Member authorized with a Hot credential acting on behalf of their Cold credential
     CommitteeHotCredential !(Credential 'HotCommitteeRole)