Improve certificate performance

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

@@ -63,7 +63,6 @@ import Control.State.Transition (
 import Data.Map (Map)
 import qualified Data.Map.Strict as Map
 import Data.Maybe (isJust)
-import qualified Data.Set as Set
 import Data.Void (Void)
 import GHC.Generics (Generic)
 import Lens.Micro ((^.))
@@ -193,22 +192,22 @@ conwayDelegTransition = do
       checkStakeKeyNotRegistered stakeCred
       pure $ dState {dsUnified = registerStakeCredential stakeCred}
     ConwayUnRegCert stakeCred sMayRefund -> do
-      let mRDPair = UM.lookup stakeCred $ UM.RewDepUView dsUnified
+      let (mUMElem, umap) = UM.extractStakingCredential stakeCred dsUnified
           checkInvalidRefund = do
             SJust suppliedRefund <- Just sMayRefund
             -- we don't want to report invalid refund when stake credential is not registered:
-            UM.RDPair _ actualRefund <- mRDPair
+            UM.UMElem (SJust rd) _ _ _ <- mUMElem
             -- we return offending refund only when it doesn't match the expected one:
-            guard (suppliedRefund /= UM.fromCompact actualRefund)
+            guard (suppliedRefund /= UM.fromCompact (UM.rdDeposit rd))
             Just suppliedRefund
           checkStakeKeyHasZeroRewardBalance = do
-            UM.RDPair compactReward _ <- mRDPair
-            guard (compactReward /= mempty)
-            Just $ UM.fromCompact compactReward
+            UM.UMElem (SJust rd) _ _ _ <- mUMElem
+            guard (UM.rdReward rd /= mempty)
+            Just $ UM.fromCompact (UM.rdReward rd)
       failOnJust checkInvalidRefund IncorrectDepositDELEG
-      isJust mRDPair ?! StakeKeyNotRegisteredDELEG stakeCred
+      isJust mUMElem ?! StakeKeyNotRegisteredDELEG stakeCred
       failOnJust checkStakeKeyHasZeroRewardBalance StakeKeyHasNonZeroRewardAccountBalanceDELEG
-      pure $ dState {dsUnified = UM.domDeleteAll (Set.singleton stakeCred) dsUnified}
+      pure $ dState {dsUnified = umap}
     ConwayDelegCert stakeCred delegatee -> do
       checkStakeKeyIsRegistered stakeCred
       checkStakeDelegateeRegistered delegatee

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

@@ -19,7 +19,13 @@ module Cardano.Ledger.Shelley.Rules.Deleg (
 )
 where
 
-import Cardano.Ledger.BaseTypes (Globals (..), ShelleyBase, epochInfoPure, invalidKey)
+import Cardano.Ledger.BaseTypes (
+  Globals (..),
+  ShelleyBase,
+  StrictMaybe (..),
+  epochInfoPure,
+  invalidKey,
+ )
 import Cardano.Ledger.Binary (
   DecCBOR (..),
   EncCBOR (..),
@@ -60,9 +66,10 @@ import Cardano.Ledger.Slot (
   (*-),
   (+*),
  )
-import Cardano.Ledger.UMap (RDPair (..), UView (..), compactCoinOrError, fromCompact)
+import Cardano.Ledger.UMap (RDPair (..), UView (..), compactCoinOrError)
 import qualified Cardano.Ledger.UMap as UM
 import Control.DeepSeq
+import Control.Monad (guard)
 import Control.Monad.Trans.Reader (asks)
 import Control.SetAlgebra (eval, range, singleton, (∉), (∪), (⨃))
 import Control.State.Transition
@@ -276,19 +283,16 @@ delegationTransition = do
           u2 = RewDepUView u1 UM.∪ (hk, RDPair (UM.CompactCoin 0) deposit)
           u3 = PtrUView u2 UM.∪ (ptr, hk)
       pure (ds {dsUnified = u3})
-    UnRegTxCert hk -> do
-      -- note that pattern match is used instead of cwitness, as in the spec
+    UnRegTxCert cred -> do
       -- (hk ∈ dom (rewards ds))
-      UM.member hk (rewards ds) ?! StakeKeyNotRegisteredDELEG hk
-      let rewardCoin = rdReward <$> UM.lookup hk (rewards ds)
-      rewardCoin == Just mempty ?! StakeKeyNonZeroAccountBalanceDELEG (fromCompact <$> rewardCoin)
-
-      let u0 = dsUnified ds
-          u1 = Set.singleton hk UM.⋪ RewDepUView u0
-          u2 = Set.singleton hk UM.⋪ SPoolUView u1
-          u3 = PtrUView u2 UM.⋫ Set.singleton hk
-          u4 = ds {dsUnified = u3}
-      pure u4
+      let (mUMElem, umap) = UM.extractStakingCredential cred (dsUnified ds)
+          checkStakeKeyHasZeroRewardBalance = do
+            UM.UMElem (SJust rd) _ _ _ <- mUMElem
+            guard (UM.rdReward rd /= mempty)
+            Just $ UM.fromCompact (UM.rdReward rd)
+      isJust mUMElem ?! StakeKeyNotRegisteredDELEG cred
+      failOnJust checkStakeKeyHasZeroRewardBalance (StakeKeyNonZeroAccountBalanceDELEG . Just)
+      pure $ ds {dsUnified = umap}
     DelegStakeTxCert hk dpool -> do
       -- note that pattern match is used instead of cwitness and dpool, as in the spec
       -- (hk ∈ dom (rewards ds))

libs/cardano-ledger-core/CHANGELOG.md

@@ -4,6 +4,7 @@
 
 * Add `credKeyHash` to `Credential`
 * Remove `maxMajorPV` from `Globals`
+* Add `deleteStakingCredential` and `extractStakingCredential` to `UMap` module.
 
 ### `testlib`
 

libs/cardano-ledger-core/src/Cardano/Ledger/UMap.hs

@@ -109,6 +109,8 @@ module Cardano.Ledger.UMap (
   findWithDefault,
   size,
   domDeleteAll,
+  deleteStakingCredential,
+  extractStakingCredential,
 )
 where
 
@@ -127,7 +129,7 @@ import qualified Data.Aeson as Aeson
 import Data.Foldable (Foldable (..))
 import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
-import Data.MapExtras (intersectDomPLeft)
+import Data.MapExtras as MapExtras (extract, intersectDomPLeft)
 import Data.Maybe as Maybe (fromMaybe, isNothing, mapMaybe)
 import Data.Maybe.Strict (StrictMaybe (..))
 import Data.Set (Set)
@@ -944,11 +946,25 @@ domDelete = (⋪)
 -- | Delete the stake credentials in the domain and all associated ranges from the `UMap`
 -- This can be expensive when there are many pointers associated with the credential.
 domDeleteAll :: Set (Credential 'Staking c) -> UMap c -> UMap c
-domDeleteAll ks UMap {umElems, umPtrs} =
-  UMap
-    { umElems = Map.withoutKeys umElems ks
-    , umPtrs = Map.filter (`Set.notMember` ks) umPtrs
-    }
+domDeleteAll ks umap = Set.foldr' deleteStakingCredential umap ks
+
+-- | Completely remove the staking credential from the UMap, including all associated
+-- pointers.
+deleteStakingCredential :: Credential 'Staking c -> UMap c -> UMap c
+deleteStakingCredential cred = snd . extractStakingCredential cred
+
+-- | Just like `deleteStakingCredential`, but also returned the removed element.
+extractStakingCredential :: Credential 'Staking c -> UMap c -> (Maybe (UMElem c), UMap c)
+extractStakingCredential cred umap@UMap {umElems, umPtrs} =
+  case MapExtras.extract cred umElems of
+    (Nothing, _) -> (Nothing, umap)
+    (e@(Just (UMElem _ ptrs _ _)), umElems') ->
+      ( e
+      , UMap
+          { umElems = umElems'
+          , umPtrs = umPtrs `Map.withoutKeys` ptrs
+          }
+      )
 
 -- | Delete all elements in the given `Set` from the range of the given map-like `UView`.
 -- This is slow for SPoolUView, RewDepUView, and DReps UViews, better hope the sets are small