Update StakePoolState delegs in shelley

Author: teodanciu

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

@@ -271,17 +271,33 @@ delegationTransition = do
             let accountBalance = accountState ^. balanceAccountStateL
             guard (accountBalance /= mempty)
             Just $ fromCompact accountBalance
-      -- (hk ∈ dom (rewards ds))
-      isJust mAccountState ?! StakeKeyNotRegisteredDELEG cred
       failOnJust checkStakeKeyHasZeroRewardBalance StakeKeyNonZeroAccountBalanceDELEG
-      pure $ certState & certDStateL . accountsL .~ accounts
+      -- (hk ∈ dom (rewards ds))
+      case mAccountState of
+        Nothing -> do
+          failBecause $ StakeKeyNotRegisteredDELEG cred
+          pure certState
+        Just accountState ->
+          pure $
+            certState
+              & certDStateL . accountsL .~ accounts
+              & certPStateL
+                %~ unDelegStakePool cred (accountState ^. stakePoolDelegationAccountStateL) Nothing
     DelegStakeTxCert cred stakePool -> do
       -- note that pattern match is used instead of cwitness and dpool, as in the spec
       -- (hk ∈ dom (rewards ds))
-      isAccountRegistered cred (ds ^. accountsL) ?! StakeDelegationImpossibleDELEG cred
-      pure $
-        certState
-          & certDStateL . accountsL %~ adjustAccountState (stakePoolDelegationAccountStateL ?~ stakePool) cred
+      case lookupAccountState cred (ds ^. accountsL) of
+        Nothing -> do
+          failBecause $ StakeDelegationImpossibleDELEG cred
+          pure certState
+        Just currentAccount ->
+          pure $
+            certState
+              & certDStateL . accountsL %~ adjustAccountState (stakePoolDelegationAccountStateL ?~ stakePool) cred
+              & certPStateL %~ \ps ->
+                ps
+                  & unDelegStakePool cred (currentAccount ^. stakePoolDelegationAccountStateL) (Just stakePool)
+                  & psStakePoolsL %~ Map.adjust (spsDelegsL %~ Set.insert cred) stakePool
     GenesisDelegTxCert gkh vkh vrf -> do
       sp <- liftSTS $ asks stabilityWindow
       -- note that pattern match is used instead of genesisDeleg, as in the spec
@@ -404,3 +420,19 @@ 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 (spsDelegsL %~ Set.delete stakeCred) oldPool)
+    (mCurStakePool >>= stakePoolToUnDeleg)
+  where
+    stakePoolToUnDeleg oldPool
+      | Just oldPool /= mNewPool = Just oldPool
+      | otherwise = Nothing

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

@@ -95,6 +95,7 @@ import Data.Foldable (fold)
 import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
 import Data.Maybe (fromJust, maybeToList)
+import qualified Data.Set as Set
 import Data.Word (Word64)
 import GHC.Stack (HasCallStack)
 import Lens.Micro
@@ -211,12 +212,19 @@ deregStakeCred cred cs = cs {chainNes = nes}
       chainNes cs
         & nesEsL . esLStateL . lsCertStateL . certDStateL . accountsL .~ accounts'
         & nesEsL . esLStateL . lsUTxOStateL . utxosDepositedL %~ (<-> refund)
+        & nesEsL . esLStateL . lsCertStateL . certPStateL %~ adjustPState
     accounts =
       chainNes cs
         ^. nesEsL . esLStateL . lsCertStateL . certDStateL . accountsL
     (mAccountState, accounts') =
       unregisterShelleyAccount cred accounts
     refund = fromCompact (fromJust mAccountState ^. depositAccountStateL)
+    currentDeleg = Map.lookup cred (accounts ^. accountsMapL) >>= (^. stakePoolDelegationAccountStateL)
+    adjustPState =
+      maybe
+        id
+        (\oldPool -> psStakePoolsL %~ Map.adjust (spsDelegsL %~ Set.delete cred) oldPool)
+        currentDeleg
 
 -- | = New Delegation
 --
@@ -234,6 +242,8 @@ delegation cred poolId cs = cs {chainNes = nes}
       chainNes cs
         & nesEsL . esLStateL . lsCertStateL . certDStateL . accountsL . accountsMapL
           %~ Map.adjust (stakePoolDelegationAccountStateL .~ Just poolId) cred
+        & nesEsL . esLStateL . lsCertStateL . certPStateL . psStakePoolsL
+          %~ Map.adjust (spsDelegsL %~ Set.insert cred) poolId
 
 -- | Register a stake pool.
 regPool ::