From 10ebb5fd1d9f2b5f86a52d5e7f589d8b7eb66b61 Mon Sep 17 00:00:00 2001
From: Heinrich Apfelmus <apfelmus@quantentunnel.de>
Date: Fri, 18 Oct 2024 15:47:33 +0200
Subject: [PATCH] Extend public module `Cardano.Write.Tx`

with the necessary types to be able to call `balanceTx`.
---
 lib/balance-tx/cardano-balance-tx.cabal     |  1 +
 lib/balance-tx/lib/main/Cardano/Write/Tx.hs | 51 +++++++++++++++------
 lib/wallet/src/Cardano/Wallet.hs            | 45 +++++++++---------
 3 files changed, 62 insertions(+), 35 deletions(-)

diff --git a/lib/balance-tx/cardano-balance-tx.cabal b/lib/balance-tx/cardano-balance-tx.cabal
index c23edfa841a..ed5158a5526 100644
--- a/lib/balance-tx/cardano-balance-tx.cabal
+++ b/lib/balance-tx/cardano-balance-tx.cabal
@@ -39,6 +39,7 @@ library
   hs-source-dirs:  lib/main
   build-depends:
     , cardano-balance-tx:internal
+    , cardano-ledger-api
   exposed-modules:
     Cardano.Write.Eras
     Cardano.Write.Tx
diff --git a/lib/balance-tx/lib/main/Cardano/Write/Tx.hs b/lib/balance-tx/lib/main/Cardano/Write/Tx.hs
index 0c53ade9849..2e7ec41eeda 100644
--- a/lib/balance-tx/lib/main/Cardano/Write/Tx.hs
+++ b/lib/balance-tx/lib/main/Cardano/Write/Tx.hs
@@ -1,11 +1,31 @@
 -- |
--- Copyright: © 2023 Cardano Foundation
+-- Copyright: © 2023-2024 Cardano Foundation
 -- License: Apache-2.0
 --
+--
 module Cardano.Write.Tx
     (
-    -- * Balancing transactions
-      balanceTx
+    -- * UTxO
+      UTxO (..)
+
+    -- * Tx
+    , Tx
+    , TxBody
+
+    -- ** PartialTx
+    , PartialTx (..)
+    , Redeemer (..)
+    , StakeKeyDepositLookup (..)
+    , TimelockKeyWitnessCounts (..)
+
+    -- ** Balancing
+    , balanceTx
+    , UTxOAssumptions (..)
+    , UTxOIndex
+    , constructUTxOIndex
+    , ChangeAddressGen (..)
+
+    -- ** Balancing Errors
     , ErrAssignRedeemers (..)
     , ErrBalanceTx (..)
     , ErrBalanceTxAssetsInsufficientError (..)
@@ -14,19 +34,18 @@ module Cardano.Write.Tx
     , ErrBalanceTxOutputError (..)
     , ErrBalanceTxOutputErrorInfo (..)
     , ErrBalanceTxUnableToCreateChangeError (..)
-
-    -- * UTxO-related types and functions
-    , UTxO
-    , UTxOAssumptions
-    , UTxOIndex
-    , constructUTxOIndex
     ) where
 
-import Internal.Cardano.Write.Tx
-    ( UTxO
+import Cardano.Ledger.Api
+    ( Tx
+    , TxBody
+    )
+import Cardano.Ledger.Api.UTxO
+    ( UTxO (..)
     )
 import Internal.Cardano.Write.Tx.Balance
-    ( ErrAssignRedeemers (..)
+    ( ChangeAddressGen (..)
+    , ErrAssignRedeemers (..)
     , ErrBalanceTx (..)
     , ErrBalanceTxAssetsInsufficientError (..)
     , ErrBalanceTxInsufficientCollateralError (..)
@@ -34,8 +53,14 @@ import Internal.Cardano.Write.Tx.Balance
     , ErrBalanceTxOutputError (..)
     , ErrBalanceTxOutputErrorInfo (..)
     , ErrBalanceTxUnableToCreateChangeError (..)
-    , UTxOAssumptions
+    , PartialTx (..)
+    , Redeemer (..)
+    , StakeKeyDepositLookup (..)
+    , UTxOAssumptions (..)
     , UTxOIndex
     , balanceTx
     , constructUTxOIndex
     )
+import Internal.Cardano.Write.Tx.Sign
+    ( TimelockKeyWitnessCounts (..)
+    )
diff --git a/lib/wallet/src/Cardano/Wallet.hs b/lib/wallet/src/Cardano/Wallet.hs
index 7d48babc2b2..b8ecdf859c4 100644
--- a/lib/wallet/src/Cardano/Wallet.hs
+++ b/lib/wallet/src/Cardano/Wallet.hs
@@ -633,7 +633,8 @@ import Cardano.Write.Eras
     , recentEra
     )
 import Cardano.Write.Tx
-    ( ErrBalanceTx (..)
+    ( ChangeAddressGen (..)
+    , ErrBalanceTx (..)
     , ErrBalanceTxUnableToCreateChangeError (..)
     )
 import Control.Arrow
@@ -810,10 +811,7 @@ import Internal.Cardano.Write.Tx
     ( toRecentEraGADT
     )
 import Internal.Cardano.Write.Tx.Balance
-    ( ChangeAddressGen (..)
-    , PartialTx (..)
-    , UTxOAssumptions (..)
-    , stakeCredentialsWithRefunds
+    ( stakeCredentialsWithRefunds
     )
 import Internal.Cardano.Write.Tx.SizeEstimation
     ( TxWitnessTag (..)
@@ -875,6 +873,16 @@ import qualified Cardano.Write.Eras as Write
     , RecentEra (..)
     , cardanoEraFromRecentEra
     )
+import qualified Cardano.Write.Tx as Write
+    ( PartialTx (..)
+    , StakeKeyDepositLookup (..)
+    , Tx
+    , UTxO (UTxO)
+    , UTxOAssumptions (..)
+    , UTxOIndex
+    , balanceTx
+    , constructUTxOIndex
+    )
 import qualified Data.ByteArray as BA
 import qualified Data.Delta.Update as Delta
 import qualified Data.Foldable as F
@@ -888,8 +896,6 @@ import qualified Internal.Cardano.Write.Tx as Write
     , FeePerByte
     , PParams
     , PParamsInAnyRecentEra (PParamsInAnyRecentEra)
-    , Tx
-    , UTxO (UTxO)
     , feeOfBytes
     , forceUTxOToEra
     , fromCardanoApiTx
@@ -898,12 +904,7 @@ import qualified Internal.Cardano.Write.Tx as Write
     , toCardanoApiTx
     )
 import qualified Internal.Cardano.Write.Tx.Balance as Write
-    ( PartialTx
-    , StakeKeyDepositLookup (StakeKeyDepositAssumeCurrent, StakeKeyDepositMap)
-    , UTxOIndex
-    , balanceTx
-    , constructUTxOIndex
-    , fromWalletUTxO
+    ( fromWalletUTxO
     )
 
 -- $Development
@@ -2186,7 +2187,7 @@ balanceTx
     => WalletLayer IO s
     -> Write.PParams era
     -> TimeTranslation
-    -> PartialTx era
+    -> Write.PartialTx era
     -> IO (Write.Tx era)
 balanceTx wrk pp timeTranslation partialTx = do
     (utxo, wallet, _txs) <- liftIO $ readWalletUTxO wrk
@@ -2206,8 +2207,8 @@ balanceTx wrk pp timeTranslation partialTx = do
         <$> getStakeDelegDeposits netLayer deregCreds
 
     let utxoAssumptions = case walletFlavor @s of
-            ShelleyWallet -> AllKeyPaymentCredentials
-            SharedWallet -> AllScriptPaymentCredentialsFrom
+            ShelleyWallet -> Write.AllKeyPaymentCredentials
+            SharedWallet -> Write.AllScriptPaymentCredentialsFrom
                 (Shared.paymentTemplate (getState wallet))
                 (sharedWalletScriptLookup (getState wallet)
                     . Convert.toWalletAddress)
@@ -2553,7 +2554,7 @@ buildTransactionPure
             utxoIndex
             changeAddrGen
             (getState wallet)
-            PartialTx
+            Write.PartialTx
                 { tx = Write.fromCardanoApiTx (Cardano.Tx unsignedTxBody [])
                 , extraUTxO = Write.UTxO mempty
                 , redeemers = []
@@ -3248,7 +3249,7 @@ transactionFee DBLayer{atomically, walletState} protocolParams
                 (Left preSelection)
 
         let ptx :: Write.PartialTx era
-            ptx = PartialTx
+            ptx = Write.PartialTx
                 { tx = Write.fromCardanoApiTx (Cardano.Tx unsignedTxBody [])
                 , extraUTxO = Write.UTxO mempty
                 , redeemers = []
@@ -4174,8 +4175,8 @@ utxoAssumptionsForWallet
     :: forall s
      . (Excluding '[SharedKey] (KeyOf s))
     => WalletFlavorS s
-    -> UTxOAssumptions
+    -> Write.UTxOAssumptions
 utxoAssumptionsForWallet = keyOfWallet >>> \case
-    ByronKeyS -> AllByronKeyPaymentCredentials
-    IcarusKeyS -> AllByronKeyPaymentCredentials
-    ShelleyKeyS -> AllKeyPaymentCredentials
+    ByronKeyS -> Write.AllByronKeyPaymentCredentials
+    IcarusKeyS -> Write.AllByronKeyPaymentCredentials
+    ShelleyKeyS -> Write.AllKeyPaymentCredentials