wallet: unify outputs grouping process

furszy · furszy · commit bd91ed1cb2cc · 2023-03-06T09:45:40.000-03:00
The 'GroupOutputs()' function performs the same
calculations for only-positive and mixed groups,
the only difference is that when we look for
only-positive groups, we discard negative utxos.

So, instead of wasting resources calling GroupOutputs()
for positive-only first, then call it again to include
the negative ones in the result, we can execute
GroupOutputs() only once, including in the response
both group types (positive-only and mixed).
diff --git a/src/wallet/coinselection.cpp b/src/wallet/coinselection.cpp
@@ -372,6 +372,28 @@ CAmount OutputGroup::GetSelectionAmount() const
     return m_subtract_fee_outputs ? m_value : effective_value;
 }
 
+void OutputGroupTypeMap::Push(const OutputGroup& group, OutputType type, bool insert_positive, bool insert_mixed)
+{
+    if (group.m_outputs.empty()) return;
+
+    Groups& groups = groups_by_type[type];
+    if (insert_positive && group.GetSelectionAmount() > 0) {
+        groups.positive_group.emplace_back(group);
+        all_groups.positive_group.emplace_back(group);
+    }
+    if (insert_mixed) {
+        groups.mixed_group.emplace_back(group);
+        all_groups.mixed_group.emplace_back(group);
+    }
+}
+
+std::optional<Groups> OutputGroupTypeMap::Find(OutputType type)
+{
+    auto it_by_type = groups_by_type.find(type);
+    if (it_by_type == groups_by_type.end()) return std::nullopt;
+    return it_by_type->second;
+}
+
 CAmount GetSelectionWaste(const std::set<std::shared_ptr<COutput>>& inputs, CAmount change_cost, CAmount target, bool use_effective_value)
 {
     // This function should not be called with empty inputs as that would mean the selection failed
diff --git a/src/wallet/coinselection.h b/src/wallet/coinselection.h
@@ -7,6 +7,7 @@
 
 #include <consensus/amount.h>
 #include <consensus/consensus.h>
+#include <outputtype.h>
 #include <policy/feerate.h>
 #include <primitives/transaction.h>
 #include <random.h>
@@ -249,6 +250,21 @@ struct Groups {
     std::vector<OutputGroup> mixed_group;
 };
 
+/** Stores several 'Groups' whose were mapped by output type. */
+struct OutputGroupTypeMap
+{
+    // Maps output type to output groups.
+    std::map<OutputType, Groups> groups_by_type;
+    // All inserted groups, no type distinction.
+    Groups all_groups;
+
+    // Based on the insert flag; appends group to the 'mixed_group' and, if value > 0, to the 'positive_group'.
+    // This affects both; the groups filtered by type and the overall groups container.
+    void Push(const OutputGroup& group, OutputType type, bool insert_positive, bool insert_mixed);
+    // Retrieves 'Groups' filtered by type
+    std::optional<Groups> Find(OutputType type);
+};
+
 /** Compute the waste for this result given the cost of change
  * and the opportunity cost of spending these inputs now vs in the future.
  * If change exists, waste = change_cost + inputs * (effective_feerate - long_term_feerate)
diff --git a/src/wallet/spend.cpp b/src/wallet/spend.cpp
@@ -404,29 +404,33 @@ std::map<CTxDestination, std::vector<COutput>> ListCoins(const CWallet& wallet)
     return result;
 }
 
-std::vector<OutputGroup> GroupOutputs(const CWallet& wallet, const std::vector<COutput>& outputs, const CoinSelectionParams& coin_sel_params, const CoinEligibilityFilter& filter, bool positive_only)
+OutputGroupTypeMap GroupOutputs(const CWallet& wallet,
+                          const CoinsResult& coins,
+                          const CoinSelectionParams& coin_sel_params,
+                          const CoinEligibilityFilter& filter)
 {
-    std::vector<OutputGroup> groups_out;
+    OutputGroupTypeMap output_groups;
 
     if (!coin_sel_params.m_avoid_partial_spends) {
-        // Allowing partial spends  means no grouping. Each COutput gets its own OutputGroup.
-        for (const COutput& output : outputs) {
-            // Skip outputs we cannot spend
-            if (!output.spendable) continue;
-
-            size_t ancestors, descendants;
-            wallet.chain().getTransactionAncestry(output.outpoint.hash, ancestors, descendants);
-
-            // If 'positive_only' is set, filter for positive only before adding the coin
-            if (!positive_only || output.GetEffectiveValue() > 0) {
-                // Make an OutputGroup containing just this output
-                OutputGroup group{coin_sel_params};
+        // Allowing partial spends means no grouping. Each COutput gets its own OutputGroup
+        for (const auto& [type, outputs] : coins.coins) {
+            for (const COutput& output : outputs) {
+                // Skip outputs we cannot spend
+                if (!output.spendable) continue;
+
+                // Get mempool info
+                size_t ancestors, descendants;
+                wallet.chain().getTransactionAncestry(output.outpoint.hash, ancestors, descendants);
+
+                // Create a new group per output and add it to the all groups vector
+                OutputGroup group(coin_sel_params);
                 group.Insert(std::make_shared<COutput>(output), ancestors, descendants);
 
-                if (group.EligibleForSpending(filter)) groups_out.push_back(group);
+                if (!group.EligibleForSpending(filter)) continue;
+                output_groups.Push(group, type, /*insert_positive=*/true, /*insert_mixed=*/true);
             }
         }
-        return groups_out;
+        return output_groups;
     }
 
     // We want to combine COutputs that have the same scriptPubKey into single OutputGroups
@@ -435,16 +439,12 @@ std::vector<OutputGroup> GroupOutputs(const CWallet& wallet, const std::vector<C
     // For each COutput, we check if the scriptPubKey is in the map, and if it is, the COutput is added
     // to the last OutputGroup in the vector for the scriptPubKey. When the last OutputGroup has
     // OUTPUT_GROUP_MAX_ENTRIES COutputs, a new OutputGroup is added to the end of the vector.
-    std::map<CScript, std::vector<OutputGroup>> spk_to_groups_map;
-    for (const auto& output : outputs) {
-        // Skip outputs we cannot spend
-        if (!output.spendable) continue;
-
-        size_t ancestors, descendants;
-        wallet.chain().getTransactionAncestry(output.outpoint.hash, ancestors, descendants);
-        CScript spk = output.txout.scriptPubKey;
-
-        std::vector<OutputGroup>& groups = spk_to_groups_map[spk];
+    typedef std::map<std::pair<CScript, OutputType>, std::vector<OutputGroup>> ScriptPubKeyToOutgroup;
+    const auto& group_outputs = [](
+            const COutput& output, OutputType type, size_t ancestors, size_t descendants,
+            ScriptPubKeyToOutgroup& groups_map, const CoinSelectionParams& coin_sel_params,
+            bool positive_only) {
+        std::vector<OutputGroup>& groups = groups_map[std::make_pair(output.txout.scriptPubKey,type)];
 
         if (groups.size() == 0) {
             // No OutputGroups for this scriptPubKey yet, add one
@@ -467,28 +467,49 @@ std::vector<OutputGroup> GroupOutputs(const CWallet& wallet, const std::vector<C
         if (!positive_only || output.GetEffectiveValue() > 0) {
             group->Insert(std::make_shared<COutput>(output), ancestors, descendants);
         }
+    };
+
+    ScriptPubKeyToOutgroup spk_to_groups_map;
+    ScriptPubKeyToOutgroup spk_to_positive_groups_map;
+    for (const auto& [type, outs] : coins.coins) {
+        for (const COutput& output : outs) {
+            // Skip outputs we cannot spend
+            if (!output.spendable) continue;
+
+            size_t ancestors, descendants;
+            wallet.chain().getTransactionAncestry(output.outpoint.hash, ancestors, descendants);
+
+            group_outputs(output, type, ancestors, descendants, spk_to_groups_map, coin_sel_params, /*positive_only=*/ false);
+            group_outputs(output, type, ancestors, descendants, spk_to_positive_groups_map,
+                          coin_sel_params, /*positive_only=*/ true);
+        }
     }
 
-    // Now we go through the entire map and pull out the OutputGroups
-    for (const auto& spk_and_groups_pair: spk_to_groups_map) {
-        const std::vector<OutputGroup>& groups_per_spk= spk_and_groups_pair.second;
+    // Now we go through the entire maps and pull out the OutputGroups
+    const auto& push_output_groups = [&](const ScriptPubKeyToOutgroup& groups_map, bool positive_only) {
+        for (const auto& [script, groups] : groups_map) {
+            // Go through the vector backwards. This allows for the first item we deal with being the partial group.
+            for (auto group_it = groups.rbegin(); group_it != groups.rend(); group_it++) {
+                const OutputGroup& group = *group_it;
 
-        // Go through the vector backwards. This allows for the first item we deal with being the partial group.
-        for (auto group_it = groups_per_spk.rbegin(); group_it != groups_per_spk.rend(); group_it++) {
-            const OutputGroup& group = *group_it;
+                if (!group.EligibleForSpending(filter)) continue;
 
-            // Don't include partial groups if there are full groups too and we don't want partial groups
-            if (group_it == groups_per_spk.rbegin() && groups_per_spk.size() > 1 && !filter.m_include_partial_groups) {
-                continue;
-            }
+                // Don't include partial groups if there are full groups too and we don't want partial groups
+                if (group_it == groups.rbegin() && groups.size() > 1 && !filter.m_include_partial_groups) {
+                    continue;
+                }
 
-            // Check the OutputGroup's eligibility. Only add the eligible ones.
-            if (positive_only && group.GetSelectionAmount() <= 0) continue;
-            if (group.m_outputs.size() > 0 && group.EligibleForSpending(filter)) groups_out.push_back(group);
+                OutputType type = script.second;
+                // Either insert the group into the positive-only groups or the mixed ones.
+                output_groups.Push(group, type, positive_only, /*insert_mixed=*/!positive_only);
+            }
         }
-    }
+    };
 
-    return groups_out;
+    push_output_groups(spk_to_groups_map, /*positive_only=*/ false);
+    push_output_groups(spk_to_positive_groups_map, /*positive_only=*/ true);
+
+    return output_groups;
 }
 
 // Returns true if the result contains an error and the message is not empty
@@ -497,13 +518,15 @@ static bool HasErrorMsg(const util::Result<SelectionResult>& res) { return !util
 util::Result<SelectionResult> AttemptSelection(const CWallet& wallet, const CAmount& nTargetValue, const CoinEligibilityFilter& eligibility_filter, const CoinsResult& available_coins,
                                const CoinSelectionParams& coin_selection_params, bool allow_mixed_output_types)
 {
+    // Calculate all the output groups filtered by type at once
+    OutputGroupTypeMap groups = GroupOutputs(wallet, available_coins, coin_selection_params, {eligibility_filter});
+
     // Run coin selection on each OutputType and compute the Waste Metric
     std::vector<SelectionResult> results;
     for (const auto& [type, coins] : available_coins.coins) {
-        Groups groups;
-        groups.positive_group = GroupOutputs(wallet, coins, coin_selection_params, eligibility_filter, true /* positive_only */);
-        groups.mixed_group = GroupOutputs(wallet, coins, coin_selection_params, eligibility_filter, false /* positive_only */);
-        auto result{ChooseSelectionResult(nTargetValue, groups, coin_selection_params)};
+        auto group_for_type = groups.Find(type);
+        if (!group_for_type) continue;
+        auto result{ChooseSelectionResult(nTargetValue, *group_for_type, coin_selection_params)};
         // If any specific error message appears here, then something particularly wrong happened.
         if (HasErrorMsg(result)) return result; // So let's return the specific error.
         // Append the favorable result.
@@ -517,11 +540,7 @@ util::Result<SelectionResult> AttemptSelection(const CWallet& wallet, const CAmo
     // over all available coins, which would allow mixing.
     // If TypesCount() <= 1, there is nothing to mix.
     if (allow_mixed_output_types && available_coins.TypesCount() > 1) {
-        const auto& all = available_coins.All();
-        Groups groups;
-        groups.positive_group = GroupOutputs(wallet, all, coin_selection_params, eligibility_filter, true /* positive_only */);
-        groups.mixed_group = GroupOutputs(wallet, all, coin_selection_params, eligibility_filter, false /* positive_only */);
-        return ChooseSelectionResult(nTargetValue, groups, coin_selection_params);
+        return ChooseSelectionResult(nTargetValue, groups.all_groups, coin_selection_params);
     }
     // Either mixing is not allowed and we couldn't find a solution from any single OutputType, or mixing was allowed and we still couldn't
     // find a solution using all available coins
diff --git a/src/wallet/spend.h b/src/wallet/spend.h
@@ -106,7 +106,14 @@ const CTxOut& FindNonChangeParentOutput(const CWallet& wallet, const COutPoint&
  */
 std::map<CTxDestination, std::vector<COutput>> ListCoins(const CWallet& wallet) EXCLUSIVE_LOCKS_REQUIRED(wallet.cs_wallet);
 
-std::vector<OutputGroup> GroupOutputs(const CWallet& wallet, const std::vector<COutput>& outputs, const CoinSelectionParams& coin_sel_params, const CoinEligibilityFilter& filter, bool positive_only);
+/**
+* Group coins by the provided filters.
+*/
+OutputGroupTypeMap GroupOutputs(const CWallet& wallet,
+                          const CoinsResult& coins,
+                          const CoinSelectionParams& coin_sel_params,
+                          const CoinEligibilityFilter& filter);
+
 /**
  * Attempt to find a valid input set that preserves privacy by not mixing OutputTypes.
  * `ChooseSelectionResult()` will be called on each OutputType individually and the best
diff --git a/src/wallet/test/coinselector_tests.cpp b/src/wallet/test/coinselector_tests.cpp
@@ -153,9 +153,9 @@ inline std::vector<OutputGroup>& KnapsackGroupOutputs(const CoinsResult& availab
         /*tx_noinputs_size=*/ 0,
         /*avoid_partial=*/ false,
     };
-    static std::vector<OutputGroup> static_groups;
-    static_groups = GroupOutputs(wallet, available_coins.All(), coin_selection_params, filter, /*positive_only=*/false);
-    return static_groups;
+    static OutputGroupTypeMap static_groups;
+    static_groups = GroupOutputs(wallet, available_coins, coin_selection_params, {filter});
+    return static_groups.all_groups.mixed_group;
 }
 
 // Branch and bound coin selection tests
diff --git a/src/wallet/test/group_outputs_tests.cpp b/src/wallet/test/group_outputs_tests.cpp
@@ -80,28 +80,28 @@ class GroupVerifier
     CoinsResult coins_pool;
     FastRandomContext rand;
 
-    void GroupVerify(const CoinEligibilityFilter& filter,
+    void GroupVerify(const OutputType type,
+                     const CoinEligibilityFilter& filter,
                      bool avoid_partial_spends,
                      bool positive_only,
                      int expected_size)
     {
-        std::vector<OutputGroup> groups = GroupOutputs(*wallet,
-                                                       coins_pool.All(),
-                                                       makeSelectionParams(rand, avoid_partial_spends),
-                                                       filter,
-                                                       positive_only);
-        BOOST_CHECK_EQUAL(groups.size(), expected_size);
+        OutputGroupTypeMap groups = GroupOutputs(*wallet, coins_pool, makeSelectionParams(rand, avoid_partial_spends), filter);
+        std::vector<OutputGroup>& groups_out = positive_only ? groups.groups_by_type[type].positive_group :
+                                               groups.groups_by_type[type].mixed_group;
+        BOOST_CHECK_EQUAL(groups_out.size(), expected_size);
     }
 
-    void GroupAndVerify(const CoinEligibilityFilter& filter,
+    void GroupAndVerify(const OutputType type,
+                        const CoinEligibilityFilter& filter,
                         int expected_with_partial_spends_size,
                         int expected_without_partial_spends_size,
                         bool positive_only)
     {
         // First avoid partial spends
-        GroupVerify(filter, /*avoid_partial_spends=*/false, positive_only,  expected_with_partial_spends_size);
+        GroupVerify(type, filter, /*avoid_partial_spends=*/false, positive_only,  expected_with_partial_spends_size);
         // Second don't avoid partial spends
-        GroupVerify(filter, /*avoid_partial_spends=*/true, positive_only, expected_without_partial_spends_size);
+        GroupVerify(type, filter, /*avoid_partial_spends=*/true, positive_only, expected_without_partial_spends_size);
     }
 };
 
@@ -125,7 +125,8 @@ BOOST_AUTO_TEST_CASE(outputs_grouping_tests)
         addCoin(group_verifier.coins_pool, *wallet, dest, 10 * COIN, /*is_from_me=*/true);
     }
 
-    group_verifier.GroupAndVerify(BASIC_FILTER,
+    group_verifier.GroupAndVerify(OutputType::BECH32,
+                                  BASIC_FILTER,
                                   /*expected_with_partial_spends_size=*/ GROUP_SIZE,
                                   /*expected_without_partial_spends_size=*/ 1,
                                   /*positive_only=*/ true);
@@ -140,7 +141,8 @@ BOOST_AUTO_TEST_CASE(outputs_grouping_tests)
         addCoin(group_verifier.coins_pool, *wallet, dest2, 5 * COIN, /*is_from_me=*/true);
     }
 
-    group_verifier.GroupAndVerify(BASIC_FILTER,
+    group_verifier.GroupAndVerify(OutputType::BECH32,
+            BASIC_FILTER,
             /*expected_with_partial_spends_size=*/ GROUP_SIZE * 2,
             /*expected_without_partial_spends_size=*/ 2,
             /*positive_only=*/ true);
@@ -154,13 +156,15 @@ BOOST_AUTO_TEST_CASE(outputs_grouping_tests)
     BOOST_CHECK(group_verifier.coins_pool.coins[OutputType::BECH32].back().GetEffectiveValue() <= 0);
 
     // First expect no changes with "positive_only" enabled
-    group_verifier.GroupAndVerify(BASIC_FILTER,
+    group_verifier.GroupAndVerify(OutputType::BECH32,
+            BASIC_FILTER,
             /*expected_with_partial_spends_size=*/ GROUP_SIZE * 2,
             /*expected_without_partial_spends_size=*/ 2,
             /*positive_only=*/ true);
 
     // Then expect changes with "positive_only" disabled
-    group_verifier.GroupAndVerify(BASIC_FILTER,
+    group_verifier.GroupAndVerify(OutputType::BECH32,
+            BASIC_FILTER,
             /*expected_with_partial_spends_size=*/ GROUP_SIZE * 2 + 1,
             /*expected_without_partial_spends_size=*/ 3,
             /*positive_only=*/ false);
@@ -176,7 +180,8 @@ BOOST_AUTO_TEST_CASE(outputs_grouping_tests)
             /*is_from_me=*/false, CFeeRate(0), /*depth=*/5);
 
     // Expect no changes from this round and the previous one (point 4)
-    group_verifier.GroupAndVerify(BASIC_FILTER,
+    group_verifier.GroupAndVerify(OutputType::BECH32,
+            BASIC_FILTER,
             /*expected_with_partial_spends_size=*/ GROUP_SIZE * 2 + 1,
             /*expected_without_partial_spends_size=*/ 3,
             /*positive_only=*/ false);
@@ -192,7 +197,8 @@ BOOST_AUTO_TEST_CASE(outputs_grouping_tests)
             /*is_from_me=*/true, CFeeRate(0), /*depth=*/0);
 
     // Expect no changes from this round and the previous one (point 5)
-    group_verifier.GroupAndVerify(BASIC_FILTER,
+    group_verifier.GroupAndVerify(OutputType::BECH32,
+            BASIC_FILTER,
             /*expected_with_partial_spends_size=*/ GROUP_SIZE * 2 + 1,
             /*expected_without_partial_spends_size=*/ 3,
             /*positive_only=*/ false);
@@ -209,14 +215,16 @@ BOOST_AUTO_TEST_CASE(outputs_grouping_tests)
 
     // Exclude partial groups only adds one more group to the previous test case (point 6)
     int PREVIOUS_ROUND_COUNT = GROUP_SIZE * 2 + 1;
-    group_verifier.GroupAndVerify(BASIC_FILTER,
+    group_verifier.GroupAndVerify(OutputType::BECH32,
+            BASIC_FILTER,
             /*expected_with_partial_spends_size=*/ PREVIOUS_ROUND_COUNT + NUM_SINGLE_ENTRIES,
             /*expected_without_partial_spends_size=*/ 4,
             /*positive_only=*/ false);
 
     // Include partial groups should add one more group inside the "avoid partial spends" count
     const CoinEligibilityFilter& avoid_partial_groups_filter{1, 6, 0, 0, /*include_partial=*/ true};
-    group_verifier.GroupAndVerify(avoid_partial_groups_filter,
+    group_verifier.GroupAndVerify(OutputType::BECH32,
+            avoid_partial_groups_filter,
             /*expected_with_partial_spends_size=*/ PREVIOUS_ROUND_COUNT + NUM_SINGLE_ENTRIES,
             /*expected_without_partial_spends_size=*/ 5,
             /*positive_only=*/ false);
