From ee57bc631824fa71e096af8e2eef7d8f4352d4f1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Hurlin?= <clement.hurlin@moduscreate.com>
Date: Mon, 9 Dec 2024 21:33:05 +0100
Subject: [PATCH] Augment spo-stake-distribution to include DRep delegation
 choice of pool's reward account

---
 .../src/Cardano/CLI/EraBased/Run/Query.hs     | 63 +++++++++++++++++--
 1 file changed, 58 insertions(+), 5 deletions(-)

diff --git a/cardano-cli/src/Cardano/CLI/EraBased/Run/Query.hs b/cardano-cli/src/Cardano/CLI/EraBased/Run/Query.hs
index f7c87a72de..2f47ddb2ec 100644
--- a/cardano-cli/src/Cardano/CLI/EraBased/Run/Query.hs
+++ b/cardano-cli/src/Cardano/CLI/EraBased/Run/Query.hs
@@ -10,6 +10,7 @@
 {-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TupleSections #-}
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE TypeOperators #-}
 
@@ -57,6 +58,7 @@ import           Cardano.CLI.Types.Output (QueryDRepStateOutput (..))
 import qualified Cardano.CLI.Types.Output as O
 import           Cardano.Crypto.Hash (hashToBytesAsHex)
 import qualified Cardano.Crypto.Hash.Blake2b as Blake2b
+import           Cardano.Prelude (catMaybes)
 import           Cardano.Slotting.EpochInfo (EpochInfo (..), epochInfoSlotToUTCTime, hoistEpochInfo)
 import           Cardano.Slotting.Time (RelativeTime (..), toRelativeTime)
 
@@ -917,7 +919,6 @@ callQueryStakeAddressInfoCmd
         , Cmd.target
         }
     , Cmd.addr = StakeAddress _ addr
-    , Cmd.mOutFile
     } = do
     let localNodeConnInfo = LocalNodeConnectInfo consensusModeParams networkId nodeSocketPath
 
@@ -1344,6 +1345,7 @@ runQueryStakeDistributionCmd
         & onLeft (left . QueryCmdAcquireFailure)
         & onLeft left
 
+
 writeStakeDistribution
   :: OutputFormatJsonOrText
   -> Maybe (File () Out)
@@ -1670,7 +1672,7 @@ runQuerySPOStakeDistribution
   Cmd.QuerySPOStakeDistributionCmdArgs
     { Cmd.eon
     , Cmd.commons =
-      Cmd.QueryCommons
+      commons@Cmd.QueryCommons
         { Cmd.nodeSocketPath
         , Cmd.consensusModeParams
         , Cmd.networkId
@@ -1687,9 +1689,60 @@ runQuerySPOStakeDistribution
 
     spos <- fromList <$> mapM spoFromSource spoHashSources
 
-    spoStakeDistribution <- runQuery localNodeConnInfo target $ querySPOStakeDistribution eon spos
-    writeOutput mOutFile $
-      Map.assocs spoStakeDistribution
+    let beo = convert eon
+
+    spoStakeDistribution :: Map (L.KeyHash L.StakePool StandardCrypto) L.Coin <-
+      runQuery localNodeConnInfo target $ querySPOStakeDistribution eon spos
+    let poolIds :: Maybe (Set (Hash StakePoolKey)) = Just $ Set.fromList $ map StakePoolKeyHash $ Map.keys spoStakeDistribution
+
+    serialisedPoolState :: SerialisedPoolState era <-
+      runQuery localNodeConnInfo target $ queryPoolState beo poolIds
+
+    PoolState (poolState :: (L.PState (ShelleyLedgerEra era))) <-
+      pure (decodePoolState serialisedPoolState)
+        & onLeft (left . QueryCmdPoolStateDecodeError)
+
+    let spoToPoolParams
+          :: Map
+              (L.KeyHash L.StakePool StandardCrypto)
+              (L.PoolParams StandardCrypto) = L.psStakePoolParams poolState
+        rewardsAccounts
+          :: Map
+              (L.KeyHash L.StakePool StandardCrypto)
+              StakeCredential = Map.map (fromShelleyStakeCredential . L.raCredential . L.ppRewardAccount) spoToPoolParams
+        rewardsAddresses
+          :: Map
+              (L.KeyHash L.StakePool StandardCrypto)
+              StakeAddress = Map.map (makeStakeAddress networkId) rewardsAccounts
+        addressesAndRewards
+          :: Map
+              StakeAddress
+              (L.KeyHash L.StakePool StandardCrypto) = Map.fromList [(addr, keyHash) | (keyHash, addr) <- Map.toList rewardsAddresses]
+        mkQueryStakeAddressInfoCmdArgs addr =
+          Cmd.QueryStakeAddressInfoCmdArgs
+            { Cmd.commons = commons
+            , addr
+            , mOutFile -- unused anyway. TODO tighten this by removing the field.
+            }
+    infos <-
+      mapM (callQueryStakeAddressInfoCmd . mkQueryStakeAddressInfoCmdArgs) $ Map.elems rewardsAddresses
+    let spoToDelegatee :: Map (L.KeyHash L.StakePool StandardCrypto) (L.DRep StandardCrypto) =
+          Map.fromList $
+            catMaybes $
+              [ fmap (,delegatee) mSpo
+              | info <- infos
+              , (addr, delegatee) <- Map.toList $ delegatees info
+              , let mSpo = Map.lookup addr addressesAndRewards
+              ]
+        toWrite =
+          [ ( spo
+            , coin
+            , Map.lookup spo spoToDelegatee
+            )
+          | (spo, coin) <- Map.assocs spoStakeDistribution
+          ]
+
+    writeOutput mOutFile toWrite
 
 runQueryCommitteeMembersState
   :: Cmd.QueryCommitteeMembersStateCmdArgs era