Factor out pure balanceTransaction logic

for easier testability.

This does break the pattern of the Wallet module of using
`HasTransactionLayer`, `HasLogger`, etc, rather than real arguments,
but:
1. I didn't want to require the full NetworkLayer, only the pparams.
2. I was unsure about adding a `m` parameter to the `HasLogger`
   constraint.

lib/core/src/Cardano/Wallet.hs

@@ -154,6 +154,7 @@ module Cardano.Wallet
     , getTransaction
     , submitExternalTx
     , submitTx
+    , balanceTransaction
     , LocalTxSubmissionConfig (..)
     , defaultLocalTxSubmissionConfig
     , runLocalTxSubmissionPool
@@ -297,6 +298,7 @@ import Cardano.Wallet.Primitive.AddressDiscovery.Shared
     )
 import Cardano.Wallet.Primitive.CoinSelection
     ( Selection
+    , SelectionCollateralRequirement (SelectionCollateralNotRequired, SelectionCollateralRequired)
     , SelectionConstraints (..)
     , SelectionError (..)
     , SelectionOf (..)
@@ -307,9 +309,10 @@ import Cardano.Wallet.Primitive.CoinSelection
     , makeSelectionReportDetailed
     , makeSelectionReportSummarized
     , performSelection
+    , selectionDelta
     )
 import Cardano.Wallet.Primitive.CoinSelection.Balance
-    ( emptySkeleton )
+    ( SelectionSkeleton (SelectionSkeleton), emptySkeleton )
 import Cardano.Wallet.Primitive.Collateral
     ( asCollateral )
 import Cardano.Wallet.Primitive.Migration
@@ -343,6 +346,7 @@ import Cardano.Wallet.Primitive.Types
     , Block (..)
     , BlockHeader (..)
     , DelegationCertificate (..)
+    , FeePolicy (LinearFee)
     , GenesisParameters (..)
     , IsDelegatingTo (..)
     , NetworkParameters (..)
@@ -380,7 +384,7 @@ import Cardano.Wallet.Primitive.Types.Tx
     , LocalTxSubmissionStatus
     , SealedTx (..)
     , TransactionInfo (..)
-    , Tx
+    , Tx (..)
     , TxChange (..)
     , TxIn (..)
     , TxMeta (..)
@@ -407,6 +411,7 @@ import Cardano.Wallet.Transaction
     , ErrUpdateSealedTx (..)
     , TransactionCtx (..)
     , TransactionLayer (..)
+    , TxUpdate (..)
     , Withdrawal (..)
     , defaultTransactionCtx
     , withdrawalToCoin
@@ -518,6 +523,7 @@ import UnliftIO.Exception
 import UnliftIO.MVar
     ( modifyMVar_, newMVar )
 
+import qualified Cardano.Api.Shelley as Cardano
 import qualified Cardano.Crypto.Wallet as CC
 import qualified Cardano.Wallet.Primitive.AddressDiscovery.Random as Rnd
 import qualified Cardano.Wallet.Primitive.AddressDiscovery.Sequential as Seq
@@ -531,6 +537,8 @@ import qualified Cardano.Wallet.Primitive.Types.TokenMap as TokenMap
 import qualified Cardano.Wallet.Primitive.Types.UTxO as UTxO
 import qualified Cardano.Wallet.Primitive.Types.UTxOIndex as UTxOIndex
 import qualified Cardano.Wallet.Primitive.Types.UTxOSelection as UTxOSelection
+import Control.Monad.Random
+    ( MonadRandom )
 import qualified Data.ByteArray as BA
 import qualified Data.Foldable as F
 import qualified Data.List as L
@@ -1264,6 +1272,139 @@ normalizeDelegationAddress s addr = do
                                   Transaction
 -------------------------------------------------------------------------------}
 
+balanceTransaction
+    :: forall m s k. (GenChange s, MonadRandom m)
+    => Tracer m WalletLog
+    -> ArgGenChange s
+    -> (W.ProtocolParameters, Cardano.ProtocolParameters)
+    -> (UTxOIndex, Wallet s, Set Tx)
+    -> TransactionLayer k SealedTx
+    -> [(TxIn, TxOut)] -- ^ Resolved inputs
+    -> SealedTx
+    -> ExceptT ErrBalanceTx m SealedTx
+balanceTransaction
+    tr
+    generateChange
+    (pp, nodePParams)
+    (internalUtxoAvailable, wallet, pendingTxs)
+    tl
+    resolvedInputs
+    partialTx
+    = do
+    let (outputs, txWithdrawal, txMetadata) = extractFromTx partialTx
+
+    let externalSelectedUtxo =
+            UTxOIndex.fromSequence resolvedInputs
+
+    let utxoAvailableForInputs = UTxOSelection.fromIndexPair
+            (internalUtxoAvailable, externalSelectedUtxo)
+
+    let utxoAvailableForCollateral =
+            UTxOIndex.toUTxO internalUtxoAvailable
+
+    -- NOTE: It is not possible to know the script execution cost in
+    -- advance because it actually depends on the final transaction. Inputs
+    -- selected as part of the fee balancing might have an influence on the
+    -- execution cost.
+    -- However, they are bounded so it is possible to balance the
+    -- transaction considering only the maximum cost, and only after, try to
+    -- adjust the change and ExUnits of each redeemer to something more
+    -- sensible than the max execution cost.
+    let txPlutusScriptExecutionCost = maxScriptExecutionCost tl pp partialTx
+    let txContext = defaultTransactionCtx
+            { txPlutusScriptExecutionCost
+            , txMetadata
+            , txWithdrawal
+            , txCollateralRequirement =
+                if txPlutusScriptExecutionCost > Coin 0 then
+                    SelectionCollateralRequired
+                else
+                    SelectionCollateralNotRequired
+            } & padFeeEstimation partialTx
+
+        -- FIXME: The coin selection and reported fees will likely be wrong in
+        -- the presence of certificates (and deposits / refunds). An immediate
+        -- "fix" is to return a proper error from the handler when any key or
+        -- pool registration (resp. deregistration) certificate is found in the
+        -- transaction. A long-term fix is to handle this case properly during
+        -- balancing.
+    let transform s sel =
+            let (sel', _) = assignChangeAddresses @s generateChange sel s
+             in ( const (selectionDelta txOutCoin sel')
+                , fst <$> F.toList (sel' ^. #inputs)
+                , fst <$> (sel' ^. #collateral)
+                , sel' ^. #change
+                )
+    (newFee, extraInputs, extraCollateral, extraOutputs) <-
+        withExceptT ErrBalanceTxSelectAssets $ selectAssets'
+            tr
+            tl
+            pp
+            SelectAssetsParams
+                { outputs
+                , pendingTxs
+                , txContext
+                , utxoAvailableForInputs
+                , utxoAvailableForCollateral
+                , wallet
+                }
+            transform
+
+    let txUpdate =
+            TxUpdate { extraInputs, extraCollateral, extraOutputs, newFee, newExUnits }
+          where
+            -- FIXME: At this stage, we set all execution units for all redeemers to the
+            -- max cost, which is guaranteed to succeed (given the coin selection above
+            -- was done with the same assumption) but also terribly ineffective when it
+            -- comes to reducing the cost. This is however sufficient to start
+            -- preliminary integration work.
+            newExUnits = const (const (pp ^. #txParameters . #getMaxExecutionUnits))
+
+    withExceptT ErrBalanceTxUpdateError . ExceptT . return $
+        updateTx tl nodePParams partialTx txUpdate
+  where
+    -- | Wallet coin selection is unaware of many kinds of transaction content
+    -- (e.g. datums, redeemers), which could be included in the input to
+    -- `balanceTransaction`. As a workaround we add some padding using
+    -- `evaluateMinimumFee`.
+    --
+    -- TODO: This logic needs to be consistent with how we call `selectAssets`,
+    -- so it would be good to join them into some single helper.
+    padFeeEstimation
+        :: SealedTx
+        -> TransactionCtx
+        -> TransactionCtx
+    padFeeEstimation sealedTx txCtx =
+        let
+            walletTx = decodeTx tl sealedTx
+            worseEstimate = calcMinimumCost tl pp txCtx skeleton
+            skeleton = SelectionSkeleton
+                { skeletonInputCount = length (view #resolvedInputs walletTx)
+                , skeletonOutputs = view #outputs walletTx
+                , skeletonChange = mempty
+                }
+            LinearFee _ (Quantity b) = pp ^. #txParameters . #getFeePolicy
+            -- NOTE: Coping with the later additions of script integrity hash and
+            -- redeemers ex units increased from 0 to their actual values.
+            extraMargin = Coin $ ceiling (100 * b)
+            txFeePadding = (<> extraMargin) $ fromMaybe (Coin 0) $ do
+                betterEstimate <- evaluateMinimumFee tl nodePParams sealedTx
+                betterEstimate `Coin.subtractCoin` worseEstimate
+        in
+            txCtx { txFeePadding }
+    extractFromTx tx =
+        let (Tx _id _fee _coll _inps outs wdrlMap meta _vldt) = decodeTx tl tx
+            -- TODO: Find a better abstraction that can cover this case.
+            wdrl = WithdrawalSelf
+                (error $ "WithdrawalSelf: reward-account should never been use "
+                      <> "when balancing transactions but it was!"
+                )
+                (error $ "WithdrawalSelf: derivation path should never been use "
+                      <> "when balancing transactions but it was!"
+                )
+                (sumCoins wdrlMap)
+         in (outs, wdrl, meta)
+
 -- | Augments the given outputs with new outputs. These new outputs correspond
 -- to change outputs to which new addresses have been assigned. This updates
 -- the wallet state as it needs to keep track of new pending change addresses.
@@ -1441,9 +1582,25 @@ selectAssets
     -> (s -> Selection -> result)
     -> ExceptT ErrSelectAssets IO result
 selectAssets ctx params transform = do
-    guardPendingWithdrawal
     pp <- liftIO $ currentProtocolParameters nl
-    liftIO $ traceWith tr $ MsgSelectionStart
+    selectAssets' tr tl pp params transform
+  where
+    nl = ctx ^. networkLayer
+    tl = ctx ^. transactionLayer @k
+    tr = contramap MsgWallet $ ctx ^. logger
+
+-- | Less convenient and restrictive version of `selectAssets`.
+selectAssets'
+    :: forall m s k result. MonadRandom m
+    => Tracer m WalletLog
+    -> TransactionLayer k SealedTx
+    -> ProtocolParameters
+    -> SelectAssetsParams s result
+    -> (s -> Selection -> result)
+    -> ExceptT ErrSelectAssets m result
+selectAssets' tr tl pp params transform = do
+    guardPendingWithdrawal
+    lift $ traceWith tr $ MsgSelectionStart
         (UTxOSelection.availableUTxO $ params ^. #utxoAvailableForInputs)
         (params ^. #outputs)
     mSel <- runExceptT $ performSelection
@@ -1491,26 +1648,23 @@ selectAssets ctx params transform = do
                 params ^. #utxoAvailableForInputs
             }
     case mSel of
-        Left e -> liftIO $
+        Left e -> lift $
             traceWith tr $ MsgSelectionError e
-        Right sel -> liftIO $ do
+        Right sel -> lift $ do
             traceWith tr $ MsgSelectionReportSummarized
                 $ makeSelectionReportSummarized sel
             traceWith tr $ MsgSelectionReportDetailed
                 $ makeSelectionReportDetailed sel
     withExceptT ErrSelectAssetsSelectionError $ except $
         transform (getState $ params ^. #wallet) <$> mSel
   where
-    nl = ctx ^. networkLayer
-    tl = ctx ^. transactionLayer @k
-    tr = contramap MsgWallet $ ctx ^. logger
 
     -- Ensure that there's no existing pending withdrawals. Indeed, a withdrawal
     -- is necessarily withdrawing rewards in their totality. So, after a first
     -- withdrawal is executed, the reward pot is empty. So, to prevent two
     -- transactions with withdrawals to go through (which will inevitably cause
     -- one of them to never be inserted), we warn users early on about it.
-    guardPendingWithdrawal :: ExceptT ErrSelectAssets IO ()
+    guardPendingWithdrawal :: ExceptT ErrSelectAssets m ()
     guardPendingWithdrawal =
         case Set.lookupMin $ Set.filter hasWithdrawal $ params ^. #pendingTxs of
             Just pendingWithdrawal
@@ -2575,6 +2729,7 @@ data ErrSignPayment
 data ErrBalanceTx
     = ErrBalanceTxTxAlreadyBalanced
     | ErrBalanceTxUpdateError ErrUpdateSealedTx
+    | ErrBalanceTxSelectAssets ErrSelectAssets
     | ErrBalanceTxNotImplemented
     deriving (Show, Eq)