#490 Cache collateral utxos

satran004 · satran004 · commit ba45b8caa0a9 · 2025-02-17T10:50:33.000+08:00
diff --git a/function/src/main/java/com/bloxbean/cardano/client/function/TxBuilderContext.java b/function/src/main/java/com/bloxbean/cardano/client/function/TxBuilderContext.java
@@ -47,6 +47,7 @@ public class TxBuilderContext {
     //Stores utxos used in the transaction.
     //This list is cleared after each build() call.
     private Set<Utxo> utxos = new HashSet<>();
+    private Set<Utxo> collateralUtxos = new HashSet<>();
 
     @Setter(AccessLevel.NONE)
     private Map<String, Tuple<PlutusScript, byte[]>> refScripts = new HashMap<>();
@@ -205,6 +206,24 @@ public void clearUtxos() {
         utxos.clear();
     }
 
+    public void addCollateralUtxo(Utxo utxo) {
+        collateralUtxos.add(utxo);
+    }
+
+    public Set<Utxo> getCollateralUtxos() {
+        return collateralUtxos;
+    }
+
+    public void clearCollateralUtxos() {
+        collateralUtxos.clear();
+    }
+
+    public Set<Utxo> getAllUtxos() {
+        Set<Utxo> allUtxos = new HashSet<>(utxos);
+        allUtxos.addAll(collateralUtxos);
+        return allUtxos;
+    }
+
     @SneakyThrows
     public void addRefScripts(PlutusScript plutusScript) {
         refScripts.put(HexUtil.encodeHexString(plutusScript.getScriptHash()), new Tuple<>(plutusScript, plutusScript.scriptRefBytes()));
@@ -307,6 +326,7 @@ private Transaction buildTransaction(TxBuilder txBuilder) {
     private void clearTempStates() {
         clearMintMultiAssets();
         clearUtxos();
+        clearCollateralUtxos();
         clearRefScripts();
     }
 }
diff --git a/function/src/main/java/com/bloxbean/cardano/client/function/helper/ScriptBalanceTxProviders.java b/function/src/main/java/com/bloxbean/cardano/client/function/helper/ScriptBalanceTxProviders.java
@@ -1,6 +1,8 @@
 package com.bloxbean.cardano.client.function.helper;
 
+import com.bloxbean.cardano.client.api.AddressIterator;
 import com.bloxbean.cardano.client.api.TransactionEvaluator;
+import com.bloxbean.cardano.client.api.common.AddressIterators;
 import com.bloxbean.cardano.client.function.TxBuilder;
 import com.bloxbean.cardano.client.function.exception.TxBuildException;
 import com.bloxbean.cardano.client.transaction.spec.Value;
@@ -15,20 +17,38 @@
 public class ScriptBalanceTxProviders {
 
     //TODO -- Unit tests pending
+
     /**
      * Function to balance an unbalanced transaction using Automatic Utxo Discovery with Additional Signers.
      * This function invokes {@link BalanceTxBuilders#balanceTxWithAdditionalSigners(String, int)} to balance the transaction.
      * If any new inputs are added to the transaction during balancing, the script cost and fee will be recomputed.
      *
      * @param feePayer Fee payer address
+     * @param additionalSigners o of Additional signers. This is required for accurate fee calculation.
+     * @param containsScript If the transaction contains script
+     * @return
+     */
+    public static TxBuilder balanceTx(String feePayer, int additionalSigners, boolean containsScript) {
+        return balanceTx(AddressIterators.of(feePayer), additionalSigners, containsScript);
+    }
+
+    /**
+     * Function to balance an unbalanced transaction using Automatic Utxo Discovery with Additional Signers.
+     * This function invokes {@link BalanceTxBuilders#balanceTxWithAdditionalSigners(String, int)} to balance the transaction.
+     * If any new inputs are added to the transaction during balancing, the script cost and fee will be recomputed.
+     *
+     * @param feePayerAddrIter Fee payer address iterator
      * @param additionalSigners No of Additional signers. This is required for accurate fee calculation.
      * @param containsScript If the transaction contains script
      * @return TxBuilder
      */
-    public static TxBuilder balanceTx(String feePayer, int additionalSigners, boolean containsScript) {
+    public static TxBuilder balanceTx(AddressIterator feePayerAddrIter, int additionalSigners, boolean containsScript) {
         return (ctx, transaction) -> {
+
+            String feePayerAddr = feePayerAddrIter.getFirst().getAddress();
+
             int inputSize = transaction.getBody().getInputs().size();
-            BalanceTxBuilders.balanceTxWithAdditionalSigners(feePayer, additionalSigners).apply(ctx, transaction);
+            BalanceTxBuilders.balanceTxWithAdditionalSigners(feePayerAddrIter.clone(), additionalSigners).apply(ctx, transaction);
             int newInputSize = transaction.getBody().getInputs().size();
 
             if (!containsScript || (newInputSize == inputSize)) { //TODO -- check for contains script
@@ -45,7 +65,7 @@ public static TxBuilder balanceTx(String feePayer, int additionalSigners, boolea
             BigInteger fee = transaction.getBody().getFee();
             transaction.getBody().setFee(BigInteger.valueOf(170000));
             transaction.getBody().getOutputs().stream()
-                    .filter(output -> feePayer.equals(output.getAddress()))
+                    .filter(output -> feePayerAddr.equals(output.getAddress()))
                     .max((to1, to2) -> to1.getValue().getCoin().compareTo(to2.getValue().getCoin()))
                     .ifPresent(transactionOutput -> transactionOutput.setValue(transactionOutput.getValue().add(Value.builder().coin(fee).build())));
 
@@ -56,7 +76,7 @@ public static TxBuilder balanceTx(String feePayer, int additionalSigners, boolea
 
             //Recompute script cost
             ScriptCostEvaluators.evaluateScriptCost().apply(ctx, transaction);
-            BalanceTxBuilders.balanceTxWithAdditionalSigners(feePayer, additionalSigners).apply(ctx, transaction);
+            BalanceTxBuilders.balanceTxWithAdditionalSigners(feePayerAddrIter, additionalSigners).apply(ctx, transaction);
         };
     }
 }
diff --git a/quicktx/src/main/java/com/bloxbean/cardano/client/quicktx/QuickTxBuilder.java b/quicktx/src/main/java/com/bloxbean/cardano/client/quicktx/QuickTxBuilder.java
@@ -23,6 +23,8 @@
 import com.bloxbean.cardano.client.transaction.spec.TransactionInput;
 import com.bloxbean.cardano.client.util.JsonUtil;
 import com.bloxbean.cardano.client.util.Tuple;
+import com.bloxbean.cardano.hdwallet.Wallet;
+import com.bloxbean.cardano.hdwallet.util.HDWalletAddressIterator;
 import lombok.NonNull;
 import lombok.extern.slf4j.Slf4j;
 
@@ -147,7 +149,9 @@ public TxContext compose(AbstractTx... txs) {
     public class TxContext {
         private AbstractTx[] txList;
         private String feePayer;
+        private Wallet feePayerWallet;
         private String collateralPayer;
+        private Wallet collateralPayerWallet;
         private Set<byte[]> requiredSigners;
         private Set<TransactionInput> collateralInputs;
 
@@ -183,18 +187,70 @@ public class TxContext {
          * When there are more than one txs, fee payer address is mandatory.
          *
          * @param address fee payer address
-         * @return TxContext
+         * @return TxContext - the updated transaction context with the fee payer set
+         * @throws TxBuildException if the fee payer has already been set
          */
         public TxContext feePayer(String address) {
+            if (feePayerWallet != null || feePayer != null)
+                throw new TxBuildException("The fee payer has already been set. It can only be set once.");
+
             this.feePayer = address;
             return this;
         }
 
+        /**
+         * Sets the fee payer wallet for the transaction. When there is only one tx, sender's address is used as fee payer address.
+         * When there are more than one txs, fee payer address/wallet is mandatory.
+         *
+         * @param feePayerWallet the wallet that will act as the fee payer for the transaction
+         * @return TxContext - the updated transaction context with the fee payer set
+         * @throws TxBuildException if the fee payer has already been set
+         */
+        public TxContext feePayer(Wallet feePayerWallet) {
+            if (feePayerWallet != null || feePayer != null)
+                throw new TxBuildException("The fee payer has already been set. It can only be set once.");
+
+            this.feePayerWallet = feePayerWallet;
+            // TODO feePayer is not used in this scenarios, but it must be set to avoid breaking other things.
+            this.feePayer = this.feePayerWallet.getBaseAddress(0).getAddress();
+
+            return this;
+        }
+
+        /**
+         * Sets the provided collateral payer address. This method ensures that the collateral payer can only be set once.
+         *
+         * @param address the address of the collateral payer to be set
+         * @return TxContext
+         * @throws TxBuildException if the collateral payer has already been set
+         */
         public TxContext collateralPayer(String address) {
+            if (collateralPayerWallet != null || collateralPayer != null)
+                throw new TxBuildException("The collateral payer has already been set. It can only be set once.");
+
             this.collateralPayer = address;
             return this;
         }
 
+        /**
+         * Sets the collateral payer using the provided wallet. This method ensures that the collateral payer
+         * is set only once.
+         *
+         * @param wallet the wallet from which the collateral payer address will be derived
+         * @return TxContext
+         * @throws TxBuildException if the collateral payer has already been set
+         */
+        public TxContext collateralPayer(Wallet wallet) {
+            if (collateralPayerWallet != null || collateralPayer != null)
+                throw new TxBuildException("The collateral payer has already been set. It can only be set once.");
+
+            this.collateralPayerWallet = wallet;
+            // TODO collateralPayer is not used in this scenarios, but it must be set to avoid breaking other things.
+            this.collateralPayer = this.collateralPayerWallet.getBaseAddress(0).getAddress();
+
+            return this;
+        }
+
         /**
          * Set a TxBuilder function to transform the transaction before balance calculation.
          * This is useful when additional transformation logic is required before balance calculation.
@@ -281,7 +337,10 @@ private Tuple<TxBuilderContext, TxBuilder> _build() {
                     ((ScriptTx) tx).withChangeAddress(feePayer);
                 }
                 if (tx.getFromAddress() == null && tx instanceof ScriptTx) {
-                    ((ScriptTx) tx).from(feePayer);
+                    if (feePayerWallet != null)
+                        ((ScriptTx) tx).from(feePayerWallet);
+                    else
+                        ((ScriptTx) tx).from(feePayer);
                 }
 
                 txBuilder = txBuilder.andThen(tx.complete());
@@ -338,7 +397,7 @@ private Tuple<TxBuilderContext, TxBuilder> _build() {
             if (preBalanceTrasformer != null)
                 txBuilder = txBuilder.andThen(preBalanceTrasformer);
 
-            if (feePayer == null) {
+            if (feePayer == null && feePayerWallet == null) {
                 if (txList.length == 1) {
                     feePayer = txList[0].getFeePayer();
                     if (feePayer == null)
@@ -352,9 +411,18 @@ private Tuple<TxBuilderContext, TxBuilder> _build() {
             txBuilder = buildValidityIntervalTxBuilder(txBuilder);
 
             if (containsScriptTx) {
-                if (collateralPayer == null)
-                    collateralPayer = feePayer;
-                txBuilder = txBuilder.andThen(buildCollateralOutput(collateralPayer));
+                if (collateralPayer == null && collateralPayerWallet == null) {
+                    if (feePayerWallet != null)
+                        collateralPayerWallet = feePayerWallet;
+                    else
+                        collateralPayer = feePayer;
+                }
+
+                if (collateralPayerWallet != null) {
+                    txBuilder = txBuilder.andThen(buildCollateralOutput(collateralPayerWallet));
+                } else {
+                    txBuilder = txBuilder.andThen(buildCollateralOutput(collateralPayer));
+                }
             }
 
             if (containsScriptTx) {
@@ -393,7 +461,11 @@ private Tuple<TxBuilderContext, TxBuilder> _build() {
             }
 
             //Balance outputs
-            txBuilder = txBuilder.andThen(ScriptBalanceTxProviders.balanceTx(feePayer, totalSigners, containsScriptTx));
+            if (feePayerWallet != null) {
+                var walletAddrIterator = new HDWalletAddressIterator(feePayerWallet, utxoSupplier);
+                txBuilder = txBuilder.andThen(ScriptBalanceTxProviders.balanceTx(walletAddrIterator, totalSigners, containsScriptTx));
+            } else
+                txBuilder = txBuilder.andThen(ScriptBalanceTxProviders.balanceTx(feePayer, totalSigners, containsScriptTx));
 
             if ((containsScriptTx || hasMultiAssetMint) && removeDuplicateScriptWitnesses) {
                 txBuilder = txBuilder.andThen(DuplicateScriptWitnessChecker.removeDuplicateScriptWitnesses());
@@ -420,22 +492,45 @@ private int getTotalSigners() {
             return totalSigners;
         }
 
-        private TxBuilder buildCollateralOutput(String feePayer) {
+        private TxBuilder buildCollateralOutput(String payingAddress) {
+            if (collateralInputs != null && !collateralInputs.isEmpty()) {
+                List<Utxo> collateralUtxos = collateralInputs.stream()
+                        .map(input -> utxoSupplier.getTxOutput(input.getTransactionId(), input.getIndex()))
+                        .map(optionalUtxo -> optionalUtxo.get())
+                        .collect(Collectors.toList());
+                return CollateralBuilders.collateralOutputs(payingAddress, List.copyOf(collateralUtxos));
+            } else {
+                UtxoSelectionStrategy utxoSelectionStrategy = new DefaultUtxoSelectionStrategyImpl(utxoSupplier);
+                Set<Utxo> collateralUtxos = utxoSelectionStrategy.select(payingAddress, DEFAULT_COLLATERAL_AMT, null);
+                if (collateralUtxos.size() > MAX_COLLATERAL_INPUTS) {
+                    utxoSelectionStrategy = new LargestFirstUtxoSelectionStrategy(utxoSupplier);
+                    collateralUtxos = utxoSelectionStrategy.select(payingAddress, DEFAULT_COLLATERAL_AMT, null);
+                }
+
+                return CollateralBuilders.collateralOutputs(payingAddress, List.copyOf(collateralUtxos));
+            }
+        }
+
+        private TxBuilder buildCollateralOutput(Wallet payingWallet) {
+            String collateralPayerAddress = payingWallet.getBaseAddressString(0); //TODO: first addr as collateral output addr
+
             if (collateralInputs != null && !collateralInputs.isEmpty()) {
                 List<Utxo> collateralUtxos = collateralInputs.stream()
                         .map(input -> utxoSupplier.getTxOutput(input.getTransactionId(), input.getIndex()))
                         .map(optionalUtxo -> optionalUtxo.get())
                         .collect(Collectors.toList());
-                return CollateralBuilders.collateralOutputs(feePayer, List.copyOf(collateralUtxos));
+                return CollateralBuilders.collateralOutputs(collateralPayerAddress, List.copyOf(collateralUtxos));
             } else {
                 UtxoSelectionStrategy utxoSelectionStrategy = new DefaultUtxoSelectionStrategyImpl(utxoSupplier);
-                Set<Utxo> collateralUtxos = utxoSelectionStrategy.select(feePayer, DEFAULT_COLLATERAL_AMT, null);
+                var hdWalletAddressIterator = new HDWalletAddressIterator(payingWallet, utxoSupplier);
+
+                List<Utxo> collateralUtxos = utxoSelectionStrategy.selectUtxos(hdWalletAddressIterator, List.of(DEFAULT_COLLATERAL_AMT), null);
                 if (collateralUtxos.size() > MAX_COLLATERAL_INPUTS) {
                     utxoSelectionStrategy = new LargestFirstUtxoSelectionStrategy(utxoSupplier);
-                    collateralUtxos = utxoSelectionStrategy.select(feePayer, DEFAULT_COLLATERAL_AMT, null);
+                    collateralUtxos = utxoSelectionStrategy.selectUtxos(hdWalletAddressIterator, List.of(DEFAULT_COLLATERAL_AMT), null);
                 }
 
-                return CollateralBuilders.collateralOutputs(feePayer, List.copyOf(collateralUtxos));
+                return CollateralBuilders.collateralOutputs(collateralPayerAddress, List.copyOf(collateralUtxos));
             }
         }
 
diff --git a/quicktx/src/main/java/com/bloxbean/cardano/client/quicktx/ScriptTx.java b/quicktx/src/main/java/com/bloxbean/cardano/client/quicktx/ScriptTx.java
@@ -46,6 +46,7 @@ public class ScriptTx extends AbstractTx<ScriptTx> {
     protected List<TransactionInput> referenceInputs;
 
     protected String fromAddress;
+    protected Wallet fromWallet;
     private StakeTx stakeTx;
     private GovTx govTx;
 
@@ -641,13 +642,19 @@ protected String getFromAddress() {
 
     @Override
     protected Wallet getFromWallet() {
-        return null;
+        return fromWallet;
     }
 
     void from(String address) {
         this.fromAddress = address;
     }
 
+    void from(Wallet wallet) {
+        this.fromWallet = wallet;
+        // TODO fromAddress is not used in this scenarios, but it must be set to avoid breaking other things.
+        this.fromAddress = this.fromWallet.getBaseAddressString(0);
+    }
+
     @Override
     protected void postBalanceTx(Transaction transaction) {
         if (spendingContexts != null && !spendingContexts.isEmpty()) {
