Refactor VState update logic to handle dRep undelegation/redelegation

Author: teodanciu

eras/conway/impl/CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## 1.21.0.0
 
+* Add `unDelegReDelegDRep` to `VState` module
 * Expose `conwayRegisterInitialAccounts`
 * Add `TxLevel` argument to `Tx` and `TxBody`
 * Add `HasEraTxLevel` instances for `Tx` and `TxBody`

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

@@ -227,7 +227,6 @@ conwayDelegTransition = do
             %~ registerConwayAccount stakeCred ppKeyDepositCompact Nothing
     ConwayUnRegCert stakeCred sMayRefund -> do
       let (mAccountState, newAccounts) = unregisterConwayAccount stakeCred accounts
-          mCurDelegatee = mAccountState >>= accountStateDelegatee
           checkInvalidRefund = do
             SJust suppliedRefund <- Just sMayRefund
             -- we don't want to report invalid refund when stake credential is not registered:
@@ -259,7 +258,7 @@ conwayDelegTransition = do
           pure $
             certState
               & certDStateL . accountsL .~ newAccounts
-              & certVStateL %~ unDelegDRep stakeCred mCurDelegatee
+              & certVStateL %~ unDelegReDelegDRep stakeCred accountState Nothing
               & certPStateL %~ unDelegReDelegStakePool stakeCred accountState Nothing
     ConwayDelegCert stakeCred delegatee -> do
       let mAccountState = lookupAccountState stakeCred accounts
@@ -312,7 +311,6 @@ processDelegationInternal preserveIncorrectDelegation stakeCred mAccountState ne
     DelegVote dRep -> delegVote dRep
     DelegStakeVote sPool dRep -> delegVote dRep . delegStake sPool
   where
-    mCurDelegatee = mAccountState >>= accountStateDelegatee
     delegStake stakePool cState =
       cState
         & certDStateL . accountsL
@@ -322,33 +320,13 @@ processDelegationInternal preserveIncorrectDelegation stakeCred mAccountState ne
           (\accountState -> certPStateL %~ unDelegReDelegStakePool stakeCred accountState (Just stakePool))
           mAccountState
     delegVote dRep cState =
-      let cState' =
-            cState
-              & certDStateL . accountsL
-                %~ adjustAccountState (dRepDelegationAccountStateL ?~ dRep) stakeCred
-              & certVStateL %~ unDelegDRep stakeCred mCurDelegatee
-          dReps
-            | preserveIncorrectDelegation = cState ^. certVStateL . vsDRepsL
-            | otherwise = cState' ^. certVStateL . vsDRepsL
-       in case dRep of
-            DRepCredential targetDRep
-              | Just dRepState <- Map.lookup targetDRep dReps ->
-                  let dRepState' = dRepState {drepDelegs = Set.insert stakeCred (drepDelegs dRepState)}
-                   in cState' & certVStateL . vsDRepsL .~ Map.insert targetDRep dRepState' dReps
-            _ -> cState'
-
-unDelegDRep ::
-  Credential 'Staking ->
-  Maybe Delegatee ->
-  VState era ->
-  VState era
-unDelegDRep stakeCred mCurDelegatee =
-  maybe
-    id
-    (\dRepCred -> vsDRepsL %~ Map.adjust (drepDelegsL %~ Set.delete stakeCred) dRepCred)
-    (mCurDelegatee >>= drepToUndeleg)
-  where
-    drepToUndeleg = \case
-      DelegVote (DRepCredential dRepCred) -> Just dRepCred
-      DelegStakeVote _ (DRepCredential dRepCred) -> Just dRepCred
-      _ -> Nothing
+      cState
+        & certDStateL . accountsL %~ adjustAccountState (dRepDelegationAccountStateL ?~ dRep) stakeCred
+        & maybe
+          (certVStateL %~ insertDRepDeleg dRep)
+          (\accountState -> certVStateL %~ unDelegReDelegDRep stakeCred accountState (Just dRep))
+          (guard (not preserveIncorrectDelegation) >> mAccountState)
+    insertDRepDeleg dRep = case dRep of
+      DRepCredential dRepCred ->
+        vsDRepsL %~ Map.adjust (drepDelegsL %~ Set.insert stakeCred) dRepCred
+      _ -> id

eras/conway/impl/src/Cardano/Ledger/Conway/State/VState.hs

@@ -17,6 +17,7 @@ module Cardano.Ledger.Conway.State.VState (
   vsNumDormantEpochsL,
   vsActualDRepExpiry,
   lookupDepositVState,
+  unDelegReDelegDRep,
 ) where
 
 import Cardano.Ledger.BaseTypes (KeyValuePairs (..), ToKeyValuePairs (..), binOpEpochNo)
@@ -33,17 +34,21 @@ import Cardano.Ledger.Binary (
 import Cardano.Ledger.Binary.Coders (Decode (..), Encode (..), decode, encode, (!>), (<!))
 import Cardano.Ledger.Coin (Coin (..))
 import Cardano.Ledger.Compactible (Compactible (..))
+import Cardano.Ledger.Conway.State.Account
 import Cardano.Ledger.Core
 import Cardano.Ledger.Credential (Credential (..))
+import Cardano.Ledger.DRep (drepDelegsL)
 import Cardano.Ledger.Shelley.State
 import Cardano.Ledger.Slot (EpochNo (..))
 import Control.DeepSeq (NFData (..))
 import Data.Aeson (ToJSON (..), (.=))
 import Data.Default (Default (def))
 import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
+import Data.Maybe (fromMaybe)
+import qualified Data.Set as Set
 import GHC.Generics (Generic)
-import Lens.Micro (Lens', lens, (^.))
+import Lens.Micro
 import NoThunks.Class (NoThunks (..))
 
 -- | The state that tracks the voting entities (DReps and Constitutional Committee
@@ -109,6 +114,34 @@ instance ToKeyValuePairs (VState era) where
         , "numDormantEpochs" .= vsNumDormantEpochs
         ]
 
+-- | Reverses DRep 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.
+unDelegReDelegDRep ::
+  ConwayEraAccounts era =>
+  Credential 'Staking ->
+  -- | Account that is losing its current delegation and/or acquiring a new one
+  AccountState era ->
+  -- | Potential new delegation. In case when stake credential unregisters this must be `Nothing`.
+  Maybe DRep ->
+  VState era ->
+  VState era
+unDelegReDelegDRep stakeCred accountState mNewDRep =
+  fromMaybe (vsDRepsL %~ addNewDelegation) $ do
+    dRep@(DRepCredential dRepCred) <- accountState ^. dRepDelegationAccountStateL
+    pure $
+      -- There is no need to update set of delegations if delegation is unchanged
+      if Just dRep == mNewDRep
+        then id
+        else
+          vsDRepsL %~ addNewDelegation . Map.adjust (drepDelegsL %~ Set.delete stakeCred) dRepCred
+  where
+    addNewDelegation =
+      case mNewDRep of
+        Just (DRepCredential dRepCred) ->
+          Map.adjust (drepDelegsL %~ Set.insert stakeCred) dRepCred
+        _ -> id
+
 vsDRepsL :: Lens' (VState era) (Map (Credential 'DRepRole) DRepState)
 vsDRepsL = lens vsDReps (\vs u -> vs {vsDReps = u})