Skip to content

Commit fef1669

Browse files
committed
feat: improve UTxO selection
- count coins from UTxOs that were already selected for different coin - optimize and defragment UTxOs
1 parent 282856d commit fef1669

File tree

1 file changed

+173
-29
lines changed

1 file changed

+173
-29
lines changed

cardano_clusterlib/txtools.py

Lines changed: 173 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,21 @@ def _organize_utxos_by_id(
4141
return db
4242

4343

44+
def _organize_utxos_by_coin_and_id(
45+
tx_list: list[structs.UTXOData],
46+
) -> dict[str, dict[str, int]]:
47+
"""Organize UTxOs by coin and ID (hash#ix)."""
48+
db: dict[str, dict[str, int]] = {}
49+
for r in tx_list:
50+
utxo_id = f"{r.utxo_hash}#{r.utxo_ix}"
51+
db_rec = db.get(r.coin)
52+
if db_rec is None:
53+
db[r.coin] = {utxo_id: r.amount}
54+
continue
55+
db_rec[utxo_id] = r.amount
56+
return db
57+
58+
4459
def _get_usable_utxos(
4560
address_utxos: list[structs.UTXOData], coins: set[str]
4661
) -> list[structs.UTXOData]:
@@ -67,31 +82,151 @@ def _get_usable_utxos(
6782
return txins
6883

6984

70-
def _collect_utxos_amount(
71-
utxos: list[structs.UTXOData], amount: int, min_change_value: int
72-
) -> list[structs.UTXOData]:
73-
"""Collect UTxOs so their total combined amount >= `amount`."""
74-
collected_utxos: list[structs.UTXOData] = []
75-
collected_amount = 0
76-
# `_min_change_value` applies only to ADA
77-
amount_plus_change = (
78-
amount + min_change_value if utxos and utxos[0].coin == consts.DEFAULT_COIN else amount
79-
)
80-
for utxo in utxos:
85+
def _pick_coins_from_already_selected_utxos(
86+
coin_txins: dict[str, int],
87+
already_selected_utxos: set[str],
88+
target_amount: int,
89+
target_with_change: int,
90+
) -> tuple[set[str], int, bool]:
91+
"""Pick coins from already selected UTxOs if they have the desired coin.
92+
93+
Args:
94+
coin_txins (dict): A dictionary of coin UTxOs.
95+
already_selected_utxos (set): A set of already selected UTxOs (for different coins).
96+
target_amount (int): The desired amount.
97+
target_with_change (int): The desired amount with minimal change.
98+
99+
Returns:
100+
tuple: A tuple with selected UTxO IDs, accumulated amount and a bool indicating if the
101+
desired amount was met.
102+
"""
103+
picked_utxos: set[str] = set()
104+
accumulated_amount = 0
105+
106+
# See if the coin exists in UTxOs that were already selected
107+
for utxo_id in already_selected_utxos:
108+
utxo_amount = coin_txins.get(utxo_id)
109+
if utxo_amount is None:
110+
continue
111+
accumulated_amount += utxo_amount
112+
81113
# If we were able to collect exact amount, no change is needed
82-
if collected_amount == amount:
114+
if accumulated_amount == target_amount:
83115
break
84116
# Make sure the change is higher than `_min_change_value`
85-
if collected_amount >= amount_plus_change:
117+
if accumulated_amount >= target_with_change:
86118
break
87-
collected_utxos.append(utxo)
88-
collected_amount += utxo.amount
119+
else:
120+
return picked_utxos, accumulated_amount, False
121+
122+
return picked_utxos, accumulated_amount, True
123+
124+
125+
def _pick_utxos_with_defragmentation(
126+
utxos: list[tuple[str, int]],
127+
target_amount: int,
128+
target_with_change: int,
129+
accumulated_amount: int,
130+
) -> tuple[set[str], int, bool]:
131+
"""Pick UTxOs to meet or exceed the target amount while prioritizing defragmentation.
132+
133+
Args:
134+
utxos (list of tuple): A list of tuples (utxo_id, coin_amount).
135+
target_amount (int): The desired amount.
136+
target_with_change (int): The desired amount with minimal change.
137+
accumulated_amount (int): The accumulated amount.
138+
139+
Returns:
140+
tuple: A tuple with selected UTxO IDs, accumulated amount and a bool indicating if the
141+
desired amount was met.
142+
"""
143+
# Sort UTxOs by amount in ascending order
144+
sorted_utxos = sorted(enumerate(utxos), key=lambda x: x[1][1]) # Keep original indices
145+
selected_indices = set()
146+
picked_utxos = set()
147+
148+
# Step 1: Select up to 10 smallest UTxOs
149+
for i, (utxo_id, coin_amount) in sorted_utxos[:10]:
150+
picked_utxos.add(utxo_id)
151+
selected_indices.add(i)
152+
accumulated_amount += coin_amount
153+
154+
# If we were able to collect exact amount, no change is needed
155+
if accumulated_amount == target_amount:
156+
return picked_utxos, accumulated_amount, True
157+
# Make sure the change is higher than `_min_change_value`
158+
if accumulated_amount >= target_with_change:
159+
return picked_utxos, accumulated_amount, True
160+
161+
# Step 2: If target is not met, select UTxO closest to remaining amount
162+
while accumulated_amount < target_with_change:
163+
# If we were able to collect exact amount, no change is needed
164+
if accumulated_amount == target_amount:
165+
return picked_utxos, accumulated_amount, True
166+
167+
# We target exact amount, but if we are already over it, we need at least additional
168+
# `_min_change_value` for change.
169+
if accumulated_amount > target_amount:
170+
remaining_amount = target_with_change - accumulated_amount
171+
else:
172+
remaining_amount = target_amount - accumulated_amount
173+
174+
# Find the index of the UTxO closest to the remaining amount
175+
closest_index = min(
176+
(i for i, _ in sorted_utxos if i not in selected_indices),
177+
key=lambda i: abs(utxos[i][1] - remaining_amount),
178+
default=None,
179+
)
180+
181+
# If all UTxOs have been considered, the target was not met
182+
if closest_index is None:
183+
return picked_utxos, accumulated_amount, False
184+
185+
# Select the closest UTxO
186+
utxo_id, coin_amount = utxos[closest_index]
187+
picked_utxos.add(utxo_id)
188+
selected_indices.add(closest_index)
189+
accumulated_amount += coin_amount
190+
191+
return picked_utxos, accumulated_amount, True
192+
193+
194+
def _select_utxos_per_coin(
195+
coin_txins: dict[str, int],
196+
coin: str,
197+
target_amount: int,
198+
target_with_change: int,
199+
already_selected_utxos: set[str],
200+
) -> set[str]:
201+
"""Select UTxOs for a given coin so their total combined amount >= `amount`."""
202+
selected_utxos, accumulated_amount, target_met = _pick_coins_from_already_selected_utxos(
203+
coin_txins=coin_txins,
204+
already_selected_utxos=already_selected_utxos,
205+
target_amount=target_amount,
206+
target_with_change=target_with_change,
207+
)
89208

90-
return collected_utxos
209+
# Pick more UTxOs if the amount is not satisfied yet
210+
if not target_met:
211+
ids_and_amounts = [(i, a) for i, a in coin_txins.items() if i not in already_selected_utxos]
212+
more_utxos, _, target_met = _pick_utxos_with_defragmentation(
213+
utxos=ids_and_amounts,
214+
target_amount=target_amount,
215+
target_with_change=target_with_change,
216+
accumulated_amount=accumulated_amount,
217+
)
218+
selected_utxos.update(more_utxos)
219+
220+
if not target_met:
221+
LOGGER.warning(
222+
f"Could not meet target amount {target_amount} for coin '{coin}' with the given UTxOs."
223+
)
224+
225+
return selected_utxos
91226

92227

93228
def _select_utxos(
94-
txins_db: dict[str, list[structs.UTXOData]],
229+
txins_by_coin_and_id: dict[str, dict[str, int]],
95230
txouts_passed_db: dict[str, list[structs.TxOut]],
96231
txouts_mint_db: dict[str, list[structs.TxOut]],
97232
fee: int,
@@ -107,8 +242,8 @@ def _select_utxos(
107242
utxo_ids: set[str] = set()
108243

109244
# Iterate over coins both in txins and txouts
110-
for coin in set(txins_db).union(txouts_passed_db).union(txouts_mint_db):
111-
coin_txins = txins_db.get(coin) or []
245+
for coin in set(txins_by_coin_and_id).union(txouts_passed_db).union(txouts_mint_db):
246+
coin_txins = txins_by_coin_and_id.get(coin) or {}
112247
coin_txouts = txouts_passed_db.get(coin) or []
113248

114249
total_output_amount = functools.reduce(lambda x, y: x + y.amount, coin_txouts, 0)
@@ -117,26 +252,35 @@ def _select_utxos(
117252
# The value "-1" means all available funds
118253
max_index = [idx for idx, val in enumerate(coin_txouts) if val.amount == -1]
119254
if max_index:
120-
utxo_ids.update(f"{rec.utxo_hash}#{rec.utxo_ix}" for rec in coin_txins)
255+
utxo_ids.update(r for r in coin_txins)
121256
continue
122257

123258
tx_fee = max(1, fee)
124259
funds_needed = total_output_amount + tx_fee + deposit + treasury_donation
125260
total_withdrawals_amount = functools.reduce(lambda x, y: x + y.amount, withdrawals, 0)
126261
# Fee needs an input, even if withdrawal would cover all needed funds
127262
input_funds_needed = max(funds_needed - total_withdrawals_amount, tx_fee)
263+
# `_min_change_value` applies only to ADA
264+
target_with_change = input_funds_needed + min_change_value
128265
else:
129266
coin_txouts_minted = txouts_mint_db.get(coin) or []
130267
total_minted_amount = functools.reduce(lambda x, y: x + y.amount, coin_txouts_minted, 0)
131268
# In case of token burning, `total_minted_amount` might be negative.
132269
# Try to collect enough funds to satisfy both token burning and token
133270
# transfers, even though there might be an overlap.
134271
input_funds_needed = total_output_amount - total_minted_amount
135-
136-
filtered_coin_utxos = _collect_utxos_amount(
137-
utxos=coin_txins, amount=input_funds_needed, min_change_value=min_change_value
138-
)
139-
utxo_ids.update(f"{rec.utxo_hash}#{rec.utxo_ix}" for rec in filtered_coin_utxos)
272+
target_with_change = input_funds_needed
273+
274+
if input_funds_needed:
275+
utxo_ids.update(
276+
_select_utxos_per_coin(
277+
coin_txins=txins_by_coin_and_id.get(coin) or {},
278+
coin=coin,
279+
target_amount=input_funds_needed,
280+
target_with_change=target_with_change,
281+
already_selected_utxos=utxo_ids,
282+
)
283+
)
140284

141285
return utxo_ids
142286

@@ -536,11 +680,11 @@ def _get_tx_ins_outs(
536680
msg = "No input UTxO."
537681
raise exceptions.CLIError(msg)
538682

539-
txins_db_all: dict[str, list[structs.UTXOData]] = _organize_tx_ins_outs_by_coin(txins_all)
683+
txins_by_coin_and_id = _organize_utxos_by_coin_and_id(txins_all)
540684

541685
# All output coins, except those minted by this transaction, need to be present in
542686
# transaction inputs
543-
if not set(outcoins_passed).difference(txouts_mint_db).issubset(txins_db_all):
687+
if not set(outcoins_passed).difference(txouts_mint_db).issubset(txins_by_coin_and_id):
544688
msg = "Not all output coins are present in input UTxOs."
545689
raise exceptions.CLIError(msg)
546690

@@ -555,11 +699,11 @@ def _get_tx_ins_outs(
555699
if txins:
556700
# Don't touch txins that were passed to the function
557701
txins_filtered = txins_all
558-
txins_db_filtered = txins_db_all
702+
txins_db_filtered = _organize_tx_ins_outs_by_coin(txins_all)
559703
else:
560704
# Select only UTxOs that are needed to satisfy all outputs, deposits and fee
561705
selected_utxo_ids = _select_utxos(
562-
txins_db=txins_db_all,
706+
txins_by_coin_and_id=txins_by_coin_and_id,
563707
txouts_passed_db=txouts_passed_db,
564708
txouts_mint_db=txouts_mint_db,
565709
fee=fee,

0 commit comments

Comments
 (0)