Optimise query stake-snapshot command

cabal.project

@@ -263,8 +263,8 @@ source-repository-package
 source-repository-package
   type: git
   location: https://github.com/input-output-hk/ouroboros-network
-  tag: c764553561bed8978d2c6753d1608dc65449617a
-  --sha256: 0hdh7xdrvxw943r6qr0xr4kwszindh5mnsn1lww6qdnxnmn7wcsc
+  tag: 01549ff9eaa895a9e728dbdd200ec95d5767b9e9
+  --sha256: 19jwqvvdxdhcsshq29qlrq3y13d2yj8vc5pxngdl8izkhgcd8p2i
   subdir:
     monoidal-synchronisation
     network-mux

cardano-api/src/Cardano/Api/KeysShelley.hs

@@ -1184,7 +1184,9 @@ instance SerialiseAsBech32 (SigningKey StakePoolKey) where
     bech32PrefixesPermitted _ = ["pool_sk"]
 
 newtype instance Hash StakePoolKey =
-    StakePoolKeyHash (Shelley.KeyHash Shelley.StakePool StandardCrypto)
+    StakePoolKeyHash
+    { unStakePoolKeyHash :: Shelley.KeyHash Shelley.StakePool StandardCrypto
+    }
   deriving stock (Eq, Ord)
   deriving (Show, IsString) via UsingRawBytesHex (Hash StakePoolKey)
   deriving (ToCBOR, FromCBOR) via UsingRawBytes (Hash StakePoolKey)

cardano-api/src/Cardano/Api/Orphans.hs

@@ -4,6 +4,7 @@
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE StandaloneDeriving #-}
 {-# LANGUAGE TypeApplications #-}
@@ -16,7 +17,7 @@ module Cardano.Api.Orphans () where
 
 import           Prelude
 
-import           Data.Aeson (FromJSON (..), ToJSON (..), object, (.=))
+import           Data.Aeson (FromJSON (..), ToJSON (..), object, pairs, (.=))
 import qualified Data.Aeson as Aeson
 import           Data.Aeson.Types (ToJSONKey (..), toJSONKeyText)
 import           Data.BiMap (BiMap (..), Bimap)
@@ -56,12 +57,13 @@ import qualified Cardano.Ledger.Shelley.Constraints as Shelley
 import qualified Cardano.Ledger.Shelley.EpochBoundary as ShelleyEpoch
 import qualified Cardano.Ledger.Shelley.LedgerState as ShelleyLedger
 import           Cardano.Ledger.Shelley.PParams (PParamsUpdate)
-import qualified Cardano.Ledger.Shelley.RewardUpdate as Shelley
 import qualified Cardano.Ledger.Shelley.Rewards as Shelley
+import qualified Cardano.Ledger.Shelley.RewardUpdate as Shelley
 import qualified Ouroboros.Consensus.Shelley.Eras as Consensus
 
 import           Cardano.Api.Script
 import           Cardano.Api.SerialiseRaw (serialiseToRawBytesHexText)
+import qualified Ouroboros.Consensus.Shelley.Ledger.Query as Consensus
 
 -- Orphan instances involved in the JSON output of the API queries.
 -- We will remove/replace these as we provide more API wrapper types
@@ -418,3 +420,40 @@ instance Crypto.Crypto crypto => ToJSON (VMap VB VB (Shelley.KeyHash    'Shelley
 
 instance Crypto.Crypto crypto => ToJSON (VMap VB VP (Shelley.Credential 'Shelley.Staking   crypto) (Shelley.CompactForm Shelley.Coin)) where
   toJSON = toJSON . fmap fromCompact . VMap.toMap
+
+-----
+
+instance ToJSON (Consensus.StakeSnapshot crypto) where
+  toJSON
+    Consensus.StakeSnapshot
+    { Consensus.sMarkPool
+    , Consensus.sSetPool
+    , Consensus.sGoPool
+    , Consensus.sMarkTotal
+    , Consensus.sSetTotal
+    , Consensus.sGoTotal
+    } = object
+    [ "poolStakeMark" .= sMarkPool
+    , "poolStakeSet" .= sSetPool
+    , "poolStakeGo" .= sGoPool
+    , "activeStakeMark" .= sMarkTotal
+    , "activeStakeSet" .= sSetTotal
+    , "activeStakeGo" .= sGoTotal
+    ]
+
+  toEncoding
+    Consensus.StakeSnapshot
+    { Consensus.sMarkPool
+    , Consensus.sSetPool
+    , Consensus.sGoPool
+    , Consensus.sMarkTotal
+    , Consensus.sSetTotal
+    , Consensus.sGoTotal
+    } = pairs $ mconcat
+    [ "poolStakeMark" .= sMarkPool
+    , "poolStakeSet" .= sSetPool
+    , "poolStakeGo" .= sGoPool
+    , "activeStakeMark" .= sMarkTotal
+    , "activeStakeSet" .= sSetTotal
+    , "activeStakeGo" .= sGoTotal
+    ]

cardano-api/src/Cardano/Api/Query.hs

@@ -47,6 +47,10 @@ module Cardano.Api.Query (
     PoolState(..),
     decodePoolState,
 
+    SerialisedStakeSnapshot(..),
+    StakeSnapshot(..),
+    decodeStakeSnapshot,
+
     EraHistory(..),
     SystemStart(..),
 
@@ -242,6 +246,10 @@ data QueryInShelleyBasedEra era result where
     :: Maybe (Set PoolId)
     -> QueryInShelleyBasedEra era (SerialisedPoolState era)
 
+  QueryStakeSnapshot
+    :: PoolId
+    -> QueryInShelleyBasedEra era (SerialisedStakeSnapshot era)
+
 deriving instance Show (QueryInShelleyBasedEra era result)
 
 
@@ -377,6 +385,18 @@ decodePoolState
   -> Either DecoderError (PoolState era)
 decodePoolState (SerialisedPoolState (Serialised ls)) = PoolState <$> decodeFull ls
 
+newtype SerialisedStakeSnapshot era
+  = SerialisedStakeSnapshot (Serialised (Consensus.StakeSnapshot (Ledger.Crypto (ShelleyLedgerEra era))))
+
+newtype StakeSnapshot era = StakeSnapshot (Consensus.StakeSnapshot (Ledger.Crypto (ShelleyLedgerEra era)))
+
+decodeStakeSnapshot
+  :: forall era. ()
+  => FromCBOR (Consensus.StakeSnapshot (Ledger.Crypto (ShelleyLedgerEra era)))
+  => SerialisedStakeSnapshot era
+  -> Either DecoderError (StakeSnapshot era)
+decodeStakeSnapshot (SerialisedStakeSnapshot (Serialised ls)) = StakeSnapshot <$> decodeFull ls
+
 toShelleyAddrSet :: CardanoEra era
                  -> Set AddressAny
                  -> Set (Shelley.Addr Consensus.StandardCrypto)
@@ -545,7 +565,7 @@ toConsensusQueryShelleyBased erainmode (QueryStakePoolParameters poolids) =
     Some (consensusQueryInEraInMode erainmode (Consensus.GetStakePoolParams poolids'))
   where
     poolids' :: Set (Shelley.KeyHash Shelley.StakePool Consensus.StandardCrypto)
-    poolids' = Set.map (\(StakePoolKeyHash kh) -> kh) poolids
+    poolids' = Set.map unStakePoolKeyHash poolids
 
 toConsensusQueryShelleyBased erainmode QueryDebugLedgerState =
     Some (consensusQueryInEraInMode erainmode (Consensus.GetCBOR Consensus.DebugNewEpochState))
@@ -557,10 +577,10 @@ toConsensusQueryShelleyBased erainmode QueryCurrentEpochState =
     Some (consensusQueryInEraInMode erainmode (Consensus.GetCBOR Consensus.DebugEpochState))
 
 toConsensusQueryShelleyBased erainmode (QueryPoolState poolIds) =
-    Some (consensusQueryInEraInMode erainmode (Consensus.GetCBOR (Consensus.GetPoolState (getPoolIds <$> poolIds))))
-  where
-    getPoolIds :: Set PoolId -> Set (Shelley.KeyHash Shelley.StakePool Consensus.StandardCrypto)
-    getPoolIds = Set.map (\(StakePoolKeyHash kh) -> kh)
+    Some (consensusQueryInEraInMode erainmode (Consensus.GetCBOR (Consensus.GetPoolState (Set.map unStakePoolKeyHash <$> poolIds))))
+
+toConsensusQueryShelleyBased erainmode (QueryStakeSnapshot poolId) =
+    Some (consensusQueryInEraInMode erainmode (Consensus.GetCBOR (Consensus.GetStakeSnapshot (unStakePoolKeyHash poolId))))
 
 consensusQueryInEraInMode
   :: forall era mode erablock modeblock result result' xs.
@@ -797,6 +817,11 @@ fromConsensusQueryResultShelleyBased _ QueryPoolState{} q' r' =
     Consensus.GetCBOR Consensus.GetPoolState {} -> SerialisedPoolState r'
     _                                           -> fromConsensusQueryResultMismatch
 
+fromConsensusQueryResultShelleyBased _ QueryStakeSnapshot{} q' r' =
+  case q' of
+    Consensus.GetCBOR Consensus.GetStakeSnapshot {} -> SerialisedStakeSnapshot r'
+    _                                               -> fromConsensusQueryResultMismatch
+
 -- | This should /only/ happen if we messed up the mapping in 'toConsensusQuery'
 -- and 'fromConsensusQueryResult' so they are inconsistent with each other.
 --

cardano-api/src/Cardano/Api/Shelley.hs

@@ -207,9 +207,15 @@ module Cardano.Api.Shelley
     CurrentEpochState(..),
     SerialisedCurrentEpochState(..),
     decodeCurrentEpochState,
+
     PoolState(..),
     SerialisedPoolState(..),
     decodePoolState,
+
+    StakeSnapshot(..),
+    SerialisedStakeSnapshot(..),
+    decodeStakeSnapshot,
+
     UTxO(..),
 
     -- ** Various calculations

cardano-cli/src/Cardano/CLI/Shelley/Run/Query.hs

@@ -49,18 +49,14 @@ import qualified Cardano.Crypto.Hash.Blake2b as Blake2b
 import qualified Cardano.Crypto.VRF as Crypto
 import qualified Cardano.Ledger.Alonzo.PParams as Alonzo
 import           Cardano.Ledger.BaseTypes (Seed, UnitInterval)
-import           Cardano.Ledger.Coin
-import           Cardano.Ledger.Compactible
 import qualified Cardano.Ledger.Core as Core
 import qualified Cardano.Ledger.Credential as Ledger
 import qualified Cardano.Ledger.Crypto as Crypto
 import qualified Cardano.Ledger.Era as Era
 import qualified Cardano.Ledger.Era as Ledger
 import           Cardano.Ledger.Keys (KeyHash (..), KeyRole (..))
 import           Cardano.Ledger.Shelley.Constraints
-import           Cardano.Ledger.Shelley.EpochBoundary
-import           Cardano.Ledger.Shelley.LedgerState (EpochState (esSnapshots),
-                   NewEpochState (nesEs), PState (_fPParams, _pParams, _retiring))
+import           Cardano.Ledger.Shelley.LedgerState (PState (_fPParams, _pParams, _retiring))
 import qualified Cardano.Ledger.Shelley.LedgerState as SL
 import qualified Cardano.Ledger.Shelley.PParams as Shelley
 import           Cardano.Ledger.Shelley.Scripts ()
@@ -92,7 +88,6 @@ import qualified Data.Text.Encoding as Text
 import qualified Data.Text.IO as T
 import qualified Data.Text.IO as Text
 import qualified Data.Vector as Vector
-import qualified Data.VMap as VMap
 import           Formatting.Buildable (build)
 import           Numeric (showEFloat)
 import qualified Ouroboros.Consensus.HardFork.History as Consensus
@@ -132,6 +127,7 @@ data ShelleyQueryCmdError
       FilePath
       -- ^ Operational certificate of the unknown stake pool.
   | ShelleyQueryCmdPoolStateDecodeError DecoderError
+  | ShelleyQueryCmdStakeSnapshotDecodeError DecoderError
 
   deriving Show
 
@@ -169,6 +165,8 @@ renderShelleyQueryCmdError err =
                   "in the current epoch, you must wait until the following epoch for the registration to take place."
     ShelleyQueryCmdPoolStateDecodeError decoderError ->
       "Failed to decode PoolState.  Error: " <> Text.pack (show decoderError)
+    ShelleyQueryCmdStakeSnapshotDecodeError decoderError ->
+      "Failed to decode StakeSnapshot.  Error: " <> Text.pack (show decoderError)
 
 runQueryCmd :: QueryCmd -> ExceptT ShelleyQueryCmdError IO ()
 runQueryCmd cmd =
@@ -627,7 +625,7 @@ runQueryStakeSnapshot
   -> NetworkId
   -> Hash StakePoolKey
   -> ExceptT ShelleyQueryCmdError IO ()
-runQueryStakeSnapshot (AnyConsensusModeParams cModeParams) network poolid = do
+runQueryStakeSnapshot (AnyConsensusModeParams cModeParams) network poolId = do
   SocketPath sockPath <- firstExceptT ShelleyQueryCmdEnvVarSocketErr readEnvSocketPath
   let localNodeConnInfo = LocalNodeConnectInfo cModeParams network sockPath
 
@@ -638,9 +636,9 @@ runQueryStakeSnapshot (AnyConsensusModeParams cModeParams) network poolid = do
   eInMode <- toEraInMode era cMode
     & hoistMaybe (ShelleyQueryCmdEraConsensusModeMismatch (AnyConsensusMode cMode) anyE)
 
-  let qInMode = QueryInEra eInMode . QueryInShelleyBasedEra sbe $ QueryDebugLedgerState
+  let qInMode = QueryInEra eInMode . QueryInShelleyBasedEra sbe $ QueryStakeSnapshot poolId
   result <- executeQuery era cModeParams localNodeConnInfo qInMode
-  obtainLedgerEraClassConstraints sbe (writeStakeSnapshot poolid) result
+  obtainLedgerEraClassConstraints sbe writeStakeSnapshot result
 
 
 runQueryLedgerState
@@ -791,44 +789,15 @@ writeLedgerState mOutFile qState@(SerialisedDebugLedgerState serLedgerState) =
 writeStakeSnapshot :: forall era ledgerera. ()
   => ShelleyLedgerEra era ~ ledgerera
   => Era.Crypto ledgerera ~ StandardCrypto
-  => FromCBOR (DebugLedgerState era)
-  => PoolId
-  -> SerialisedDebugLedgerState era
+  => SerialisedStakeSnapshot era
   -> ExceptT ShelleyQueryCmdError IO ()
-writeStakeSnapshot (StakePoolKeyHash hk) qState =
-  case decodeDebugLedgerState qState of
-    -- In the event of decode failure print the CBOR instead
-    Left bs -> firstExceptT ShelleyQueryCmdHelpersError $ pPrintCBOR bs
-
-    Right ledgerState -> do
-      -- Ledger State
-      let (DebugLedgerState snapshot) = ledgerState
-
-      -- The three stake snapshots, obtained from the ledger state
-      let (SnapShots markS setS goS _) = esSnapshots $ nesEs snapshot
+writeStakeSnapshot qState =
+  case decodeStakeSnapshot qState of
+    Left err -> left (ShelleyQueryCmdStakeSnapshotDecodeError err)
 
+    Right (StakeSnapshot snapshot) -> do
       -- Calculate the three pool and active stake values for the given pool
-      liftIO . LBS.putStrLn $ encodePretty $ Stakes
-        { markPool = getPoolStake hk markS
-        , setPool = getPoolStake hk setS
-        , goPool = getPoolStake hk goS
-        , markTotal = getAllStake markS
-        , setTotal = getAllStake setS
-        , goTotal = getAllStake goS
-        }
-
--- | Sum all the stake that is held by the pool
-getPoolStake :: KeyHash Cardano.Ledger.Keys.StakePool crypto -> SnapShot crypto -> Integer
-getPoolStake hash ss = pStake
-  where
-    Coin pStake = fold (Map.map fromCompact $ VMap.toMap s)
-    Stake s = poolStake hash (_delegations ss) (_stake ss)
-
--- | Sum the active stake from a snapshot
-getAllStake :: SnapShot crypto -> Integer
-getAllStake (SnapShot stake _ _) = activeStake
-  where
-    Coin activeStake = fold (fmap fromCompact (VMap.toMap (unStake stake)))
+      liftIO . LBS.putStrLn $ encodePretty snapshot
 
 -- | This function obtains the pool parameters, equivalent to the following jq query on the output of query ledger-state
 --   .nesEs.esLState._delegationState._pstate._pParams.<pool_id>