Merge pull request #5264 from IntersectMBO/aniketd/pstate-query

Add queryPoolState state query

Author: lehins

libs/cardano-ledger-api/CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## 1.12.0.0
 
+* Add `queryPoolState` and bring back `queryPoolParameters` state query.
 * Add `queryDRepDelegations` state query
 * Remove `filterStakePoolDelegsAndRewards` as unnecessary. Use `queryStakePoolDelegsAndRewards` instead
 * Expose `binaryUpgradeTx`, `binaryUpgradeTxBody`, `binaryUpgradeTxWits`, `binaryUpgradeTxAuxData`, `upgradeTx`, `upgradeTxBody`, `upgradeTxWits`, `upgradeTxAuxData`

libs/cardano-ledger-api/src/Cardano/Ledger/Api/State/Query.hs

@@ -1,8 +1,13 @@
 {-# LANGUAGE DataKinds #-}
+{-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE GADTs #-}
+{-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE UndecidableInstances #-}
 
 module Cardano.Ledger.Api.State.Query (
   -- * @GetFilteredDelegationsAndRewardAccounts@
@@ -58,10 +63,16 @@ module Cardano.Ledger.Api.State.Query (
   -- * @GetRatifyState@
   queryRatifyState,
 
-  -- * @GetStakePoolDefaultVote
+  -- * @GetStakePoolDefaultVote@
   queryStakePoolDefaultVote,
   DefaultVote (..),
 
+  -- * @GetPoolState@
+  queryPoolParameters,
+  queryPoolState,
+  QueryPoolStateResult (..),
+  mkQueryPoolStateResult,
+
   -- * For testing
   getNextEpochCommitteeMembers,
 ) where
@@ -74,6 +85,7 @@ import Cardano.Ledger.Api.State.Query.CommitteeMembersState (
   NextEpochChange (..),
  )
 import Cardano.Ledger.BaseTypes (EpochNo, strictMaybeToMaybe)
+import Cardano.Ledger.Binary
 import Cardano.Ledger.Coin (Coin (..), CompactForm (..))
 import Cardano.Ledger.Compactible (fromCompact)
 import Cardano.Ledger.Conway.Governance (
@@ -426,3 +438,61 @@ queryStakePoolDefaultVote ::
 queryStakePoolDefaultVote nes poolId =
   defaultStakePoolVote poolId (nes ^. nesEsL . epochStateStakePoolsL) $
     nes ^. nesEsL . esLStateL . lsCertStateL . certDStateL . accountsL
+
+-- | Used only for the `queryPoolState` query. This resembles the older way of
+-- representing StakePoolState in Ledger.
+data QueryPoolStateResult = QueryPoolStateResult
+  { qpsrStakePoolParams :: !(Map (KeyHash 'StakePool) PoolParams)
+  , qpsrFutureStakePoolParams :: !(Map (KeyHash 'StakePool) PoolParams)
+  , qpsrRetiring :: !(Map (KeyHash 'StakePool) EpochNo)
+  , qpsrDeposits :: !(Map (KeyHash 'StakePool) Coin)
+  }
+  deriving (Show, Eq)
+
+instance EncCBOR QueryPoolStateResult where
+  encCBOR (QueryPoolStateResult a b c d) =
+    encodeListLen 4 <> encCBOR a <> encCBOR b <> encCBOR c <> encCBOR d
+
+instance DecCBOR QueryPoolStateResult where
+  decCBOR = decodeRecordNamed "QueryPoolStateResult" (const 4) $ do
+    qpsrStakePoolParams <- decCBOR
+    qpsrFutureStakePoolParams <- decCBOR
+    qpsrRetiring <- decCBOR
+    qpsrDeposits <- decCBOR
+    pure
+      QueryPoolStateResult {qpsrStakePoolParams, qpsrFutureStakePoolParams, qpsrRetiring, qpsrDeposits}
+
+mkQueryPoolStateResult ::
+  (forall x. Map.Map (KeyHash 'StakePool) x -> Map.Map (KeyHash 'StakePool) x) ->
+  PState era ->
+  QueryPoolStateResult
+mkQueryPoolStateResult f ps =
+  QueryPoolStateResult
+    { qpsrStakePoolParams = Map.mapWithKey stakePoolStateToPoolParams restrictedStakePools
+    , qpsrFutureStakePoolParams = Map.mapWithKey stakePoolStateToPoolParams restrictedStakePools
+    , qpsrRetiring = f $ psRetiring ps
+    , qpsrDeposits = Map.map (fromCompact . spsDeposit) restrictedStakePools
+    }
+  where
+    restrictedStakePools = f $ psStakePools ps
+
+-- | Query the QueryPoolStateResult. This is slightly different from the internal
+-- representation used by Ledger and is intended to resemble how the internal
+-- representation used to be.
+queryPoolState ::
+  EraCertState era => NewEpochState era -> Maybe (Set (KeyHash 'StakePool)) -> QueryPoolStateResult
+queryPoolState nes mPoolKeys =
+  let pstate = nes ^. nesEsL . esLStateL . lsCertStateL . certPStateL
+      f :: forall x. Map.Map (KeyHash 'StakePool) x -> Map.Map (KeyHash 'StakePool) x
+      f = case mPoolKeys of
+        Nothing -> id
+        Just keys -> (`Map.restrictKeys` keys)
+   in mkQueryPoolStateResult f pstate
+
+-- | Query the current PoolParams.
+queryPoolParameters ::
+  EraCertState era =>
+  NewEpochState era -> Set (KeyHash 'StakePool) -> Map (KeyHash 'StakePool) PoolParams
+queryPoolParameters nes poolKeys =
+  let pools = nes ^. nesEsL . esLStateL . lsCertStateL . certPStateL . psStakePoolsL
+   in Map.mapWithKey stakePoolStateToPoolParams $ Map.restrictKeys pools poolKeys

libs/cardano-ledger-api/test/Test/Cardano/Ledger/Api/State/QuerySpec.hs

@@ -16,6 +16,7 @@ import Cardano.Ledger.Api.State.Query (
   HotCredAuthStatus (..),
   MemberStatus (..),
   NextEpochChange (..),
+  QueryPoolStateResult,
   getNextEpochCommitteeMembers,
   queryCommitteeMembersState,
   queryStakePoolDelegsAndRewards,
@@ -54,12 +55,22 @@ import Test.Cardano.Ledger.Common
 import Test.Cardano.Ledger.Conway.Arbitrary ()
 import Test.Cardano.Ledger.Conway.Era (ShelleyEraTest)
 import Test.Cardano.Ledger.Core.Arbitrary (genValidUMapWithCreds)
+import Test.Cardano.Ledger.Core.Binary.RoundTrip (roundTripEraExpectation)
 import Test.Cardano.Ledger.Era (accountsFromUMap)
 import Test.Cardano.Ledger.Shelley.Arbitrary ()
 import Test.Cardano.Slotting.Numeric ()
 
 spec :: Spec
 spec = do
+  describe "API Types" $ do
+    describe "Roundtrip" $ do
+      prop "Shelley" $ roundTripEraExpectation @ShelleyEra @QueryPoolStateResult
+      prop "Allegra" $ roundTripEraExpectation @AllegraEra @QueryPoolStateResult
+      prop "Mary" $ roundTripEraExpectation @MaryEra @QueryPoolStateResult
+      prop "Alonzo" $ roundTripEraExpectation @AlonzoEra @QueryPoolStateResult
+      prop "Babbage" $ roundTripEraExpectation @BabbageEra @QueryPoolStateResult
+      prop "Conway" $ roundTripEraExpectation @ConwayEra @QueryPoolStateResult
+      prop "Dijkstra" $ roundTripEraExpectation @DijkstraEra @QueryPoolStateResult
   queryStakePoolDelegsAndRewardsSpec @ShelleyEra
   queryStakePoolDelegsAndRewardsSpec @AllegraEra
   queryStakePoolDelegsAndRewardsSpec @MaryEra

libs/cardano-ledger-api/testlib/Test/Cardano/Ledger/Api/Arbitrary.hs

@@ -2,9 +2,12 @@
 
 module Test.Cardano.Ledger.Api.Arbitrary () where
 
-import Cardano.Ledger.Api.State.Query (MemberStatus)
+import Cardano.Ledger.Api.State.Query (MemberStatus, QueryPoolStateResult (..))
 import Test.Cardano.Ledger.Common
 import Test.Cardano.Ledger.Dijkstra.Arbitrary ()
 
 instance Arbitrary MemberStatus where
   arbitrary = arbitraryBoundedEnum
+
+instance Arbitrary QueryPoolStateResult where
+  arbitrary = QueryPoolStateResult <$> arbitrary <*> arbitrary <*> arbitrary <*> arbitrary

libs/cardano-ledger-api/testlib/Test/Cardano/Ledger/Api/Upgrade.hs

@@ -17,6 +17,7 @@ import Cardano.Ledger.Core
 import Cardano.Ledger.MemoBytes (EqRaw (eqRaw))
 import Data.Default (Default (def))
 import qualified Prettyprinter as Pretty
+import Test.Cardano.Ledger.Api.Arbitrary ()
 import Test.Cardano.Ledger.Binary.RoundTrip
 import Test.Cardano.Ledger.Common
 import Test.Cardano.Ledger.Core.Arbitrary ()