Optimise query leadership-schedule command

Author: newhoggy

Diff for: 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: 9f549559bd6cb714486fc50a8dbf295278f6e9d1
+  --sha256: 1c4cp0j6ry7n6ha5ix5cnmdl6n626am1mhm8c6q40iws8y4bxy05
   subdir:
     monoidal-synchronisation
     network-mux

Diff for: cardano-api/src/Cardano/Api/LedgerState.hs

@@ -105,8 +105,9 @@ import           Cardano.Api.Modes (CardanoMode, EpochSlots (..))
 import qualified Cardano.Api.Modes as Api
 import           Cardano.Api.NetworkId (NetworkId (..), NetworkMagic (NetworkMagic))
 import           Cardano.Api.ProtocolParameters
-import           Cardano.Api.Query (CurrentEpochState (..), ProtocolState,
-                   SerialisedCurrentEpochState (..), decodeCurrentEpochState, decodeProtocolState)
+import           Cardano.Api.Query (CurrentEpochState (..), PoolDistr (unPoolDistr), ProtocolState,
+                   SerialisedCurrentEpochState (..), SerialisedPoolDistr, decodeCurrentEpochState,
+                   decodePoolDistr, decodeProtocolState)
 import           Cardano.Binary (DecoderError, FromCBOR)
 import qualified Cardano.Chain.Genesis
 import qualified Cardano.Chain.Update
@@ -1384,9 +1385,10 @@ nextEpochEligibleLeadershipSlots sbe sGen serCurrEpochState ptclState poolid (Vr
                                  $ obtainDecodeEpochStateConstraints sbe
                                  $ decodeCurrentEpochState serCurrEpochState
 
-  let markSnapshotPoolDistr :: Map (SL.KeyHash 'SL.StakePool Shelley.StandardCrypto) (SL.IndividualPoolStake Shelley.StandardCrypto)
-      markSnapshotPoolDistr = ShelleyAPI.unPoolDistr . ShelleyAPI.calculatePoolDistr . ShelleyAPI._pstakeMark
-                                $ obtainIsStandardCrypto sbe $ ShelleyAPI.esSnapshots cEstate
+  let snapshot :: ShelleyAPI.SnapShot Shelley.StandardCrypto
+      snapshot = ShelleyAPI._pstakeMark $ obtainIsStandardCrypto sbe $ ShelleyAPI.esSnapshots cEstate
+      markSnapshotPoolDistr :: Map (SL.KeyHash 'SL.StakePool Shelley.StandardCrypto) (SL.IndividualPoolStake Shelley.StandardCrypto)
+      markSnapshotPoolDistr = ShelleyAPI.unPoolDistr . ShelleyAPI.calculatePoolDistr $ snapshot
 
   let slotRangeOfInterest = Set.filter
         (not . Ledger.isOverlaySlot firstSlotOfEpoch (getField @"_d" (toLedgerPParams sbe pParams)))
@@ -1510,10 +1512,10 @@ currentEpochEligibleLeadershipSlots :: forall era ledgerera. ()
   -> ProtocolState era
   -> PoolId
   -> SigningKey VrfKey
-  -> SerialisedCurrentEpochState era
+  -> SerialisedPoolDistr era
   -> EpochNo -- ^ Current EpochInfo
   -> Either LeadershipError (Set SlotNo)
-currentEpochEligibleLeadershipSlots sbe sGen eInfo pParams ptclState poolid (VrfSigningKey vrkSkey) serCurrEpochState currentEpoch = do
+currentEpochEligibleLeadershipSlots sbe sGen eInfo pParams ptclState poolid (VrfSigningKey vrkSkey) serPoolDistr currentEpoch = do
 
   chainDepState :: ChainDepState (Api.ConsensusProtocol era) <-
     first LeaderErrDecodeProtocolStateFailure $ decodeProtocolState ptclState
@@ -1526,17 +1528,10 @@ currentEpochEligibleLeadershipSlots sbe sGen eInfo pParams ptclState poolid (Vrf
   (firstSlotOfEpoch, lastSlotofEpoch) :: (SlotNo, SlotNo) <- first LeaderErrSlotRangeCalculationFailure
     $ Slot.epochInfoRange eInfo currentEpoch
 
-  CurrentEpochState (cEstate :: ShelleyAPI.EpochState (ShelleyLedgerEra era)) <-
-    first LeaderErrDecodeProtocolEpochStateFailure
+  setSnapshotPoolDistr <-
+    first LeaderErrDecodeProtocolEpochStateFailure . fmap (SL.unPoolDistr . unPoolDistr)
       $ obtainDecodeEpochStateConstraints sbe
-      $ decodeCurrentEpochState serCurrEpochState
-
-  -- We need the "set" stake distribution (distribution of the previous epoch)
-  -- in order to calculate the leadership schedule of the current epoch.
-  let setSnapshotPoolDistr :: Map (SL.KeyHash 'SL.StakePool Shelley.StandardCrypto) (SL.IndividualPoolStake Shelley.StandardCrypto)
-      setSnapshotPoolDistr = ShelleyAPI.unPoolDistr . ShelleyAPI.calculatePoolDistr
-                                . ShelleyAPI._pstakeSet . obtainIsStandardCrypto sbe
-                                $ ShelleyAPI.esSnapshots cEstate
+      $ decodePoolDistr serPoolDistr
 
   let slotRangeOfInterest = Set.filter
         (not . Ledger.isOverlaySlot firstSlotOfEpoch (getField @"_d" (toLedgerPParams sbe pParams)))

Diff for: cardano-api/src/Cardano/Api/Query.hs

@@ -48,6 +48,10 @@ module Cardano.Api.Query (
     PoolState(..),
     decodePoolState,
 
+    SerialisedPoolDistr(..),
+    PoolDistr(..),
+    decodePoolDistr,
+
     EraHistory(..),
     SystemStart(..),
 
@@ -245,6 +249,10 @@ data QueryInShelleyBasedEra era result where
     :: Maybe (Set PoolId)
     -> QueryInShelleyBasedEra era (SerialisedPoolState era)
 
+  QueryPoolDistr
+    :: Maybe (Set PoolId)
+    -> QueryInShelleyBasedEra era (SerialisedPoolDistr era)
+
 deriving instance Show (QueryInShelleyBasedEra era result)
 
 
@@ -403,6 +411,20 @@ decodePoolState
   -> Either DecoderError (PoolState era)
 decodePoolState (SerialisedPoolState (Serialised ls)) = PoolState <$> decodeFull ls
 
+newtype SerialisedPoolDistr era
+  = SerialisedPoolDistr (Serialised (Shelley.PoolDistr (Ledger.Crypto (ShelleyLedgerEra era))))
+
+newtype PoolDistr era = PoolDistr
+  { unPoolDistr :: Shelley.PoolDistr (Ledger.Crypto (ShelleyLedgerEra era))
+  }
+
+decodePoolDistr
+  :: forall era. ()
+  => FromCBOR (Shelley.PoolDistr (Ledger.Crypto (ShelleyLedgerEra era)))
+  => SerialisedPoolDistr era
+  -> Either DecoderError (PoolDistr era)
+decodePoolDistr (SerialisedPoolDistr (Serialised ls)) = PoolDistr <$> decodeFull ls
+
 toShelleyAddrSet :: CardanoEra era
                  -> Set AddressAny
                  -> Set (Shelley.Addr Consensus.StandardCrypto)
@@ -588,6 +610,12 @@ toConsensusQueryShelleyBased erainmode (QueryPoolState poolIds) =
     getPoolIds :: Set PoolId -> Set (Shelley.KeyHash Shelley.StakePool Consensus.StandardCrypto)
     getPoolIds = Set.map (\(StakePoolKeyHash kh) -> kh)
 
+toConsensusQueryShelleyBased erainmode (QueryPoolDistr poolIds) =
+    Some (consensusQueryInEraInMode erainmode (Consensus.GetCBOR (Consensus.GetPoolDistr (getPoolIds <$> poolIds))))
+  where
+    getPoolIds :: Set PoolId -> Set (Shelley.KeyHash Shelley.StakePool Consensus.StandardCrypto)
+    getPoolIds = Set.map (\(StakePoolKeyHash kh) -> kh)
+
 consensusQueryInEraInMode
   :: forall era mode erablock modeblock result result' xs.
      ConsensusBlockForEra era   ~ erablock
@@ -823,6 +851,11 @@ fromConsensusQueryResultShelleyBased _ QueryPoolState{} q' r' =
     Consensus.GetCBOR Consensus.GetPoolState {} -> SerialisedPoolState r'
     _                                           -> fromConsensusQueryResultMismatch
 
+fromConsensusQueryResultShelleyBased _ QueryPoolDistr{} q' r' =
+  case q' of
+    Consensus.GetCBOR Consensus.GetPoolDistr {} -> SerialisedPoolDistr r'
+    _                                           -> fromConsensusQueryResultMismatch
+
 -- | This should /only/ happen if we messed up the mapping in 'toConsensusQuery'
 -- and 'fromConsensusQueryResult' so they are inconsistent with each other.
 --

Diff for: cardano-api/src/Cardano/Api/Shelley.hs

@@ -210,6 +210,9 @@ module Cardano.Api.Shelley
     PoolState(..),
     SerialisedPoolState(..),
     decodePoolState,
+    PoolDistr(..),
+    SerialisedPoolDistr(..),
+    decodePoolDistr,
     UTxO(..),
     AcquireFailure(..),
     SystemStart(..),
@@ -235,8 +238,8 @@ import           Cardano.Api.Address
 import           Cardano.Api.Block
 import           Cardano.Api.Certificate
 import           Cardano.Api.Eras
-import           Cardano.Api.IPC
 import           Cardano.Api.InMode
+import           Cardano.Api.IPC
 import           Cardano.Api.KeysByron
 import           Cardano.Api.KeysPraos
 import           Cardano.Api.KeysShelley

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

@@ -1181,9 +1181,9 @@ runQueryLeadershipSchedule (AnyConsensusModeParams cModeParams) network
       schedule :: Set SlotNo
         <- case whichSchedule of
              CurrentEpoch -> do
-               let currentEpochStateQuery = QueryInEra eInMode $ QueryInShelleyBasedEra sbe QueryCurrentEpochState
-                   currentEpochQuery = QueryInEra eInMode $ QueryInShelleyBasedEra sbe QueryEpoch
-               serCurrentEpochState <- executeQuery era cModeParams localNodeConnInfo currentEpochStateQuery
+               let currentEpochQuery = QueryInEra eInMode $ QueryInShelleyBasedEra sbe QueryEpoch
+               serCurrentEpochState <- executeQuery era cModeParams localNodeConnInfo $
+                QueryInEra eInMode $ QueryInShelleyBasedEra sbe (QueryPoolDistr (Just (Set.singleton poolid)))
                curentEpoch <- executeQuery era cModeParams localNodeConnInfo currentEpochQuery
                firstExceptT ShelleyQueryCmdLeaderShipError $ hoistEither
                 $ eligibleLeaderSlotsConstaints sbe