Generalize distributeSurplus.

lib/core/src/Cardano/Wallet.hs

@@ -433,7 +433,6 @@ import Cardano.Wallet.Primitive.Types.Tx
     , TxStatus (..)
     , UnsignedTx (..)
     , fromTransactionInfo
-    , txOutAddCoin
     , txOutCoin
     , withdrawals
     )
@@ -461,8 +460,6 @@ import Cardano.Wallet.Transaction
     , defaultTransactionCtx
     , withdrawalToCoin
     )
-import Cardano.Wallet.Util
-    ( mapFirst )
 import Control.Applicative
     ( (<|>) )
 import Control.Arrow
@@ -539,7 +536,7 @@ import Data.List.NonEmpty
 import Data.Map.Strict
     ( Map )
 import Data.Maybe
-    ( fromMaybe, isJust, listToMaybe, mapMaybe )
+    ( fromMaybe, isJust, mapMaybe )
 import Data.Proxy
     ( Proxy )
 import Data.Quantity
@@ -1702,26 +1699,21 @@ balanceTransactionWithSelectionStrategy
 
     let feeAndChange = TxFeeAndChange
             (unsafeFromLovelace candidateMinFee)
-            (txOutCoin <$> listToMaybe extraOutputs)
+            (extraOutputs)
     let feePolicy = view (#txParameters . #getFeePolicy) pp
 
     -- @distributeSurplus@ should never fail becase we have provided enough
     -- padding in @selectAssets'@.
-    TxFeeAndChange extraFee extraChange <-
-        withExceptT
-            (\(ErrMoreSurplusNeeded c) ->
-                ErrBalanceTxNotYetSupported $ UnderestimatedFee c candidateTx)
-            (ExceptT . pure $
-                distributeSurplus tl feePolicy surplus feeAndChange)
+    TxFeeAndChange updatedFee updatedChange <- withExceptT
+        (\(ErrMoreSurplusNeeded c) ->
+            ErrBalanceTxNotYetSupported $ UnderestimatedFee c candidateTx)
+        (ExceptT . pure $ distributeSurplus tl feePolicy surplus feeAndChange)
 
     guardTxSize =<< guardTxBalanced =<< (assembleTransaction $ TxUpdate
         { extraInputs
         , extraCollateral
-        , extraOutputs = mapFirst
-            (txOutAddCoin $ fromMaybe (Coin 0) extraChange)
-            extraOutputs
-        , feeUpdate = UseNewTxFee
-            (unsafeFromLovelace candidateMinFee <> extraFee)
+        , extraOutputs = updatedChange
+        , feeUpdate = UseNewTxFee updatedFee
         })
   where
     tl = ctx ^. transactionLayer @k

lib/core/src/Cardano/Wallet/Primitive/Types/Tx.hs

@@ -67,6 +67,7 @@ module Cardano.Wallet.Primitive.Types.Tx
     , txMetadataIsNull
     , txOutCoin
     , txOutAddCoin
+    , txOutSubtractCoin
     , failedScriptValidation
 
     -- * Constants
@@ -144,7 +145,7 @@ import Data.Either
 import Data.Function
     ( on, (&) )
 import Data.Generics.Internal.VL.Lens
-    ( view )
+    ( over, view )
 import Data.Generics.Labels
     ()
 import Data.Int
@@ -330,11 +331,29 @@ data TxOut = TxOut
 txOutCoin :: TxOut -> Coin
 txOutCoin = TokenBundle.getCoin . view #tokens
 
--- Add a fixed coin value to an existing output.
+-- | Increments the 'Coin' value of a 'TxOut'.
+--
+-- Satisfies the following property for all values of 'c':
+--
+-- >>> txOutSubtractCoin c . txOutAddCoin c == id
+--
 txOutAddCoin :: Coin -> TxOut -> TxOut
 txOutAddCoin val (TxOut addr tokens) =
     TxOut addr (tokens <> TokenBundle.fromCoin val)
 
+-- | Decrements the 'Coin' value of a 'TxOut'.
+--
+-- Satisfies the following property for all values of 'c':
+--
+-- >>> txOutSubtractCoin c . txOutAddCoin c == id
+--
+-- If the given 'Coin' is greater than the 'Coin' value of the given 'TxOut',
+-- the resulting 'TxOut' will have a 'Coin' value of zero.
+--
+txOutSubtractCoin :: Coin -> TxOut -> TxOut
+txOutSubtractCoin toSubtract =
+    over (#tokens . #coin) (`Coin.difference` toSubtract)
+
 -- Since the 'TokenBundle' type deliberately does not provide an 'Ord' instance
 -- (as that would lead to arithmetically invalid orderings), this means we can't
 -- automatically derive an 'Ord' instance for the 'TxOut' type.

lib/core/src/Cardano/Wallet/Transaction.hs

@@ -6,6 +6,7 @@
 {-# LANGUAGE DuplicateRecordFields #-}
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeFamilies #-}
@@ -35,6 +36,7 @@ module Cardano.Wallet.Transaction
     , PlutusScriptInfo (..)
     , PlutusVersion (..)
     , TxFeeAndChange (..)
+    , mapTxFeeAndChange
 
     -- * Errors
     , ErrSignTx (..)
@@ -257,32 +259,26 @@ data TransactionLayer k tx = TransactionLayer
 
     , distributeSurplus
         :: FeePolicy
-        -> Coin -- Surplus to distribute
-        -> TxFeeAndChange -- Fee and value of relevant change output (if any)
-        -> Either ErrMoreSurplusNeeded TxFeeAndChange
-        -- ^ Distribute a surplus transaction balance between a given change
-        -- output (if one exists present) and the transaction fee. The function
-        -- is aware of the fact that any increase of 'Coin' values could
-        -- increase the size and fee-requirement of the transaction.
-        --
-        -- This helper is used from 'balanceTransaction'.
+        -> Coin
+        -- ^ Surplus transaction balance to distribute.
+        -> TxFeeAndChange [TxOut]
+        -- ^ Original fee and change outputs.
+        -> Either ErrMoreSurplusNeeded (TxFeeAndChange [TxOut])
+        -- ^ Adjusted fee and change outputs.
         --
-        -- >>> distributeSurplus feePolicy (Coin 100) (TxFeeAndChange (Coin 200) (Coin 200))
-        -- TxFeeAndChange
-        --    { fee = Coin 1
-        --    , change = Coin 99
-        --    }
+        -- Distributes a surplus transaction balance between the given change
+        -- outputs and the given fee. This function is aware of the fact that
+        -- any increase in a 'Coin' value could increase the size and fee
+        -- requirement of a transaction.
         --
-        -- >>> distributeSurplus feePolicy (Coin 100) (TxFeeAndChange (Coin 255) (Coin 200))
-        -- TxFeeAndChange
-        --    { fee = Coin 2
-        --    , change = Coin 98
-        --    }
+        -- When comparing the original fee and change outputs to the adjusted
+        -- fee and change outputs, this function guarantees that:
         --
-        -- Important note: the return value is a delta. In particular a returned
-        -- change value of @Nothing@ or @Just (Coin 0)@ does **not** mean the
-        -- change should be set to @Coin 0@, but rather that the change should
-        -- not be increased!
+        --  - The number of the change outputs remains constant;
+        --  - The fee quantity either remains the same or increases.
+        --  - For each change output:
+        --      - the ada quantity either remains constant or increases.
+        --      - non-ada quantities remain the same.
 
     , computeSelectionLimit
         :: ProtocolParameters
@@ -518,7 +514,22 @@ newtype ErrMoreSurplusNeeded = ErrMoreSurplusNeeded Coin
 
 -- | Small helper record to disambiguate between a fee and change Coin values.
 -- Used by 'distributeSurplus'.
-data TxFeeAndChange = TxFeeAndChange
+data TxFeeAndChange change = TxFeeAndChange
     { fee :: Coin
-    , change :: Maybe Coin
-    } deriving (Show, Eq)
+    , change :: change
+    }
+    deriving (Eq, Show)
+
+-- | Manipulates a 'TxFeeAndChange' value.
+--
+mapTxFeeAndChange
+    :: (Coin -> Coin)
+    -- ^ A function to transform the fee
+    -> (change1 -> change2)
+    -- ^ A function to transform the change
+    -> TxFeeAndChange change1
+    -- ^ The original fee and change
+    -> TxFeeAndChange change2
+    -- ^ The transformed fee and change
+mapTxFeeAndChange mapFee mapChange TxFeeAndChange {fee, change} =
+    TxFeeAndChange (mapFee fee) (mapChange change)

lib/shelley/src/Cardano/Wallet/Shelley/Transaction.hs

@@ -55,6 +55,7 @@ module Cardano.Wallet.Shelley.Transaction
     , txConstraints
     , costOfIncreasingCoin
     , _distributeSurplus
+    , distributeSurplusDelta
     , sizeOfCoin
     , maximumCostOfIncreasingCoin
     ) where
@@ -143,6 +144,7 @@ import Cardano.Wallet.Primitive.Types.Tx
     , TxSize (..)
     , sealedTxFromCardano'
     , sealedTxFromCardanoBody
+    , txOutAddCoin
     , txOutCoin
     , txSizeDistance
     )
@@ -186,6 +188,7 @@ import Cardano.Wallet.Transaction
     , TxFeeAndChange (..)
     , TxFeeUpdate (..)
     , TxUpdate (..)
+    , mapTxFeeAndChange
     , withdrawalToCoin
     )
 import Cardano.Wallet.Util
@@ -207,7 +210,7 @@ import Data.Bifunctor
 import Data.Function
     ( (&) )
 import Data.Functor
-    ( ($>) )
+    ( ($>), (<&>) )
 import Data.Functor.Identity
     ( runIdentity )
 import Data.Generics.Internal.VL.Lens
@@ -626,7 +629,6 @@ newTransactionLayer networkId = TransactionLayer
     , maxScriptExecutionCost =
         _maxScriptExecutionCost
 
-
     , distributeSurplus = _distributeSurplus
 
     , assignScriptRedeemers =
@@ -1491,33 +1493,71 @@ sizeOfCoin (Coin c)
     | c >=            24 = TxSize 2
     | otherwise          = TxSize 1
 
--- | Actual implementation for 'distributeSurplus'.
+-- | Distributes a surplus transaction balance between the given change outputs
+--   and the given fee.
+--
+-- See documentation for 'TransactionLayer.distributeSurplus' for more details.
+--
 _distributeSurplus
+    :: FeePolicy
+    -> Coin
+    -- ^ Surplus transaction balance to distribute.
+    -> TxFeeAndChange [TxOut]
+    -- ^ Original fee and change outputs.
+    -> Either ErrMoreSurplusNeeded (TxFeeAndChange [TxOut])
+    -- ^ Adjusted fee and change outputs.
+_distributeSurplus feePolicy surplus fc@(TxFeeAndChange fee change) =
+    distributeSurplusDelta feePolicy surplus
+        (mapTxFeeAndChange id (fmap txOutCoin) fc)
+    <&> mapTxFeeAndChange
+        (fee <>)
+        (zipWith (flip txOutAddCoin) change)
+
+distributeSurplusDelta
+    :: FeePolicy
+    -> Coin
+    -- ^ Surplus to distribute
+    -> TxFeeAndChange [Coin]
+    -> Either ErrMoreSurplusNeeded (TxFeeAndChange [Coin])
+distributeSurplusDelta feePolicy surplus (TxFeeAndChange fee change) =
+    case change of
+        changeHead : changeTail ->
+            distributeSurplusDeltaWithOneChangeCoin feePolicy surplus
+                (TxFeeAndChange fee changeHead)
+            <&> mapTxFeeAndChange id
+                (: (Coin 0 <$ changeTail))
+        [] ->
+            burnSurplusAsFees feePolicy surplus
+                (TxFeeAndChange fee ())
+            <&> mapTxFeeAndChange id
+                (\() -> [])
+
+distributeSurplusDeltaWithOneChangeCoin
     :: FeePolicy
     -> Coin -- ^ Surplus to distribute
-    -> TxFeeAndChange
-    -> Either ErrMoreSurplusNeeded TxFeeAndChange
-_distributeSurplus feePolicy surplus fc@(TxFeeAndChange _fee0 Nothing) =
-    burnSurplusAsFees feePolicy surplus fc
-_distributeSurplus feePolicy surplus fc@(TxFeeAndChange fee0 (Just change0)) =
+    -> TxFeeAndChange Coin
+    -> Either ErrMoreSurplusNeeded (TxFeeAndChange Coin)
+distributeSurplusDeltaWithOneChangeCoin
+    feePolicy surplus fc@(TxFeeAndChange fee0 change0) =
     let
         -- We calculate the maximum possible fee increase, by assuming the
         -- **entire** surplus is added to the change.
         extraFee = findFixpointIncreasingFeeBy $
             costOfIncreasingCoin feePolicy change0 surplus
-
     in
         case surplus `Coin.subtract` extraFee of
             Just extraChange ->
                 Right $ TxFeeAndChange
                     { fee = extraFee
-                    , change = Just extraChange
+                    , change = extraChange
                     }
             Nothing ->
                 -- The fee increase from adding the surplus to the change was
                 -- greater than the surplus itself. This could happen if the
                 -- surplus is small.
-                burnSurplusAsFees feePolicy surplus fc
+                burnSurplusAsFees feePolicy surplus
+                    (mapTxFeeAndChange id (const ()) fc)
+                        <&> mapTxFeeAndChange id (\() -> Coin 0)
   where
     -- Increasing the fee may itself increase the fee. If that is the case, this
     -- function will increase the fee further. The process repeats until the fee
@@ -1567,13 +1607,13 @@ _distributeSurplus feePolicy surplus fc@(TxFeeAndChange fee0 (Just change0)) =
 burnSurplusAsFees
     :: FeePolicy
     -> Coin -- Surplus
-    -> TxFeeAndChange
-    -> Either ErrMoreSurplusNeeded TxFeeAndChange
-burnSurplusAsFees feePolicy surplus (TxFeeAndChange fee0 _) =
+    -> TxFeeAndChange ()
+    -> Either ErrMoreSurplusNeeded (TxFeeAndChange ())
+burnSurplusAsFees feePolicy surplus (TxFeeAndChange fee0 ()) =
     case costOfBurningSurplus `Coin.subtract` surplus of
         Just shortfall -> Left $ ErrMoreSurplusNeeded shortfall
         Nothing ->
-            Right $ TxFeeAndChange surplus Nothing
+            Right $ TxFeeAndChange surplus ()
   where
     costOfBurningSurplus = costOfIncreasingCoin feePolicy fee0 surplus