@@ -26,10 +26,10 @@ import Control.Monad.Trans.Either (EitherT, hoistEither, newEitherT, runEitherT)
2626import Data.Coerce (coerce )
2727import Data.Either.Combinators (rightToMaybe )
2828import Data.Kind (Type )
29- import Data.List (partition , (\\) )
29+ import Data.List ((\\) )
3030import Data.Map (Map )
3131import Data.Map qualified as Map
32- import Data.Maybe (fromMaybe , mapMaybe )
32+ import Data.Maybe (fromJust , fromMaybe , mapMaybe )
3333import Data.Set (Set )
3434import Data.Set qualified as Set
3535import Data.Text (Text )
@@ -311,17 +311,17 @@ addTxCollaterals utxos tx =
311311handleNonAdaChange :: Address -> Map TxOutRef TxOut -> Tx -> Either Text Tx
312312handleNonAdaChange changeAddr utxos tx =
313313 let nonAdaChange = getNonAdaChange utxos tx
314+ newOutput =
315+ TxOut
316+ { txOutAddress = changeAddr
317+ , txOutValue = nonAdaChange
318+ , txOutDatumHash = Nothing
319+ }
314320 outputs =
315- case partition ((==) changeAddr . Tx. txOutAddress) $ txOutputs tx of
316- ([] , txOuts) ->
317- TxOut
318- { txOutAddress = changeAddr
319- , txOutValue = nonAdaChange
320- , txOutDatumHash = Nothing
321- } :
322- txOuts
323- (txOut@ TxOut {txOutValue = v} : txOuts, txOuts') ->
324- txOut {txOutValue = v <> nonAdaChange} : (txOuts <> txOuts')
321+ modifyFirst
322+ ((==) changeAddr . Tx. txOutAddress)
323+ (maybe newOutput $ addValueToTxOut nonAdaChange)
324+ (txOutputs tx)
325325 in if isValueNat nonAdaChange
326326 then Right $ if Value. isZero nonAdaChange then tx else tx {txOutputs = outputs}
327327 else Left " Not enough inputs to balance tokens."
@@ -332,18 +332,29 @@ hasChangeUTxO changeAddr tx =
332332
333333-- | Adds ada change to a transaction, assuming there is already an output going to ownPkh. Otherwise, this is identity
334334addAdaChange :: Address -> Integer -> Tx -> Tx
335+ addAdaChange _ 0 tx = tx
335336addAdaChange changeAddr change tx =
336337 tx
337338 { txOutputs =
338- case partition ( (==) changeAddr . Tx. txOutAddress) $ txOutputs tx of
339- (txOut @ TxOut {txOutValue = v} : txOuts, txOuts') ->
340- txOut {txOutValue = v <> Ada. lovelaceValueOf change} : (txOuts <> txOuts' )
341- _ -> txOutputs tx
339+ modifyFirst
340+ ((==) changeAddr . Tx. txOutAddress)
341+ (addValueToTxOut ( Ada. lovelaceValueOf change) . fromJust )
342+ ( txOutputs tx)
342343 }
343344
345+ {- | Modifies the first element matching a predicate, or, if none found, call the modifier with Nothing
346+ Calling this function ensures the modifier will always be run once
347+ -}
348+ modifyFirst :: (a -> Bool ) -> (Maybe a -> a ) -> [a ] -> [a ]
349+ modifyFirst _ m [] = [m Nothing ]
350+ modifyFirst p m (x : xs) = if p x then m (Just x) : xs else x : modifyFirst p m xs
351+
352+ addValueToTxOut :: Value -> TxOut -> TxOut
353+ addValueToTxOut val txOut = txOut {txOutValue = txOutValue txOut <> val}
354+
344355-- | Adds a 1 lovelace output to a transaction
345356addOutput :: Address -> Tx -> Tx
346- addOutput changeAddr tx = tx {txOutputs = changeTxOut : txOutputs tx}
357+ addOutput changeAddr tx = tx {txOutputs = txOutputs tx ++ [changeTxOut] }
347358 where
348359 changeTxOut =
349360 TxOut
0 commit comments