Set delegators of updated pools in PoolReap rule

in order not to lose the delegations that happened after the update,
before the end of the epoch

Author: teodanciu

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

@@ -277,7 +277,7 @@ poolDelegationTransition = do
                       Nothing -> Map.insert ppVrf (knownNonZeroBounded @1)
                       Just futureStakePoolState
                         | futureStakePoolState ^. spsVrfL /= ppVrf ->
-                            (Map.insert ppVrf (knownNonZeroBounded @1))
+                            Map.insert ppVrf (knownNonZeroBounded @1)
                               . Map.delete (futureStakePoolState ^. spsVrfL)
                         | otherwise -> id
                 | otherwise = id
@@ -297,7 +297,9 @@ poolDelegationTransition = do
           let futureStakePoolState =
                 mkStakePoolState
                   (stakePoolState ^. spsDepositL)
-                  (stakePoolState ^. spsDelegatorsL)
+                  -- delegators are set in PoolReap,
+                  -- in order to capture delegations that happened after re-registration but before the end of the epoch
+                  mempty
                   poolParams
           pure $
             ps

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

@@ -152,7 +152,11 @@ poolReapTransition = do
     -- activate future stakePools
     ps =
       ps0
-        { psStakePools = Map.union (ps0 ^. psFutureStakePoolsL) (ps0 ^. psStakePoolsL)
+        { psStakePools =
+            Map.unionWith
+              (\newPoolState oldPoolState -> newPoolState {spsDelegators = spsDelegators oldPoolState})
+              (ps0 ^. psFutureStakePoolsL)
+              (ps0 ^. psStakePoolsL)
         , psFutureStakePools = Map.empty
         }
     cs = cs0 & certPStateL .~ ps
@@ -226,7 +230,6 @@ poolReapTransition = do
     removeVRFKeyHashOccurrence =
       -- Removes the key from the map if the value drops to 0
       Map.update (mapNonZero (\n -> n - 1))
-
     delegsToClear cState pools =
       foldMap spsDelegators $
         Map.restrictKeys (cState ^. certPStateL . psStakePoolsL) pools

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

@@ -16,7 +16,7 @@ import Cardano.Ledger.Credential (Credential (..))
 import Cardano.Ledger.Shelley.LedgerState
 import Cardano.Ledger.Shelley.Rules (ShelleyPoolPredFailure (..))
 import Cardano.Ledger.State
-import qualified Data.Map.Strict as Map
+import Data.Map.Strict as Map
 import Data.Proxy
 import Lens.Micro
 import Test.Cardano.Ledger.Binary.Arbitrary (genByteString)
@@ -52,6 +52,7 @@ spec = describe "POOL" $ do
           submitTx_ tx
         else
           submitFailingTx tx [injectFailure $ WrongNetworkPOOL (Mismatch Mainnet Testnet) kh]
+
     it "register a pool with too big metadata" $ do
       pv <- getsPParams ppProtocolVersionL
       let maxMetadataSize = sizeHash (Proxy :: Proxy HASH)
@@ -82,6 +83,39 @@ spec = describe "POOL" $ do
             expectPool khNew Nothing
       expectPool kh (Just vrf)
 
+    it "re-register a pool and change its delegations in the same epoch" $ do
+      (poolKh, _) <- registerNewPool
+      (poolKh2, _) <- registerNewPool
+      stakeCred <- KeyHashObj <$> freshKeyHash
+      _ <- registerStakeCredential stakeCred
+      stakeCred2 <- KeyHashObj <$> freshKeyHash
+      _ <- registerStakeCredential stakeCred2
+      delegateStake stakeCred poolKh
+      vrf1 <- freshKeyHashVRF
+      registerPoolTx <$> poolParams poolKh vrf1 >>= \tx -> do
+        submitTx_ tx
+        expectPoolDelegs poolKh (Just [stakeCred])
+        delegateStake stakeCred2 poolKh
+        expectPoolDelegs poolKh (Just [stakeCred, stakeCred2])
+        passEpoch
+        expectPoolDelegs poolKh (Just [stakeCred, stakeCred2])
+
+      vrf2 <- freshKeyHashVRF
+      registerPoolTx <$> poolParams poolKh vrf2 >>= \tx -> do
+        submitTx_ tx
+        expectPoolDelegs poolKh (Just [stakeCred, stakeCred2])
+
+        unRegTxCert <- genUnRegTxCert stakeCred2
+        submitTx_ $
+          mkBasicTx mkBasicTxBody
+            & bodyTxL . certsTxBodyL
+              .~ [unRegTxCert]
+        expectPoolDelegs poolKh (Just [stakeCred])
+        delegateStake stakeCred poolKh2
+        expectPoolDelegs poolKh (Just [])
+        passEpoch
+        expectPoolDelegs poolKh (Just [])
+
     it "re-register a pool with an already registered VRF" $ do
       pv <- getsPParams ppProtocolVersionL
       (kh1, vrf1) <- registerNewPool
@@ -287,6 +321,27 @@ spec = describe "POOL" $ do
       expectRetiring False khNew
       expectPool kh Nothing
 
+    it "retiring a pool clears its delegations" $ do
+      (poolKh, _) <- registerNewPool
+      let retirement = 1
+      stakeCred1 <- do
+        cred <- KeyHashObj <$> freshKeyHash
+        _ <- registerStakeCredential cred
+        delegateStake cred poolKh
+        pure cred
+
+      retirePoolTx poolKh (EpochInterval retirement) >>= submitTx_
+      expectPoolDelegs poolKh (Just [stakeCred1])
+      stakeCred2 <- do
+        cred <- KeyHashObj <$> freshKeyHash
+        _ <- registerStakeCredential cred
+        delegateStake cred poolKh
+        pure cred
+      expectPoolDelegs poolKh (Just [stakeCred1, stakeCred2])
+
+      passNEpochs (fromIntegral retirement)
+      expectPoolDelegs poolKh Nothing
+
   describe "Retired pools" $ do
     it "re-register a pool with the same keyhash and VRF " $ do
       (kh, vrf) <- registerNewPool
@@ -331,6 +386,9 @@ spec = describe "POOL" $ do
     expectFuturePool poolKh mbVrf = do
       fps <- psFutureStakePools <$> getPState
       spsVrf <$> Map.lookup poolKh fps `shouldBe` mbVrf
+    expectPoolDelegs poolKh delegs = do
+      pps <- psStakePools <$> getPState
+      spsDelegators <$> Map.lookup poolKh pps `shouldBe` delegs
     expectRetiring isRetiring poolKh = do
       retiring <- psRetiring <$> getPState
       assertBool