Add StakeCredentials and with all Maps unrolled:

* two helper functions: `toStakeCredentials` and `domRestrictedStakeCredentials`
* generators: `genValidUMapWithCreds` and `uniformSubset`

libs/cardano-ledger-core/CHANGELOG.md

@@ -49,6 +49,12 @@
 * Add new key roles: `CommitteeColdKey` and `CommitteeHotKey`
 * Remove `ConstitutionalDelegCert`. Instead it now lives in `cardano-ledger-shelley` as
   `GenesisDelegCert`
+* Add `StakeCredentials` and two helper functions: `toStakeCredentials` and
+  `domRestrictedStakeCredentials`
+
+### `testlib`
+
+* Add `genValidUMapWithCreds` and `uniformSubset`
 
 ## 1.2.0.0
 

libs/cardano-ledger-core/cardano-ledger-core.cabal

@@ -152,6 +152,7 @@ library testlib
         nothunks,
         primitive,
         QuickCheck,
+        random >= 1.2,
         text,
         vector-map
 

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

@@ -39,6 +39,11 @@ module Cardano.Ledger.UMap (
   empty,
   umInvariant,
 
+  -- * StakeCredentials
+  StakeCredentials (..),
+  toStakeCredentials,
+  domRestrictedStakeCredentials,
+
   -- * `UView` and its components
   UView (..),
   rewDepUView,
@@ -379,6 +384,20 @@ data UMap c = UMap
   }
   deriving (Show, Eq, Generic, NoThunks, NFData)
 
+-- | All maps unrolled. It is important to note that all fields are lazy, because
+-- conversion from UMap can be expensive, thus only fields that are forced will incur that
+-- conversion overhead.
+data StakeCredentials c = StakeCredentials
+  { scRewards :: Map (Credential 'Staking c) Coin
+  , scDeposits :: Map (Credential 'Staking c) Coin
+  , scSPools :: Map (Credential 'Staking c) (KeyHash 'StakePool c)
+  , scDReps :: Map (Credential 'Staking c) (Credential 'Voting c)
+  , scPtrs :: Map Ptr (Credential 'Staking c)
+  , scPtrsInverse :: Map (Credential 'Staking c) (Set Ptr)
+  -- ^ There will be no empty sets in the range
+  }
+  deriving (Show, Eq, Generic, NoThunks, NFData)
+
 instance ToExpr (UMap c)
 
 instance Crypto c => ToJSON (UMap c) where
@@ -510,50 +529,74 @@ unUnifyToVMap uview = case uview of
     toDRep (key, t) = (,) key <$> umElemDRep t
 
 -- | Extract a reward-deposit pairs `Map` from a 'UMap'
-rdPairMap :: UMap c -> Map.Map (Credential 'Staking c) RDPair
+rdPairMap :: UMap c -> Map (Credential 'Staking c) RDPair
 rdPairMap x = unUnify $ RewDepUView x
 
 -- | Extract a rewards `Map` from a 'UMap'
-rewardMap :: UMap c -> Map.Map (Credential 'Staking c) Coin
+rewardMap :: UMap c -> Map (Credential 'Staking c) Coin
 rewardMap x = Map.map (fromCompact . rdReward) $ unUnify $ RewDepUView x
 
 -- | Extract a compact rewards `Map` from a 'UMap'
-compactRewardMap :: UMap c -> Map.Map (Credential 'Staking c) (CompactForm Coin)
+compactRewardMap :: UMap c -> Map (Credential 'Staking c) (CompactForm Coin)
 compactRewardMap x = Map.map rdReward $ unUnify $ RewDepUView x
 
 -- | Extract a deposits `Map` from a 'UMap'
-depositMap :: UMap c -> Map.Map (Credential 'Staking c) Coin
+depositMap :: UMap c -> Map (Credential 'Staking c) Coin
 depositMap x = Map.map (fromCompact . rdDeposit) $ unUnify $ RewDepUView x
 
 -- | Extract a pointers `Map` from a 'UMap'
-ptrMap :: UMap c -> Map.Map Ptr (Credential 'Staking c)
+ptrMap :: UMap c -> Map Ptr (Credential 'Staking c)
 ptrMap x = unUnify $ PtrUView x
 
 -- | Extract a pointers `Map` from a 'UMap'
-revPtrMap :: UMap c -> Map.Map (Credential 'Staking c) (Set Ptr)
+revPtrMap :: UMap c -> Map (Credential 'Staking c) (Set Ptr)
 revPtrMap UMap {umElems} =
   Map.foldlWithKey'
     (\ans k (UMElem _ ptrSet _ _) -> if Set.null ptrSet then ans else Map.insert k ptrSet ans)
     Map.empty
     umElems
 
 -- | Extract a stake pool delegations `Map` from a 'UMap'
-sPoolMap :: UMap c -> Map.Map (Credential 'Staking c) (KeyHash 'StakePool c)
+sPoolMap :: UMap c -> Map (Credential 'Staking c) (KeyHash 'StakePool c)
 sPoolMap x = unUnify $ SPoolUView x
 
 -- | Extract a delegated-representatives `Map` from a 'UMap'
-dRepMap :: UMap c -> Map.Map (Credential 'Staking c) (Credential 'Voting c)
+dRepMap :: UMap c -> Map (Credential 'Staking c) (Credential 'Voting c)
 dRepMap x = unUnify $ DRepUView x
 
 -- | Extract a domain-restricted `Map` of a `UMap`.
 -- If `Set k` is small this should be efficient.
-domRestrictedMap :: Set k -> UView c k v -> Map.Map k v
+domRestrictedMap :: Set k -> UView c k v -> Map k v
 domRestrictedMap setk = \case
   RewDepUView UMap {umElems} -> Map.mapMaybe umElemRDPair (Map.restrictKeys umElems setk)
   PtrUView UMap {umPtrs} -> Map.restrictKeys umPtrs setk
   SPoolUView UMap {umElems} -> Map.mapMaybe umElemSPool (Map.restrictKeys umElems setk)
   DRepUView UMap {umElems} -> Map.mapMaybe umElemDRep (Map.restrictKeys umElems setk)
 
+toStakeCredentials :: UMap c -> StakeCredentials c
+toStakeCredentials umap =
+  StakeCredentials
+    { scRewards = rewardMap umap
+    , scDeposits = depositMap umap
+    , scSPools = sPoolMap umap
+    , scDReps = dRepMap umap
+    , scPtrs = ptrMap umap
+    , scPtrsInverse = revPtrMap umap
+    }
+
+domRestrictedStakeCredentials :: Set (Credential 'Staking c) -> UMap c -> StakeCredentials c
+domRestrictedStakeCredentials setk UMap {umElems, umPtrs} =
+  let umElems' = Map.restrictKeys umElems setk
+      ptrs = Map.mapMaybe umElemPtrs umElems'
+   in StakeCredentials
+        { scRewards = Map.mapMaybe (\e -> fromCompact . rdReward <$> umElemRDPair e) umElems'
+        , scDeposits = Map.mapMaybe (\e -> fromCompact . rdDeposit <$> umElemRDPair e) umElems'
+        , scSPools = Map.mapMaybe umElemSPool umElems'
+        , scDReps = Map.mapMaybe umElemDRep umElems'
+        , scPtrs = umPtrs `Map.restrictKeys` fold ptrs
+        , scPtrsInverse = ptrs
+        }
+
 -- | All `View`s are `Foldable`
 instance Foldable (UView c k) where
   foldMap f = \case

libs/cardano-ledger-core/test/Test/Cardano/Ledger/UMapSpec.hs

@@ -8,11 +8,13 @@ module Test.Cardano.Ledger.UMapSpec where
 
 import Cardano.Ledger.BaseTypes (StrictMaybe (SJust, SNothing))
 import Cardano.Ledger.Coin (Coin, CompactForm)
+import Cardano.Ledger.Compactible (fromCompact)
 import Cardano.Ledger.Credential (Credential, Ptr)
 import Cardano.Ledger.Crypto (StandardCrypto)
 import Cardano.Ledger.Keys (KeyHash, KeyRole (StakePool, Staking, Voting))
 import Cardano.Ledger.UMap (
   RDPair (RDPair, rdReward),
+  StakeCredentials (..),
   UMElem (UMElem),
   UMap (UMap, umElems, umPtrs),
   UView (DRepUView, PtrUView, RewDepUView, SPoolUView),
@@ -21,7 +23,10 @@ import Cardano.Ledger.UMap (
   dRepMap,
   delete,
   delete',
+  depositMap,
   domRestrict,
+  domRestrictedMap,
+  domRestrictedStakeCredentials,
   domain,
   empty,
   insert,
@@ -30,9 +35,13 @@ import Cardano.Ledger.UMap (
   nullUView,
   ptrMap,
   range,
+  rdDeposit,
   rdPairMap,
+  revPtrMap,
+  rewardMap,
   sPoolMap,
   size,
+  toStakeCredentials,
   umInvariant,
   unUView,
   unUnify,
@@ -45,6 +54,7 @@ import Cardano.Ledger.UMap (
  )
 import qualified Cardano.Ledger.UMap as UMap (lookup)
 import Control.Exception (assert)
+import Data.Foldable (Foldable (fold))
 import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
 import qualified Data.Set as Set
@@ -57,6 +67,8 @@ import Test.Cardano.Ledger.Core.Arbitrary (
   genInvariantNonEmpty,
   genRightPreferenceUMap,
   genValidTuples,
+  genValidUMap,
+  genValidUMapWithCreds,
  )
 
 data Action
@@ -532,3 +544,32 @@ spec = do
           ( \(umap, m) ->
               RewDepUView umap ∪+ (rdReward <$> m) === RewDepUView umap `oldUnionRewAgg` (rdReward <$> m)
           )
+    describe "StakeCredentials" $ do
+      prop "toStakeCredentials as full domRestrictedStakeCredentials" $
+        forAll genValidUMap $ \umap ->
+          toStakeCredentials umap
+            `shouldBe` domRestrictedStakeCredentials (Map.keysSet (umElems umap)) umap
+
+      prop "domRestrictedStakeCredentials" $ forAll genValidUMapWithCreds $ \(umap, creds) ->
+        domRestrictedStakeCredentials creds umap
+          `shouldBe` StakeCredentials
+            { scRewards = rewardMap umap `Map.restrictKeys` creds
+            , scDeposits = depositMap umap `Map.restrictKeys` creds
+            , scSPools = sPoolMap umap `Map.restrictKeys` creds
+            , scDReps = dRepMap umap `Map.restrictKeys` creds
+            , scPtrs = Map.filter (`Set.member` creds) $ ptrMap umap
+            , scPtrsInverse = revPtrMap umap `Map.restrictKeys` creds
+            }
+      prop "domRestrictedStakeCredentials with domRestrictedMap" $
+        forAll genValidUMapWithCreds $ \(umap, creds) ->
+          let rdmap = domRestrictedMap creds (RewDepUView umap)
+              ptrs = revPtrMap umap `Map.restrictKeys` creds
+           in domRestrictedStakeCredentials creds umap
+                `shouldBe` StakeCredentials
+                  { scRewards = Map.map (fromCompact . rdReward) rdmap
+                  , scDeposits = Map.map (fromCompact . rdDeposit) rdmap
+                  , scSPools = domRestrictedMap creds (SPoolUView umap)
+                  , scDReps = domRestrictedMap creds (DRepUView umap)
+                  , scPtrs = domRestrictedMap (fold ptrs) (PtrUView umap)
+                  , scPtrsInverse = ptrs
+                  }

libs/cardano-ledger-core/testlib/Test/Cardano/Ledger/Core/Arbitrary.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE FlexibleInstances #-}
@@ -18,6 +19,7 @@ module Test.Cardano.Ledger.Core.Arbitrary (
   genBadPtr,
   genValidUMap,
   genValidUMapNonEmpty,
+  genValidUMapWithCreds,
   genValidTuples,
   genValidTuplesNonEmpty,
   genInvariantNonEmpty,
@@ -26,6 +28,11 @@ module Test.Cardano.Ledger.Core.Arbitrary (
   genInsertDeleteRoundtripPtr,
   genInsertDeleteRoundtripSPool,
   genInsertDeleteRoundtripDRep,
+
+  -- * Utils
+
+  -- | Will need to find a better home in the future
+  uniformSubset,
 )
 where
 
@@ -91,16 +98,21 @@ import Cardano.Ledger.UMap (RDPair (..), UMElem (UMElem), UMap (UMap, umElems, u
 import Cardano.Ledger.UTxO (UTxO (..))
 import Control.Monad.Identity (Identity)
 import Data.GenValidity
+import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
 import Data.Ratio ((%))
+import Data.Set (Set)
+import qualified Data.Set as Set
 import qualified Data.Text as T
 import Data.Typeable
 import qualified Data.VMap as VMap
 import Data.Word (Word16, Word32, Word64)
 import GHC.Stack
 import Generic.Random (genericArbitraryU)
+import System.Random.Stateful (StatefulGen, uniformRM)
 import qualified Test.Cardano.Chain.Common.Gen as Byron
 import Test.Cardano.Ledger.Binary.Arbitrary
+import Test.Cardano.Ledger.Binary.Random (QC (..))
 import Test.Cardano.Ledger.Core.KeyPair (KeyPair (..))
 import Test.Cardano.Ledger.Core.Utils (unsafeBoundRational)
 import Test.QuickCheck
@@ -480,10 +492,10 @@ instance Crypto c => Arbitrary (UMap c) where
 -- | Used for testing UMap operations
 genValidTuples ::
   Gen
-    ( Map.Map (Credential 'Staking StandardCrypto) RDPair
-    , Map.Map Ptr (Credential 'Staking StandardCrypto)
-    , Map.Map (Credential 'Staking StandardCrypto) (KeyHash 'StakePool StandardCrypto)
-    , Map.Map (Credential 'Staking StandardCrypto) (Credential 'Voting StandardCrypto)
+    ( Map (Credential 'Staking StandardCrypto) RDPair
+    , Map Ptr (Credential 'Staking StandardCrypto)
+    , Map (Credential 'Staking StandardCrypto) (KeyHash 'StakePool StandardCrypto)
+    , Map (Credential 'Staking StandardCrypto) (Credential 'Voting StandardCrypto)
     )
 genValidTuples = scale (* 2) $ do
   creds :: [Credential 'Staking StandardCrypto] <- arbitrary
@@ -501,10 +513,10 @@ genValidTuples = scale (* 2) $ do
 
 genValidTuplesNonEmpty ::
   Gen
-    ( Map.Map (Credential 'Staking StandardCrypto) RDPair
-    , Map.Map Ptr (Credential 'Staking StandardCrypto)
-    , Map.Map (Credential 'Staking StandardCrypto) (KeyHash 'StakePool StandardCrypto)
-    , Map.Map (Credential 'Staking StandardCrypto) (Credential 'Voting StandardCrypto)
+    ( Map (Credential 'Staking StandardCrypto) RDPair
+    , Map Ptr (Credential 'Staking StandardCrypto)
+    , Map (Credential 'Staking StandardCrypto) (KeyHash 'StakePool StandardCrypto)
+    , Map (Credential 'Staking StandardCrypto) (Credential 'Voting StandardCrypto)
     )
 genValidTuplesNonEmpty = scale (* 2) $ do
   Positive nCreds <- arbitrary
@@ -531,7 +543,33 @@ genValidUMapNonEmpty = do
   (rdPairs, ptrs, sPools, dReps) <- genValidTuplesNonEmpty
   pure $ unify rdPairs ptrs sPools dReps
 
-genExcludingKey :: (Ord k, Arbitrary k) => Map.Map k a -> Gen k
+uniformSubset ::
+  (StatefulGen g m, Ord k) =>
+  -- | Size of the subset. If supplied will be clamped to @[0, Set.size s]@ interval,
+  -- otherwise will be generated randomly.
+  Maybe Int ->
+  Set k ->
+  g ->
+  m (Set k)
+uniformSubset mSubSetSize inputSet gen = do
+  subSetSize <- case mSubSetSize of
+    Nothing -> uniformRM (0, Set.size inputSet) gen
+    Just n -> pure $ max 0 $ min (Set.size inputSet) n
+  go inputSet Set.empty subSetSize
+  where
+    go !s !acc !i
+      | i <= 0 = pure acc
+      | otherwise = do
+          ix <- uniformRM (0, Set.size s - 1) gen
+          go (Set.insert (Set.elemAt ix s) acc) (Set.deleteAt ix s) (i - 1)
+
+genValidUMapWithCreds :: Gen (UMap StandardCrypto, Set (Credential 'Staking StandardCrypto))
+genValidUMapWithCreds = do
+  umap <- genValidUMap
+  creds <- uniformSubset Nothing (Map.keysSet $ umElems umap) QC
+  pure (umap, creds)
+
+genExcludingKey :: (Ord k, Arbitrary k) => Map k a -> Gen k
 genExcludingKey ks = do
   k <- arbitrary
   if k `Map.member` ks
@@ -591,7 +629,7 @@ genInvariantNonEmpty = do
       ]
   pure (cred, ptr, umap)
 
-genRightPreferenceUMap :: Gen (UMap StandardCrypto, Map.Map (Credential 'Staking StandardCrypto) RDPair)
+genRightPreferenceUMap :: Gen (UMap StandardCrypto, Map (Credential 'Staking StandardCrypto) RDPair)
 genRightPreferenceUMap = do
   umap <- genValidUMap
   let rdMap = unUnify $ RewDepUView umap