From f0561b9d0347d3fb6014c107b5e080e74b0a45d1 Mon Sep 17 00:00:00 2001
From: Pawel Jakubas <pawel.jakubas@iohk.io>
Date: Mon, 29 May 2023 17:39:06 +0200
Subject: [PATCH] fix size estimation due to optional script withdrawals

---
 lib/wallet/src/Cardano/Wallet.hs              |  6 +++--
 .../src/Cardano/Wallet/Shelley/Transaction.hs | 24 +++++++++++++++----
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/lib/wallet/src/Cardano/Wallet.hs b/lib/wallet/src/Cardano/Wallet.hs
index 2ff771be122..3b5fa65fb2a 100644
--- a/lib/wallet/src/Cardano/Wallet.hs
+++ b/lib/wallet/src/Cardano/Wallet.hs
@@ -1325,14 +1325,13 @@ readRewardAccount
     => DBLayer IO s
     -> IO (RewardAccount, Maybe XPub, NonEmpty DerivationIndex)
 readRewardAccount db = do
+    walletState <- getState <$> readWalletCheckpoint db
     case walletFlavor @s of
         ShelleyWallet -> do
-            walletState <- getState <$> readWalletCheckpoint db
             let xpub = Seq.rewardAccountKey walletState
             let path = stakeDerivationPath $ Seq.derivationPrefix walletState
             pure (toRewardAccount xpub, Just $ getRawKey ShelleyKeyS xpub, path)
         SharedWallet -> do
-            walletState <- getState <$> readWalletCheckpoint db
             let path = stakeDerivationPath $ Shared.derivationPrefix walletState
             case Shared.rewardAccountKey walletState of
                 Just rewardAcct -> pure $ (rewardAcct, Nothing, path)
@@ -1342,6 +1341,9 @@ readWalletCheckpoint
     :: DBLayer IO s -> IO (Wallet s)
 readWalletCheckpoint DBLayer{..} = liftIO $ atomically readCheckpoint
 
+-- | Unsafe version of the `readRewardAccount` function
+-- that throws error when applied to a non-shared
+-- or a non-shared wallet state.
 sharedOnlyReadRewardAccount
     :: forall s
      . WalletFlavor s
diff --git a/lib/wallet/src/Cardano/Wallet/Shelley/Transaction.hs b/lib/wallet/src/Cardano/Wallet/Shelley/Transaction.hs
index fbb88124121..1676f0980b1 100644
--- a/lib/wallet/src/Cardano/Wallet/Shelley/Transaction.hs
+++ b/lib/wallet/src/Cardano/Wallet/Shelley/Transaction.hs
@@ -1446,6 +1446,7 @@ data TxSkeleton = TxSkeleton
     , txOutputs :: ![TxOut]
     , txChange :: ![Set AssetId]
     , txPaymentTemplate :: !(Maybe (Script Cosigner))
+    , txDelegationTemplate :: !(Maybe (Script Cosigner))
     , txMintOrBurnScripts :: [Script KeyHash]
     , txAssetsToMintOrBurn :: Set AssetId
     -- ^ The set of assets to mint or burn.
@@ -1466,6 +1467,7 @@ emptyTxSkeleton txWitnessTag = TxSkeleton
     , txOutputs = []
     , txChange = []
     , txPaymentTemplate = Nothing
+    , txDelegationTemplate = Nothing
     , txMintOrBurnScripts = []
     , txAssetsToMintOrBurn = Set.empty
     }
@@ -1491,6 +1493,9 @@ mkTxSkeleton witness context skeleton = TxSkeleton
     , txPaymentTemplate =
         template <$>
         view #txPaymentCredentialScriptTemplate context
+    , txDelegationTemplate =
+        template <$>
+        view #txStakingCredentialScriptTemplate context
     , txMintOrBurnScripts = (<>)
         (Map.elems (snd $ view #txAssetsToMint context))
         (Map.elems (snd $ view #txAssetsToBurn context))
@@ -1737,6 +1742,7 @@ estimateTxSize era skeleton =
         , txOutputs
         , txChange
         , txPaymentTemplate
+        , txDelegationTemplate
         , txMintOrBurnScripts
         , txAssetsToMintOrBurn
         } = skeleton
@@ -1754,23 +1760,33 @@ estimateTxSize era skeleton =
     numberOf_MintingWitnesses
         = intCast $ sumVia estimateMaxWitnessRequiredPerInput txMintOrBurnScripts
 
-    numberOf_ScriptVkeyWitnesses
+    numberOf_ScriptVkeyWitnessesForPayment
         = intCast $ maybe 0 estimateMaxWitnessRequiredPerInput txPaymentTemplate
 
+    numberOf_ScriptVkeyWitnessesForDelegation
+        = intCast $ maybe 0 estimateMaxWitnessRequiredPerInput txDelegationTemplate
+
     numberOf_VkeyWitnesses
         = case txWitnessTag of
             TxWitnessByronUTxO -> 0
             TxWitnessShelleyUTxO ->
-                if numberOf_ScriptVkeyWitnesses == 0 then
+                -- there cannot be missing payment script if there is delegation script
+                -- the latter is optional
+                if numberOf_ScriptVkeyWitnessesForPayment == 0 then
                     numberOf_Inputs
                     + numberOf_Withdrawals
                     + numberOf_CertificateSignatures
                     + numberOf_MintingWitnesses
-                else
-                    (numberOf_Inputs * numberOf_ScriptVkeyWitnesses)
+                else if numberOf_ScriptVkeyWitnessesForDelegation == 0 then
+                    (numberOf_Inputs * numberOf_ScriptVkeyWitnessesForPayment)
                     + numberOf_Withdrawals
                     + numberOf_CertificateSignatures
                     + numberOf_MintingWitnesses
+               else
+                    (numberOf_Inputs * numberOf_ScriptVkeyWitnessesForPayment)
+                    + (numberOf_Withdrawals * numberOf_ScriptVkeyWitnessesForDelegation)
+                    + numberOf_CertificateSignatures
+                    + numberOf_MintingWitnesses
 
     numberOf_BootstrapWitnesses
         = case txWitnessTag of