refactor: consolidate ProTx wallet UTXO locking logic

src/interfaces/chain.h

@@ -143,7 +143,7 @@ class Chain
     virtual bool hasChainLock(int height, const uint256& hash) = 0;
 
     //! Return list of MN Collateral from outputs
-    virtual std::vector<COutPoint> listMNCollaterials(const std::vector<std::pair<const CTransactionRef&, unsigned int>>& outputs) = 0;
+    virtual std::vector<COutPoint> listMNCollaterials(const std::vector<std::pair<const CTransactionRef&, uint32_t>>& outputs) = 0;
     //! Return whether node has the block and optionally return block metadata
     //! or contents.
     virtual bool findBlock(const uint256& hash, const FoundBlock& block={}) = 0;

src/interfaces/wallet.h

@@ -145,10 +145,10 @@ class Wallet
     virtual bool isLockedCoin(const COutPoint& output) = 0;
 
     //! List locked coins.
-    virtual void listLockedCoins(std::vector<COutPoint>& outputs) = 0;
+    virtual std::vector<COutPoint> listLockedCoins() = 0;
 
     //! List protx coins.
-    virtual void listProTxCoins(std::vector<COutPoint>& vOutpts) = 0;
+    virtual std::vector<COutPoint> listProTxCoins() = 0;
 
     //! Create transaction.
     virtual CTransactionRef createTransaction(const std::vector<CRecipient>& recipients,

src/node/interfaces.cpp

@@ -786,7 +786,7 @@ class ChainImpl : public Chain
         if (m_node.llmq_ctx == nullptr || m_node.llmq_ctx->clhandler == nullptr) return false;
         return m_node.llmq_ctx->clhandler->HasChainLock(height, hash);
     }
-    std::vector<COutPoint> listMNCollaterials(const std::vector<std::pair<const CTransactionRef&, unsigned int>>& outputs) override
+    std::vector<COutPoint> listMNCollaterials(const std::vector<std::pair<const CTransactionRef&, uint32_t>>& outputs) override
     {
         const CBlockIndex *tip = WITH_LOCK(::cs_main, return chainman().ActiveChain().Tip());
         CDeterministicMNList mnList{};

src/qt/coincontroldialog.cpp

@@ -426,11 +426,10 @@ void CoinControlDialog::viewItemChanged(QTreeWidgetItem* item, int column)
 // shows count of locked unspent outputs
 void CoinControlDialog::updateLabelLocked()
 {
-    std::vector<COutPoint> vOutpts;
-    model->wallet().listLockedCoins(vOutpts);
-    if (vOutpts.size() > 0)
+    auto locked_coins{model->wallet().listLockedCoins().size()};
+    if (locked_coins > 0)
     {
-       ui->labelLocked->setText(tr("(%1 locked)").arg(vOutpts.size()));
+       ui->labelLocked->setText(tr("(%1 locked)").arg(locked_coins));
        ui->labelLocked->setVisible(true);
     }
     else ui->labelLocked->setVisible(false);

src/qt/masternodelist.cpp

@@ -204,9 +204,7 @@ void MasternodeList::updateDIP3List()
 
     std::set<COutPoint> setOutpts;
     if (walletModel && ui->checkBoxMyMasternodesOnly->isChecked()) {
-        std::vector<COutPoint> vOutpts;
-        walletModel->wallet().listProTxCoins(vOutpts);
-        for (const auto& outpt : vOutpts) {
+        for (const auto& outpt : walletModel->wallet().listProTxCoins()) {
             setOutpts.emplace(outpt);
         }
     }

src/rpc/evo.cpp

@@ -1390,10 +1390,8 @@ static RPCHelpMan protx_list()
             throw JSONRPCError(RPC_INVALID_PARAMETER, "invalid height specified");
         }
 
-        std::vector<COutPoint> vOutpts;
-        wallet->ListProTxCoins(vOutpts);
         std::set<COutPoint> setOutpts;
-        for (const auto& outpt : vOutpts) {
+        for (const auto& outpt : wallet->ListProTxCoins()) {
             setOutpts.emplace(outpt);
         }
 

src/wallet/interfaces.cpp

@@ -255,15 +255,15 @@ class WalletImpl : public Wallet
         LOCK(m_wallet->cs_wallet);
         return m_wallet->IsLockedCoin(output.hash, output.n);
     }
-    void listLockedCoins(std::vector<COutPoint>& outputs) override
+    std::vector<COutPoint> listLockedCoins() override
     {
         LOCK(m_wallet->cs_wallet);
-        return m_wallet->ListLockedCoins(outputs);
+        return m_wallet->ListLockedCoins();
     }
-    void listProTxCoins(std::vector<COutPoint>& outputs) override
+    std::vector<COutPoint> listProTxCoins() override
     {
         LOCK(m_wallet->cs_wallet);
-        return m_wallet->ListProTxCoins(outputs);
+        return m_wallet->ListProTxCoins();
     }
     CTransactionRef createTransaction(const std::vector<CRecipient>& recipients,
         const CCoinControl& coin_control,

src/wallet/rpcwallet.cpp

@@ -2371,14 +2371,9 @@ static RPCHelpMan listlockunspent()
 
     LOCK(pwallet->cs_wallet);
 
-    std::vector<COutPoint> vOutpts;
-    pwallet->ListLockedCoins(vOutpts);
-
     UniValue ret(UniValue::VARR);
-
-    for (const COutPoint& outpt : vOutpts) {
+    for (const COutPoint& outpt : pwallet->ListLockedCoins()) {
         UniValue o(UniValue::VOBJ);
-
         o.pushKV("txid", outpt.hash.GetHex());
         o.pushKV("vout", (int)outpt.n);
         ret.push_back(o);

src/wallet/wallet.cpp

@@ -914,27 +914,15 @@ CWalletTx* CWallet::AddToWallet(CTransactionRef tx, const CWalletTx::Confirmatio
     CWalletTx& wtx = (*ret.first).second;
     bool fInsertedNew = ret.second;
     bool fUpdated = update_wtx && update_wtx(wtx, fInsertedNew);
+    std::set<COutPoint> candidates;
     if (fInsertedNew) {
         wtx.m_confirm = confirm;
         wtx.nTimeReceived = GetTime();
         wtx.nOrderPos = IncOrderPosNext(&batch);
         wtx.m_it_wtxOrdered = wtxOrdered.insert(std::make_pair(wtx.nOrderPos, &wtx));
         wtx.nTimeSmart = ComputeTimeSmart(wtx);
         AddToSpends(hash, &batch);
-
-        std::vector<std::pair<const CTransactionRef&, unsigned int>> outputs;
-        for(unsigned int i = 0; i < wtx.tx->vout.size(); ++i) {
-            if (IsMine(wtx.tx->vout[i]) && !IsSpent(hash, i)) {
-                setWalletUTXO.insert(COutPoint(hash, i));
-                outputs.emplace_back(wtx.tx, i);
-            }
-        }
-        // TODO: refactor duplicated code between CWallet::AddToWallet and CWallet::AutoLockMasternodeCollaterals
-        if (m_chain) {
-            for (const auto& outPoint : m_chain->listMNCollaterials(outputs)) {
-                LockCoin(outPoint, &batch);
-            }
-        }
+        candidates = AddWalletUTXOs(wtx.tx, /*ret_dups=*/true);
     }
 
     if (!fInsertedNew)
@@ -950,25 +938,12 @@ CWalletTx* CWallet::AddToWallet(CTransactionRef tx, const CWalletTx::Confirmatio
             assert(wtx.m_confirm.hashBlock == confirm.hashBlock);
             assert(wtx.m_confirm.block_height == confirm.block_height);
         }
-
-        std::vector<std::pair<const CTransactionRef&, unsigned int>> outputs;
-        for(unsigned int i = 0; i < wtx.tx->vout.size(); ++i) {
-            if (IsMine(wtx.tx->vout[i]) && !IsSpent(hash, i)) {
-                bool new_utxo = setWalletUTXO.insert(COutPoint(hash, i)).second;
-                if (new_utxo) {
-                    outputs.emplace_back(wtx.tx, i);
-                    fUpdated = true;
-                }
-            }
-        }
-        // TODO: refactor duplicated code with case fInstertedNew
-        if (m_chain) {
-            for (const auto& outPoint : m_chain->listMNCollaterials(outputs)) {
-                LockCoin(outPoint);
-            }
-        }
+        candidates = AddWalletUTXOs(wtx.tx, /*ret_dups=*/false);
+        if (!candidates.empty()) fUpdated = true;
     }
 
+    LockProTxCoins(candidates, &batch);
+
     //// debug print
     WalletLogPrintf("AddToWallet %s  %s%s\n", hash.ToString(), (fInsertedNew ? "new" : ""), (fUpdated ? "update" : ""));
 
@@ -1062,6 +1037,21 @@ bool CWallet::LoadToWallet(const uint256& hash, const UpdateWalletTxFn& fill_wtx
     return true;
 }
 
+std::set<COutPoint> CWallet::AddWalletUTXOs(CTransactionRef tx, bool ret_dups)
+{
+    AssertLockHeld(cs_wallet);
+    std::set<COutPoint> ret;
+    uint256 hash{tx->GetHash()};
+    for (size_t idx = 0; idx < tx->vout.size(); ++idx) {
+        if (IsMine(tx->vout[idx]) && !IsSpent(hash, idx)) {
+            if (auto [_, inserted] = setWalletUTXO.emplace(hash, idx); inserted || ret_dups) {
+                ret.emplace(hash, idx);
+            }
+        }
+    }
+    return ret;
+}
+
 bool CWallet::AddToWalletIfInvolvingMe(const CTransactionRef& ptx, CWalletTx::Confirmation confirm, WalletBatch& batch, bool fUpdate)
 {
     const CTransaction& tx = *ptx;
@@ -2799,12 +2789,10 @@ std::map<CTxDestination, std::vector<COutput>> CWallet::ListCoins() const
         }
     }
 
-    std::vector<COutPoint> lockedCoins;
-    ListLockedCoins(lockedCoins);
     // Include watch-only for LegacyScriptPubKeyMan wallets without private keys
     const bool include_watch_only = GetLegacyScriptPubKeyMan() && IsWalletFlagSet(WALLET_FLAG_DISABLE_PRIVATE_KEYS);
     const isminetype is_mine_filter = include_watch_only ? ISMINE_WATCH_ONLY : ISMINE_SPENDABLE;
-    for (const COutPoint& output : lockedCoins) {
+    for (const COutPoint& output : setLockedCoins) {
         auto it = mapWallet.find(output.hash);
         if (it != mapWallet.end()) {
             int depth = it->second.GetDepthInMainChain();
@@ -4056,21 +4044,14 @@ DBErrors CWallet::LoadWallet()
 // This avoids accidental spending of collaterals. They can still be unlocked manually if a spend is really intended.
 void CWallet::AutoLockMasternodeCollaterals()
 {
-    if (!m_chain) return;
-
-    std::vector<std::pair<const CTransactionRef&, unsigned int>> outputs;
-
+    std::set<COutPoint> candidates;
     LOCK(cs_wallet);
-    for (const auto& pair : mapWallet) {
-        for (unsigned int i = 0; i < pair.second.tx->vout.size(); ++i) {
-            if (IsMine(pair.second.tx->vout[i]) && !IsSpent(pair.first, i)) {
-                outputs.emplace_back(pair.second.tx, i);
-            }
-        }
-    }
-    for (const auto& outPoint : m_chain->listMNCollaterials(outputs)) {
-        LockCoin(outPoint);
+    for (const auto& [txid, wtx] : mapWallet) {
+        auto tx_utxos{AddWalletUTXOs(wtx.tx, /*ret_dups=*/true)};
+        candidates.insert(tx_utxos.begin(), tx_utxos.end());
     }
+    WalletBatch batch(GetDatabase());
+    LockProTxCoins(candidates, &batch);
 }
 
 DBErrors CWallet::ZapSelectTx(std::vector<uint256>& vHashIn, std::vector<uint256>& vHashOut)
@@ -4502,34 +4483,36 @@ bool CWallet::IsLockedCoin(uint256 hash, unsigned int n) const
     return (setLockedCoins.count(outpt) > 0);
 }
 
-void CWallet::ListLockedCoins(std::vector<COutPoint>& vOutpts) const
+std::vector<COutPoint> CWallet::ListLockedCoins() const
 {
     AssertLockHeld(cs_wallet);
-    for (std::set<COutPoint>::iterator it = setLockedCoins.begin();
-         it != setLockedCoins.end(); it++) {
-        COutPoint outpt = (*it);
-        vOutpts.push_back(outpt);
-    }
+    return std::vector<COutPoint>(setLockedCoins.begin(), setLockedCoins.end());
 }
 
-void CWallet::ListProTxCoins(std::vector<COutPoint>& vOutpts) const
+std::vector<COutPoint> CWallet::ListProTxCoins() const { return ListProTxCoins(setWalletUTXO); }
+
+std::vector<COutPoint> CWallet::ListProTxCoins(const std::set<COutPoint>& utxos) const
 {
-    // TODO: refactor duplicated code between CWallet::AutoLockMasternodeCollaterals and CWallet::ListProTxCoins
-    if (!m_chain) {
-        vOutpts.clear();
-        return;
+    AssertLockHeld(cs_wallet);
+
+    if (!m_chain) return std::vector<COutPoint>();
+
+    std::vector<std::pair<const CTransactionRef&, /*index=*/uint32_t>> candidates;
+    for (const auto& output : utxos) {
+        if (auto it = mapWallet.find(output.hash); it != mapWallet.end()) {
+            const auto& [hash, wtx] = *it;
+            candidates.emplace_back(wtx.tx, output.n);
+        }
     }
-    std::vector<std::pair<const CTransactionRef&, unsigned int>> outputs;
+    return m_chain->listMNCollaterials(candidates);
+}
 
+void CWallet::LockProTxCoins(const std::set<COutPoint>& utxos, WalletBatch* batch)
+{
     AssertLockHeld(cs_wallet);
-    for (const auto &o : setWalletUTXO) {
-        auto it = mapWallet.find(o.hash);
-        if (it != mapWallet.end()) {
-            const auto &ptx = it->second;
-            outputs.emplace_back(ptx.tx, o.n);
-        }
+    for (const auto& utxo : ListProTxCoins(utxos)) {
+        LockCoin(utxo, batch);
     }
-    vOutpts = m_chain->listMNCollaterials(outputs);
 }
 
 /** @} */ // end of Actions

src/wallet/wallet.h

@@ -763,6 +763,12 @@ class CWallet final : public WalletStorage, public interfaces::Chain::Notificati
     void AddToSpends(const uint256& wtxid, WalletBatch* batch = nullptr) EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
 
     std::set<COutPoint> setWalletUTXO;
+    /** Add new UTXOs to the wallet UTXO set
+     *
+     *  @param[in] tx         Transaction to scan eligible UTXOs from
+     *  @param[in] ret_dups   Allow UTXOs already in set to be included in return value
+     *  @returns              Set of all new UTXOs (eligible to be) added to set */
+    std::set<COutPoint> AddWalletUTXOs(CTransactionRef tx, bool ret_dups) EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
     mutable std::map<COutPoint, int> mapOutpointRoundsCache GUARDED_BY(cs_wallet);
 
     /**
@@ -1038,8 +1044,10 @@ class CWallet final : public WalletStorage, public interfaces::Chain::Notificati
     bool LockCoin(const COutPoint& output, WalletBatch* batch = nullptr) EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
     bool UnlockCoin(const COutPoint& output, WalletBatch* batch = nullptr) EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
     bool UnlockAllCoins() EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
-    void ListLockedCoins(std::vector<COutPoint>& vOutpts) const EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
-    void ListProTxCoins(std::vector<COutPoint>& vOutpts) const EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
+    std::vector<COutPoint> ListLockedCoins() const EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
+    std::vector<COutPoint> ListProTxCoins() const EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
+    std::vector<COutPoint> ListProTxCoins(const std::set<COutPoint>& utxos) const EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
+    void LockProTxCoins(const std::set<COutPoint>& utxos, WalletBatch* batch = nullptr) EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
 
     /*
      * Rescan abort properties