From 2d8c70263edb697fe5fc12c450f4f7ca844a1f01 Mon Sep 17 00:00:00 2001
From: Carl Hammann <carl.hammann@tweag.io>
Date: Wed, 10 Jan 2024 12:22:26 +0100
Subject: [PATCH] Use custom sum type to turn stake info on or off

---
 cardano-cli/src/Cardano/CLI/EraBased/Commands/Query.hs | 8 +++++++-
 cardano-cli/src/Cardano/CLI/EraBased/Options/Query.hs  | 2 +-
 cardano-cli/src/Cardano/CLI/EraBased/Run/Query.hs      | 8 ++++----
 3 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/cardano-cli/src/Cardano/CLI/EraBased/Commands/Query.hs b/cardano-cli/src/Cardano/CLI/EraBased/Commands/Query.hs
index a06c4645fe..831594578e 100644
--- a/cardano-cli/src/Cardano/CLI/EraBased/Commands/Query.hs
+++ b/cardano-cli/src/Cardano/CLI/EraBased/Commands/Query.hs
@@ -25,6 +25,7 @@ module Cardano.CLI.EraBased.Commands.Query
   , QueryDRepStateCmdArgs(..)
   , QueryDRepStakeDistributionCmdArgs(..)
   , renderQueryCmds
+  , IncludeStake (..)
   ) where
 
 import           Cardano.Api.Shelley hiding (QueryInShelleyBasedEra (..))
@@ -188,10 +189,15 @@ data QueryDRepStateCmdArgs era = QueryDRepStateCmdArgs
   , consensusModeParams :: !ConsensusModeParams
   , networkId           :: !NetworkId
   , drepKeys            :: ![VerificationKeyOrHashOrFile DRepKey]
-  , includeStake        :: !Bool -- ^ whether to include the stake, as queried by drep-stake-distribution, in the output. This is (computationally) expensive, but sometimes convenient.
+  , includeStake        :: !IncludeStake
   , mOutFile            :: !(Maybe (File () Out))
   } deriving Show
 
+-- | Whether to include the stake, as queried by drep-stake-distribution, in
+-- the output of drep-state. This is (computationally) expensive, but sometimes
+-- convenient.
+data IncludeStake = WithStake | NoStake deriving Show
+
 data QueryDRepStakeDistributionCmdArgs era = QueryDRepStakeDistributionCmdArgs
   { eon                 :: !(ConwayEraOnwards era)
   , nodeSocketPath      :: !SocketPath
diff --git a/cardano-cli/src/Cardano/CLI/EraBased/Options/Query.hs b/cardano-cli/src/Cardano/CLI/EraBased/Options/Query.hs
index 859acf635c..e5d83228fd 100644
--- a/cardano-cli/src/Cardano/CLI/EraBased/Options/Query.hs
+++ b/cardano-cli/src/Cardano/CLI/EraBased/Options/Query.hs
@@ -323,7 +323,7 @@ pQueryDRepStateCmd era envCli = do
         <*> pConsensusModeParams
         <*> pNetworkId envCli
         <*> many pDRepVerificationKeyOrHashOrFile
-        <*> Opt.flag True False (mconcat
+        <*> Opt.flag WithStake NoStake (mconcat
               [ Opt.long "include-stake"
               , Opt.help $ mconcat
                  [ "Also return the stake associated with each DRep. "
diff --git a/cardano-cli/src/Cardano/CLI/EraBased/Run/Query.hs b/cardano-cli/src/Cardano/CLI/EraBased/Run/Query.hs
index 51b2905f1f..9fc34de5cc 100644
--- a/cardano-cli/src/Cardano/CLI/EraBased/Run/Query.hs
+++ b/cardano-cli/src/Cardano/CLI/EraBased/Run/Query.hs
@@ -1404,10 +1404,10 @@ runQueryDRepState
   drepState <- runQuery localNodeConnInfo $ queryDRepState eon $ Set.fromList drepCreds
 
   drepStakeDistribution <-
-    if includeStake
-    then runQuery localNodeConnInfo $
-      queryDRepStakeDistribution eon (Set.fromList $ Ledger.DRepCredential <$> drepCreds)
-    else return mempty
+    case includeStake of
+      Cmd.WithStake -> runQuery localNodeConnInfo $
+        queryDRepStakeDistribution eon (Set.fromList $ Ledger.DRepCredential <$> drepCreds)
+      Cmd.NoStake -> return mempty
 
   writeOutput mOutFile $
     drepStateToJson drepStakeDistribution <$> Map.assocs drepState