Optimise query stake snapshot command

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

@@ -1218,7 +1218,7 @@ 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

@@ -6,6 +6,7 @@
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE StandaloneDeriving #-}

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

@@ -52,6 +52,10 @@ module Cardano.Api.Query (
     PoolDistribution(..),
     decodePoolDistribution,
 
+    SerialisedStakeSnapshots(..),
+    StakeSnapshot(..),
+    decodeStakeSnapshot,
+
     EraHistory(..),
     SystemStart(..),
 
@@ -255,6 +259,10 @@ data QueryInShelleyBasedEra era result where
     :: Maybe (Set PoolId)
     -> QueryInShelleyBasedEra era (SerialisedPoolDistribution era)
 
+  QueryStakeSnapshot
+    :: PoolId
+    -> QueryInShelleyBasedEra era (SerialisedStakeSnapshots era)
+
 deriving instance Show (QueryInShelleyBasedEra era result)
 
 
@@ -429,6 +437,18 @@ decodePoolDistribution
   -> Either DecoderError (PoolDistribution era)
 decodePoolDistribution (SerialisedPoolDistribution (Serialised ls)) = PoolDistribution <$> decodeFull ls
 
+newtype SerialisedStakeSnapshots era
+  = SerialisedStakeSnapshots (Serialised (Consensus.StakeSnapshots (Ledger.Crypto (ShelleyLedgerEra era))))
+
+newtype StakeSnapshot era = StakeSnapshot (Consensus.StakeSnapshots (Ledger.Crypto (ShelleyLedgerEra era)))
+
+decodeStakeSnapshot
+  :: forall era. ()
+  => FromCBOR (Consensus.StakeSnapshots (Ledger.Crypto (ShelleyLedgerEra era)))
+  => SerialisedStakeSnapshots era
+  -> Either DecoderError (StakeSnapshot era)
+decodeStakeSnapshot (SerialisedStakeSnapshots (Serialised ls)) = StakeSnapshot <$> decodeFull ls
+
 toShelleyAddrSet :: CardanoEra era
                  -> Set AddressAny
                  -> Set (Shelley.Addr Consensus.StandardCrypto)
@@ -597,7 +617,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))
@@ -609,10 +629,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.GetStakeSnapshots (Just (Set.singleton (unStakePoolKeyHash poolId))))))
 
 toConsensusQueryShelleyBased erainmode (QueryPoolDistribution poolIds) =
     Some (consensusQueryInEraInMode erainmode (Consensus.GetCBOR (Consensus.GetPoolDistr (getPoolIds <$> poolIds))))
@@ -860,6 +880,11 @@ fromConsensusQueryResultShelleyBased _ QueryPoolDistribution{} q' r' =
     Consensus.GetCBOR Consensus.GetPoolDistr {} -> SerialisedPoolDistribution r'
     _                                           -> fromConsensusQueryResultMismatch
 
+fromConsensusQueryResultShelleyBased _ QueryStakeSnapshot{} q' r' =
+  case q' of
+    Consensus.GetCBOR Consensus.GetStakeSnapshots {} -> SerialisedStakeSnapshots 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,12 +207,19 @@ module Cardano.Api.Shelley
     CurrentEpochState(..),
     SerialisedCurrentEpochState(..),
     decodeCurrentEpochState,
+
     PoolState(..),
     SerialisedPoolState(..),
     decodePoolState,
+
     PoolDistribution(..),
     SerialisedPoolDistribution(..),
     decodePoolDistribution,
+
+    StakeSnapshot(..),
+    SerialisedStakeSnapshots(..),
+    decodeStakeSnapshot,
+
     UTxO(..),
     AcquiringFailure(..),
     SystemStart(..),

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

@@ -53,7 +53,6 @@ import qualified Data.Text.IO as Text
 import           Data.Text.Lazy.Builder (toLazyText)
 import           Data.Time.Clock
 import qualified Data.Vector as Vector
-import qualified Data.VMap as VMap
 import           Formatting.Buildable (build)
 import           Numeric (showEFloat)
 import qualified System.IO as IO
@@ -74,18 +73,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.SafeHash (HashAnnotated)
-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 ()
@@ -94,12 +89,13 @@ import           Cardano.Slotting.EpochInfo (EpochInfo (..), epochInfoSlotToUTCT
 import           Ouroboros.Consensus.BlockchainTime.WallClock.Types (RelativeTime (..),
                    toRelativeTime)
 import           Ouroboros.Consensus.Cardano.Block as Consensus (EraMismatch (..))
-import           Ouroboros.Consensus.Protocol.TPraos
+import           Ouroboros.Consensus.Protocol.TPraos ( StandardCrypto )
 import           Ouroboros.Network.Block (Serialised (..))
 
 import qualified Ouroboros.Consensus.HardFork.History as Consensus
 import qualified Ouroboros.Consensus.Protocol.Abstract as Consensus
 import qualified Ouroboros.Consensus.Protocol.Praos.Common as Consensus
+import qualified Ouroboros.Consensus.Shelley.Ledger.Query as Consensus
 
 import qualified Ouroboros.Consensus.HardFork.History.Qry as Qry
 import qualified Ouroboros.Network.Protocol.LocalStateQuery.Type as LocalStateQuery
@@ -133,6 +129,7 @@ data ShelleyQueryCmdError
       FilePath
       -- ^ Operational certificate of the unknown stake pool.
   | ShelleyQueryCmdPoolStateDecodeError DecoderError
+  | ShelleyQueryCmdStakeSnapshotDecodeError DecoderError
 
   deriving Show
 
@@ -170,6 +167,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 =
@@ -674,9 +673,8 @@ runQueryStakeSnapshot
   -> NetworkId
   -> Hash StakePoolKey
   -> ExceptT ShelleyQueryCmdError IO ()
-runQueryStakeSnapshot (AnyConsensusModeParams cModeParams) network poolid = do
-  SocketPath sockPath <- firstExceptT ShelleyQueryCmdEnvVarSocketErr
-                           $ newExceptT readEnvSocketPath
+runQueryStakeSnapshot (AnyConsensusModeParams cModeParams) network poolId = do
+  SocketPath sockPath <- firstExceptT ShelleyQueryCmdEnvVarSocketErr $ newExceptT readEnvSocketPath
   let localNodeConnInfo = LocalNodeConnectInfo cModeParams network sockPath
 
   anyE@(AnyCardanoEra era) <-
@@ -689,9 +687,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
@@ -855,44 +853,27 @@ writeLedgerState mOutFile qState@(SerialisedDebugLedgerState serLedgerState) =
 writeStakeSnapshot :: forall era ledgerera. ()
   => ShelleyLedgerEra era ~ ledgerera
   => Era.Crypto ledgerera ~ StandardCrypto
-  => FromCBOR (DebugLedgerState era)
-  => PoolId
-  -> SerialisedDebugLedgerState era
+  => SerialisedStakeSnapshots 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 $ Aeson.object $
+        [ "activeStakeMark" .= Consensus.ssMarkTotal snapshot
+        , "activeStakeSet"  .= Consensus.ssSetTotal snapshot
+        , "activeStakeGo"   .= Consensus.ssGoTotal snapshot
+        ] <> poolFields snapshot
+  where poolFields :: Consensus.StakeSnapshots (Ledger.Crypto (ShelleyLedgerEra era)) -> [Aeson.Pair]
+        poolFields snapshot = case Map.elems (Consensus.ssStakeSnapshots snapshot) of
+          [pool] -> 
+            [ "poolStakeMark"   .= Consensus.ssMarkPool pool
+            , "poolStakeSet"    .= Consensus.ssSetPool pool
+            , "poolStakeGo"     .= Consensus.ssGoPool pool
+            ]
+          _ -> []
 
 -- | 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>

cardano-cli/src/Cardano/CLI/Types.hs

@@ -36,7 +36,6 @@ module Cardano.CLI.Types
   , TxMempoolQuery (..)
   , UpdateProposalFile (..)
   , VerificationKeyFile (..)
-  , Stakes (..)
   , Params (..)
   , RequiredSigner (..)
   ) where
@@ -182,43 +181,6 @@ data OutputFormat
   | OutputFormatBech32
   deriving (Eq, Show)
 
-
--- | This data structure is used to allow nicely formatted output within the query stake-snapshot command.
---
--- "markPool", "setPool", "goPool" are the three ledger state stake snapshots (from most recent to least recent)
--- go is the snapshot that is used for the current epoch, set will be used in the next epoch,
--- mark for the epoch after that.  "markTotal", "setTotal", "goTotal" record the total active stake for each snapshot.
---
--- This information can be used by community tools to calculate upcoming leader schedules.
-data Stakes =  Stakes
-  { markPool :: Integer
-  , setPool :: Integer
-  , goPool :: Integer
-  , markTotal :: Integer
-  , setTotal :: Integer
-  , goTotal :: Integer
-  } deriving Show
-
--- | Pretty printing for stake information
-instance ToJSON Stakes where
-  toJSON (Stakes m s g mt st gt) = object
-    [ "poolStakeMark" .= m
-    , "poolStakeSet" .= s
-    , "poolStakeGo" .= g
-    , "activeStakeMark" .= mt
-    , "activeStakeSet" .= st
-    , "activeStakeGo" .= gt
-    ]
-
-  toEncoding  (Stakes m s g mt st gt) = pairs $ mconcat
-    [ "poolStakeMark" .= m
-    , "poolStakeSet" .= s
-    , "poolStakeGo" .= g
-    , "activeStakeMark" .= mt
-    , "activeStakeSet" .= st
-    , "activeStakeGo" .= gt
-    ]
-
 -- | This data structure is used to allow nicely formatted output in the query pool-params command.
 -- params are the current pool parameter settings, futureparams are new parameters, retiringEpoch is the
 -- epoch that has been set for pool retirement.  Any of these may be Nothing.