Merge pull request #4250 from input-output-hk/newhoggy/optimise-query-leadership-schedule-command

newhoggy · web-flow · commit 0ffbd5a9db91 · 2022-12-22T14:19:59.000+11:00
Optimise query leadership schedule command
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 (..), PoolDistribution (unPoolDistr), ProtocolState,
+                   SerialisedCurrentEpochState (..), SerialisedPoolDistribution,
+                   decodeCurrentEpochState, decodePoolDistribution, decodeProtocolState)
 import           Cardano.Api.Utils (textShow)
 import           Cardano.Binary (DecoderError, FromCBOR)
 import qualified Cardano.Chain.Genesis
@@ -1385,9 +1386,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)))
@@ -1515,10 +1517,10 @@ currentEpochEligibleLeadershipSlots :: forall era ledgerera. ()
   -> ProtocolState era
   -> PoolId
   -> SigningKey VrfKey
-  -> SerialisedCurrentEpochState era
+  -> SerialisedPoolDistribution 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
@@ -1531,17 +1533,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
+      $ decodePoolDistribution 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,
 
+    SerialisedPoolDistribution(..),
+    PoolDistribution(..),
+    decodePoolDistribution,
+
     EraHistory(..),
     SystemStart(..),
 
@@ -246,6 +250,10 @@ data QueryInShelleyBasedEra era result where
     :: Maybe (Set PoolId)
     -> QueryInShelleyBasedEra era (SerialisedPoolState era)
 
+  QueryPoolDistribution
+    :: Maybe (Set PoolId)
+    -> QueryInShelleyBasedEra era (SerialisedPoolDistribution era)
+
 deriving instance Show (QueryInShelleyBasedEra era result)
 
 
@@ -406,6 +414,20 @@ decodePoolState
   -> Either DecoderError (PoolState era)
 decodePoolState (SerialisedPoolState (Serialised ls)) = PoolState <$> decodeFull ls
 
+newtype SerialisedPoolDistribution era
+  = SerialisedPoolDistribution (Serialised (Shelley.PoolDistr (Ledger.Crypto (ShelleyLedgerEra era))))
+
+newtype PoolDistribution era = PoolDistribution
+  { unPoolDistr :: Shelley.PoolDistr (Ledger.Crypto (ShelleyLedgerEra era))
+  }
+
+decodePoolDistribution
+  :: forall era. ()
+  => FromCBOR (Shelley.PoolDistr (Ledger.Crypto (ShelleyLedgerEra era)))
+  => SerialisedPoolDistribution era
+  -> Either DecoderError (PoolDistribution era)
+decodePoolDistribution (SerialisedPoolDistribution (Serialised ls)) = PoolDistribution <$> decodeFull ls
+
 toShelleyAddrSet :: CardanoEra era
                  -> Set AddressAny
                  -> Set (Shelley.Addr Consensus.StandardCrypto)
@@ -591,6 +613,12 @@ toConsensusQueryShelleyBased erainmode (QueryPoolState poolIds) =
     getPoolIds :: Set PoolId -> Set (Shelley.KeyHash Shelley.StakePool Consensus.StandardCrypto)
     getPoolIds = Set.map (\(StakePoolKeyHash kh) -> kh)
 
+toConsensusQueryShelleyBased erainmode (QueryPoolDistribution 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
@@ -826,6 +854,11 @@ fromConsensusQueryResultShelleyBased _ QueryPoolState{} q' r' =
     Consensus.GetCBOR Consensus.GetPoolState {} -> SerialisedPoolState r'
     _                                           -> fromConsensusQueryResultMismatch
 
+fromConsensusQueryResultShelleyBased _ QueryPoolDistribution{} q' r' =
+  case q' of
+    Consensus.GetCBOR Consensus.GetPoolDistr {} -> SerialisedPoolDistribution 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,
+    PoolDistribution(..),
+    SerialisedPoolDistribution(..),
+    decodePoolDistribution,
     UTxO(..),
     AcquiringFailure(..),
     SystemStart(..),
diff --git a/cardano-cli/src/Cardano/CLI/Shelley/Orphans.hs b/cardano-cli/src/Cardano/CLI/Shelley/Orphans.hs
@@ -15,11 +15,11 @@
 module Cardano.CLI.Shelley.Orphans () where
 
 import           Cardano.Api.Orphans ()
-import           Cardano.Ledger.AuxiliaryData (AuxiliaryDataHash (..))
+import qualified Cardano.Ledger.AuxiliaryData as Ledger
 import qualified Cardano.Ledger.Credential as Ledger
 import qualified Cardano.Ledger.Crypto as CC (Crypto)
 import qualified Cardano.Ledger.Mary.Value as Ledger.Mary
-import           Cardano.Ledger.PoolDistr (PoolDistr (..))
+import qualified Cardano.Ledger.PoolDistr as Ledger
 import qualified Cardano.Ledger.Shelley.EpochBoundary as Ledger
 import qualified Cardano.Ledger.Shelley.PoolRank as Ledger
 import           Cardano.Ledger.TxIn (TxId (..))
@@ -78,9 +78,9 @@ deriving newtype instance CC.Crypto crypto => ToJSON (TxId crypto)
 deriving newtype instance CC.Crypto crypto => ToJSON (ShelleyHash crypto)
 deriving newtype instance CC.Crypto crypto => ToJSON (HashHeader crypto)
 
-deriving newtype instance ToJSON (AuxiliaryDataHash StandardCrypto)
+deriving newtype instance ToJSON (Ledger.AuxiliaryDataHash StandardCrypto)
 deriving newtype instance ToJSON Ledger.LogWeight
-deriving newtype instance ToJSON (PoolDistr StandardCrypto)
+deriving newtype instance ToJSON (Ledger.PoolDistr StandardCrypto)
 
 deriving newtype instance ToJSON (Ledger.Stake StandardCrypto)
 
diff --git a/cardano-cli/src/Cardano/CLI/Shelley/Run/Query.hs b/cardano-cli/src/Cardano/CLI/Shelley/Run/Query.hs
@@ -1247,40 +1247,37 @@ runQueryLeadershipSchedule (AnyConsensusModeParams cModeParams) network
       ptclState <- executeQuery era cModeParams localNodeConnInfo ptclStateQuery
       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)
+      let currentEpochQuery = QueryInEra eInMode $ QueryInShelleyBasedEra sbe QueryEpoch
+      curentEpoch <- executeQuery era cModeParams localNodeConnInfo currentEpochQuery
+
+      schedule <- case whichSchedule of
+        CurrentEpoch -> do
+          serCurrentEpochState <- executeQuery era cModeParams localNodeConnInfo $
+            QueryInEra eInMode $ QueryInShelleyBasedEra sbe (QueryPoolDistribution (Just (Set.singleton poolid)))
+          firstExceptT ShelleyQueryCmdLeaderShipError $ hoistEither
+            $ eligibleLeaderSlotsConstaints sbe
+            $ currentEpochEligibleLeadershipSlots
+              sbe
+              shelleyGenesis
+              eInfo
+              pparams
+              ptclState
+              poolid
+              vrkSkey
+              serCurrentEpochState
+              curentEpoch
+
+        NextEpoch -> do
+          let currentEpochStateQuery = QueryInEra eInMode $ QueryInShelleyBasedEra sbe QueryCurrentEpochState
+
+          tip <- liftIO $ getLocalChainTip localNodeConnInfo
+          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)
