Optimise query leadership-schedule command

newhoggy · newhoggy · commit ae0dcf641151 · 2022-08-22T13:30:13.000+10:00
diff --git a/cardano-api/src/Cardano/Api/LedgerState.hs b/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 --git a/cardano-api/src/Cardano/Api/Query.hs b/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 --git a/cardano-api/src/Cardano/Api/Shelley.hs b/cardano-api/src/Cardano/Api/Shelley.hs
@@ -210,6 +210,9 @@ module Cardano.Api.Shelley
     PoolState(..),
     SerialisedPoolState(..),
     decodePoolState,
+    PoolDistr(..),
+    SerialisedPoolDistr(..),
+    decodePoolDistr,
     UTxO(..),
     AcquireFailure(..),
     SystemStart(..),
diff --git a/cardano-cli/src/Cardano/CLI/Shelley/Run/Query.hs b/cardano-cli/src/Cardano/CLI/Shelley/Run/Query.hs
@@ -1178,39 +1178,38 @@ runQueryLeadershipSchedule (AnyConsensusModeParams cModeParams) network
       eraHistory <- firstExceptT ShelleyQueryCmdAcquireFailure . newExceptT $ queryNodeLocalState localNodeConnInfo Nothing eraHistoryQuery
       let eInfo = toEpochInfo eraHistory
 
-      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
-               curentEpoch <- executeQuery era cModeParams localNodeConnInfo currentEpochQuery
-               firstExceptT ShelleyQueryCmdLeaderShipError $ hoistEither
-                $ eligibleLeaderSlotsConstaints sbe
-                $ currentEpochEligibleLeadershipSlots
-                    sbe
-                    shelleyGenesis
-                    eInfo
-                    pparams
-                    ptclState
-                    poolid
-                    vrkSkey
-                    serCurrentEpochState
-                    curentEpoch
-
-             NextEpoch -> do
-               let currentEpochStateQuery = QueryInEra eInMode $ QueryInShelleyBasedEra sbe QueryCurrentEpochState
-                   currentEpochQuery = QueryInEra eInMode $ QueryInShelleyBasedEra sbe QueryEpoch
-               tip <- liftIO $ getLocalChainTip localNodeConnInfo
-
-               curentEpoch <- executeQuery era cModeParams localNodeConnInfo currentEpochQuery
-               serCurrentEpochState <- executeQuery era cModeParams localNodeConnInfo currentEpochStateQuery
-
-               firstExceptT ShelleyQueryCmdLeaderShipError $ hoistEither
-                $ eligibleLeaderSlotsConstaints sbe
-                $ nextEpochEligibleLeadershipSlots sbe shelleyGenesis
-                    serCurrentEpochState ptclState poolid vrkSkey pparams
-                    eInfo (tip, curentEpoch)
+      schedule <- case whichSchedule of
+        CurrentEpoch -> do
+          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
+            $ currentEpochEligibleLeadershipSlots
+              sbe
+              shelleyGenesis
+              eInfo
+              pparams
+              ptclState
+              poolid
+              vrkSkey
+              serCurrentEpochState
+              curentEpoch
+
+        NextEpoch -> do
+          let currentEpochStateQuery = QueryInEra eInMode $ QueryInShelleyBasedEra sbe QueryCurrentEpochState
+              currentEpochQuery = QueryInEra eInMode $ QueryInShelleyBasedEra sbe QueryEpoch
+          tip <- liftIO $ getLocalChainTip localNodeConnInfo
+
+          curentEpoch <- executeQuery era cModeParams localNodeConnInfo currentEpochQuery
+          serCurrentEpochState <- executeQuery era cModeParams localNodeConnInfo currentEpochStateQuery
+
+          firstExceptT ShelleyQueryCmdLeaderShipError $ hoistEither
+            $ eligibleLeaderSlotsConstaints sbe
+            $ nextEpochEligibleLeadershipSlots sbe shelleyGenesis
+              serCurrentEpochState ptclState poolid vrkSkey pparams
+              eInfo (tip, curentEpoch)
 
       case mJsonOutputFile of
         Nothing -> liftIO $ printLeadershipScheduleAsText schedule eInfo (SystemStart $ sgSystemStart shelleyGenesis)
