[wip] - removePoolsDelegations

Author: teodanciu

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

@@ -208,11 +208,12 @@ poolReapTransition = do
       a {casTreasury = casTreasury a <+> fromCompact unclaimed}
       ( cs
           & certDStateL . accountsL
-            %~ removeStakePoolDelegations retired . addToBalanceAccounts refunds
+            %~ removeStakePoolDelegations (allDelegsToClear cs retired)
+              . addToBalanceAccounts refunds
           & certPStateL . psStakePoolsL %~ (`Map.withoutKeys` retired)
           & certPStateL . psRetiringL %~ (`Map.withoutKeys` retired)
           & certPStateL . psVRFKeyHashesL
-            %~ ( removeVRFKeyHashOccurrences (retiredVRFKeyHashes)
+            %~ ( removeVRFKeyHashOccurrences retiredVRFKeyHashes
                    . (`Map.withoutKeys` danglingVRFKeyHashes)
                )
       )
@@ -226,6 +227,11 @@ poolReapTransition = do
       -- Removes the key from the map if the value drops to 0
       Map.update (mapNonZero (\n -> n - 1))
 
+    allDelegsToClear cState pools =
+      foldMap spsDelegs $
+        Map.elems $
+          Map.restrictKeys (cState ^. certPStateL . psStakePoolsL) pools
+
 renderPoolReapViolation ::
   ( EraGov era
   , State t ~ ShelleyPoolreapState era

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

@@ -94,8 +94,7 @@ import Cardano.Slotting.Slot (EpochNo, WithOrigin (..))
 import Data.Foldable (fold)
 import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
-import Data.Maybe (fromJust)
-import qualified Data.Set as Set
+import Data.Maybe (fromJust, maybeToList)
 import Data.Word (Word64)
 import GHC.Stack (HasCallStack)
 import Lens.Micro
@@ -369,7 +368,11 @@ reapPool pool cs = cs {chainNes = nes'}
            in ( accounts & accountsMapL %~ Map.insert poolAccountCred accountState'
               , mempty
               )
-    ds' = ds {dsAccounts = removeStakePoolDelegations (Set.singleton poolId) accounts'}
+    delegsToClear =
+      foldMap spsDelegs $
+        maybeToList $
+          Map.lookup poolId (dps ^. certPStateL . psStakePoolsL)
+    ds' = ds {dsAccounts = removeStakePoolDelegations delegsToClear accounts'}
     chainAccountState = esChainAccountState es
     chainAccountState' = chainAccountState {casTreasury = casTreasury chainAccountState <+> fromCompact unclaimed}
     utxoSt = lsUTxOState ls

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

@@ -45,7 +45,6 @@ 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 Lens.Micro
 import NoThunks.Class (NoThunks)
 
@@ -164,7 +163,7 @@ isAccountRegistered cred accounts = Map.member cred (accounts ^. accountsMapL)
 adjustAccountState ::
   EraAccounts era =>
   (AccountState era -> AccountState era) -> Credential 'Staking -> Accounts era -> Accounts era
-adjustAccountState cred f = accountsMapL %~ Map.adjust cred f
+adjustAccountState f cred = accountsMapL %~ Map.adjust f cred
 
 -- | In case when account state is registered and it is delegated to a stake pool this function
 -- will return that delegation.
@@ -238,18 +237,15 @@ drainAccounts (Withdrawals withdrawalsMap) accounts =
         accountsMap
         withdrawalsMap
 
--- TODO: This is an expensive operation, since it iterates over the whole accountsMap. We need to
--- start keeping track of all delegations to the stake pool in its state, then we would be able to
--- switch from `Set (KeyHash 'StakePool)` to `Map (KeyHash 'StakePool) (Set (Credential Staking))`
--- and drastically speed up this operation.
-
 -- | Remove delegations for the supplied Stake
 removeStakePoolDelegations ::
-  EraAccounts era => Set (KeyHash 'StakePool) -> Accounts era -> Accounts era
+  EraAccounts era => Set (Credential 'Staking) -> Accounts era -> Accounts era
 removeStakePoolDelegations stakeDelegationsToRemove accounts =
-  accounts & accountsMapL %~ Map.map clearAccountStateDelegation
-  where
-    clearAccountStateDelegation =
-      stakePoolDelegationAccountStateL %~ \case
-        Just poolId | poolId `Set.member` stakeDelegationsToRemove -> Nothing
-        delegation -> delegation
+  accounts
+    & accountsMapL
+      %~ ( \accountsMap ->
+             foldr
+               (Map.adjust (stakePoolDelegationAccountStateL .~ Nothing))
+               accountsMap
+               stakeDelegationsToRemove
+         )