[LABS-304] Refactor get_spendable_coins logic mostly out of RPC (#20302)

Refactor `get_spendable_coins` logic mostly out of RPC

chia/wallet/util/tx_config.py

@@ -32,6 +32,15 @@ class CoinSelectionConfig:
     def override(self, **kwargs: Any) -> CoinSelectionConfig:
         return dataclasses.replace(self, **kwargs)
 
+    def filter_coins(self, coins: set[Coin]) -> set[Coin]:
+        return {
+            coin
+            for coin in coins
+            if self.min_coin_amount <= coin.amount <= self.max_coin_amount
+            and coin.amount not in self.excluded_coin_amounts
+            and coin.name() not in self.excluded_coin_ids
+        }
+
 
 @dataclasses.dataclass(frozen=True)
 class TXConfig(CoinSelectionConfig):

chia/wallet/wallet_rpc_api.py

@@ -1674,50 +1674,42 @@ class WalletRpcApi:
 
         state_mgr = self.service.wallet_state_manager
         async with state_mgr.lock:
-            all_coin_records = await state_mgr.coin_store.get_unspent_coins_for_wallet(request.wallet_id)
-            spendable_coins = list(await state_mgr.get_spendable_coins_for_wallet(request.wallet_id, all_coin_records))
-
-            # Now we get the unconfirmed transactions and manually derive the additions and removals.
-            unconfirmed_transactions: list[TransactionRecord] = await state_mgr.tx_store.get_unconfirmed_for_wallet(
-                request.wallet_id
+            # Removals
+            unconfirmed_removals = await state_mgr.unconfirmed_additions_or_removals_for_wallet(
+                wallet_id=request.wallet_id, get="removals"
             )
-            unconfirmed_removal_ids: dict[bytes32, uint64] = {
-                coin.name(): transaction.created_at_time
-                for transaction in unconfirmed_transactions
-                for coin in transaction.removals
-            }
-            unconfirmed_additions: list[Coin] = [
-                coin
-                for transaction in unconfirmed_transactions
-                for coin in transaction.additions
-                if await state_mgr.does_coin_belong_to_wallet(coin, request.wallet_id)
-            ]
-            valid_spendable_cr: list[CoinRecord] = []
-            unconfirmed_removals: list[CoinRecord] = []
-            for coin_record in all_coin_records:
-                if coin_record.name() in unconfirmed_removal_ids:
-                    unconfirmed_removals.append(coin_record.to_coin_record(unconfirmed_removal_ids[coin_record.name()]))
+            unconfirmed_removal_ids = {coin.name() for coin in unconfirmed_removals}
+            removal_records: list[CoinRecord] = []
+            for coin_record in (
+                await state_mgr.coin_store.get_coin_records(
+                    coin_id_filter=HashFilter.include(list(unconfirmed_removal_ids))
+                )
+            ).records:
+                removal_records.append(await state_mgr.get_coin_record_by_wallet_record(coin_record))
 
-            cs_config = request.autofill(constants=self.service.wallet_state_manager.constants)
-            for coin_record in spendable_coins:  # remove all the unconfirmed coins, exclude coins and dust.
-                if coin_record.name() in unconfirmed_removal_ids:
-                    continue
-                if coin_record.coin.name() in cs_config.excluded_coin_ids:
-                    continue
-                if (coin_record.coin.amount < cs_config.min_coin_amount) or (
-                    coin_record.coin.amount > cs_config.max_coin_amount
-                ):
-                    continue
-                if coin_record.coin.amount in cs_config.excluded_coin_amounts:
-                    continue
-                c_r = await state_mgr.get_coin_record_by_wallet_record(coin_record)
-                assert c_r is not None and c_r.coin == coin_record.coin  # this should never happen
-                valid_spendable_cr.append(c_r)
+            # Additions
+            unconfirmed_additions = await state_mgr.unconfirmed_additions_or_removals_for_wallet(
+                wallet_id=request.wallet_id, get="additions"
+            )
+
+            # Spendable coins
+            unfiltered_spendable_coin_records = await state_mgr.get_spendable_coins_for_wallet(
+                request.wallet_id, pending_removals=unconfirmed_removal_ids
+            )
+            filtered_spendable_coins = request.autofill(
+                constants=self.service.wallet_state_manager.constants
+            ).filter_coins({cr.coin for cr in unfiltered_spendable_coin_records})
+            filtered_spendable_coin_records = list(
+                cr for cr in unfiltered_spendable_coin_records if cr.coin in filtered_spendable_coins
+            )
+            valid_spendable_cr: list[CoinRecord] = []
+            for coin_record in filtered_spendable_coin_records:
+                valid_spendable_cr.append(await state_mgr.get_coin_record_by_wallet_record(coin_record))
 
         return GetSpendableCoinsResponse(
             confirmed_records=valid_spendable_cr,
-            unconfirmed_removals=unconfirmed_removals,
-            unconfirmed_additions=unconfirmed_additions,
+            unconfirmed_removals=removal_records,
+            unconfirmed_additions=list(unconfirmed_additions),
         )
 
     @marshal

chia/wallet/wallet_state_manager.py

@@ -10,7 +10,7 @@ import traceback
 from collections.abc import AsyncIterator, Callable
 from contextlib import asynccontextmanager
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, TypeVar, cast
+from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast
 
 import aiosqlite
 from chia_rs import AugSchemeMPL, CoinRecord, CoinSpend, CoinState, ConsensusConstants, G1Element, G2Element, PrivateKey
@@ -2498,8 +2498,26 @@ class WalletStateManager:
         await result.commit(self)
         self.state_changed("wallet_created")
 
+    async def unconfirmed_additions_or_removals_for_wallet(
+        self, *, wallet_id: uint32, get: Literal["additions", "removals"]
+    ) -> set[Coin]:
+        unconfirmed_tx: list[TransactionRecord] = await self.tx_store.get_unconfirmed_for_wallet(wallet_id)
+        return_set: set[Coin] = set()
+        for tx in unconfirmed_tx:
+            hint_dict = tx.hint_dict()
+            checked_set = tx.removals if get == "removals" else tx.additions
+            for coin in checked_set:
+                if await self.does_coin_belong_to_wallet(coin, wallet_id, hint_dict):
+                    return_set.add(coin)
+
+        return return_set
+
     async def get_spendable_coins_for_wallet(
-        self, wallet_id: int, records: set[WalletCoinRecord] | None = None, in_one_block: bool = False
+        self,
+        wallet_id: int,
+        records: set[WalletCoinRecord] | None = None,
+        pending_removals: set[bytes32] | None = None,
+        in_one_block: bool = False,
     ) -> set[WalletCoinRecord]:
         wallet = self.wallets[uint32(wallet_id)]
         wallet_type = wallet.type()
@@ -2510,22 +2528,20 @@ class WalletStateManager:
                 records = await self.coin_store.get_unspent_coins_for_wallet(wallet_id)
 
         # Coins that are currently part of a transaction
-        unconfirmed_tx: list[TransactionRecord] = await self.tx_store.get_unconfirmed_for_wallet(wallet_id)
-        removal_dict: dict[bytes32, Coin] = {}
-        for tx in unconfirmed_tx:
-            for coin in tx.removals:
-                # TODO, "if" might not be necessary once unconfirmed tx doesn't contain coins for other wallets
-                if await self.does_coin_belong_to_wallet(coin, wallet_id, tx.hint_dict()):
-                    removal_dict[coin.name()] = coin
+        if pending_removals is None:
+            pending_removals = {
+                coin.name()
+                for coin in await self.unconfirmed_additions_or_removals_for_wallet(
+                    wallet_id=uint32(wallet_id), get="removals"
+                )
+            }
 
         # Coins that are part of the trade
         offer_locked_coins: dict[bytes32, WalletCoinRecord] = await self.trade_manager.get_locked_coins()
 
         filtered = set()
         for record in records:
-            if record.coin.name() in offer_locked_coins:
-                continue
-            if record.coin.name() in removal_dict:
+            if record.coin.name() in {*offer_locked_coins.keys(), *pending_removals}:
                 continue
             if hasattr(wallet, "is_coin_spendable") and not await wallet.is_coin_spendable(record):
                 continue