wallet: Introduce and use PreselectedInput class in CCoinControl

Instead of having different maps for selected inputs, external inputs, and input weight in CCoinControl, have a class PreselectedInput which tracks stores that information for each input.
lateminer · Dec 8, 2023 · e1abfb5 · e1abfb5
1 parent dce1dfb
commit e1abfb5
Show file tree

Hide file tree

Showing 2 changed files with 89 additions and 32 deletions.
diff --git a/src/wallet/coincontrol.cpp b/src/wallet/coincontrol.cpp
@@ -14,69 +14,104 @@ CCoinControl::CCoinControl()
 
 bool CCoinControl::HasSelected() const
 {
-    return !m_selected_inputs.empty();
+    return !m_selected.empty();
 }
 
-bool CCoinControl::IsSelected(const COutPoint& output) const
+bool CCoinControl::IsSelected(const COutPoint& outpoint) const
 {
-    return m_selected_inputs.count(output) > 0;
+    return m_selected.count(outpoint) > 0;
 }
 
-bool CCoinControl::IsExternalSelected(const COutPoint& output) const
+bool CCoinControl::IsExternalSelected(const COutPoint& outpoint) const
 {
-    return m_external_txouts.count(output) > 0;
+    const auto it = m_selected.find(outpoint);
+    return it != m_selected.end() && it->second.HasTxOut();
 }
 
 std::optional<CTxOut> CCoinControl::GetExternalOutput(const COutPoint& outpoint) const
 {
-    const auto ext_it = m_external_txouts.find(outpoint);
-    if (ext_it == m_external_txouts.end()) {
+    const auto it = m_selected.find(outpoint);
+    if (it == m_selected.end() || !it->second.HasTxOut()) {
         return std::nullopt;
     }
-
-    return std::make_optional(ext_it->second);
+    return it->second.GetTxOut();
 }
 
-void CCoinControl::Select(const COutPoint& output)
+PreselectedInput& CCoinControl::Select(const COutPoint& outpoint)
 {
-    m_selected_inputs.insert(output);
+    return m_selected[outpoint];
 }
 
 void CCoinControl::SelectExternal(const COutPoint& outpoint, const CTxOut& txout)
 {
-    m_selected_inputs.insert(outpoint);
-    m_external_txouts.emplace(outpoint, txout);
+    m_selected[outpoint].SetTxOut(txout);
 }
 
-void CCoinControl::UnSelect(const COutPoint& output)
+void CCoinControl::UnSelect(const COutPoint& outpoint)
 {
-    m_selected_inputs.erase(output);
+    m_selected.erase(outpoint);
 }
 
 void CCoinControl::UnSelectAll()
 {
-    m_selected_inputs.clear();
+    m_selected.clear();
 }
 
 std::vector<COutPoint> CCoinControl::ListSelected() const
 {
-    return {m_selected_inputs.begin(), m_selected_inputs.end()};
+    std::vector<COutPoint> outpoints;
+    std::transform(m_selected.begin(), m_selected.end(), std::back_inserter(outpoints),
+        [](const std::map<COutPoint, PreselectedInput>::value_type& pair) {
+            return pair.first;
+        });
+    return outpoints;
 }
 
 void CCoinControl::SetInputWeight(const COutPoint& outpoint, int64_t weight)
 {
-    m_input_weights[outpoint] = weight;
+    m_selected[outpoint].SetInputWeight(weight);
 }
 
 bool CCoinControl::HasInputWeight(const COutPoint& outpoint) const
 {
-    return m_input_weights.count(outpoint) > 0;
+    const auto it = m_selected.find(outpoint);
+    return it != m_selected.end() && it->second.HasInputWeight();
 }
 
 int64_t CCoinControl::GetInputWeight(const COutPoint& outpoint) const
 {
-    auto it = m_input_weights.find(outpoint);
-    assert(it != m_input_weights.end());
-    return it->second;
+    return m_selected.at(outpoint).GetInputWeight();
+}
+
+void PreselectedInput::SetTxOut(const CTxOut& txout)
+{
+    m_txout = txout;
+}
+
+CTxOut PreselectedInput::GetTxOut() const
+{
+    assert(m_txout.has_value());
+    return m_txout.value();
+}
+
+bool PreselectedInput::HasTxOut() const
+{
+    return m_txout.has_value();
+}
+
+void PreselectedInput::SetInputWeight(int64_t weight)
+{
+    m_weight = weight;
+}
+
+int64_t PreselectedInput::GetInputWeight() const
+{
+    assert(m_weight.has_value());
+    return m_weight.value();
+}
+
+bool PreselectedInput::HasInputWeight() const
+{
+    return m_weight.has_value();
 }
 } // namespace wallet
diff --git a/src/wallet/coincontrol.h b/src/wallet/coincontrol.h
@@ -24,6 +24,33 @@ const int DEFAULT_MAX_DEPTH = 9999999;
 //! Default for -avoidpartialspends
 static constexpr bool DEFAULT_AVOIDPARTIALSPENDS = false;
 
+class PreselectedInput
+{
+private:
+    //! The previous output being spent by this input
+    std::optional<CTxOut> m_txout;
+    //! The input weight for spending this input
+    std::optional<int64_t> m_weight;
+
+public:
+    /**
+     * Set the previous output for this input.
+     * Only necessary if the input is expected to be an external input.
+     */
+    void SetTxOut(const CTxOut& txout);
+    /** Retrieve the previous output for this input. */
+    CTxOut GetTxOut() const;
+    /** Return whether the previous output is set for this input. */
+    bool HasTxOut() const;
+
+    /** Set the weight for this input. */
+    void SetInputWeight(int64_t weight);
+    /** Retrieve the input weight for this input. */
+    int64_t GetInputWeight() const;
+    /** Return whether the input weight is set. */
+    bool HasInputWeight() const;
+};
+
 /** Coin Control Features. */
 class CCoinControl
 {
@@ -69,11 +96,11 @@ class CCoinControl
     /**
      * Returns true if the given output is pre-selected.
      */
-    bool IsSelected(const COutPoint& output) const;
+    bool IsSelected(const COutPoint& outpoint) const;
     /**
      * Returns true if the given output is selected as an external input.
      */
-    bool IsExternalSelected(const COutPoint& output) const;
+    bool IsExternalSelected(const COutPoint& outpoint) const;
     /**
      * Returns the external output for the given outpoint if it exists.
      */
@@ -82,7 +109,7 @@ class CCoinControl
      * Lock-in the given output for spending.
      * The output will be included in the transaction even if it's not the most optimal choice.
      */
-    void Select(const COutPoint& output);
+    PreselectedInput& Select(const COutPoint& outpoint);
     /**
      * Lock-in the given output as an external input for spending because it is not in the wallet.
      * The output will be included in the transaction even if it's not the most optimal choice.
@@ -91,7 +118,7 @@ class CCoinControl
     /**
      * Unselects the given output.
      */
-    void UnSelect(const COutPoint& output);
+    void UnSelect(const COutPoint& outpoint);
     /**
      * Unselects all outputs.
      */
@@ -115,12 +142,7 @@ class CCoinControl
 
 private:
     //! Selected inputs (inputs that will be used, regardless of whether they're optimal or not)
-    std::set<COutPoint> m_selected_inputs;
-    //! Map of external inputs to include in the transaction
-    //! These are not in the wallet, so we need to track them separately
-    std::map<COutPoint, CTxOut> m_external_txouts;
-    //! Map of COutPoints to the maximum weight for that input
-    std::map<COutPoint, int64_t> m_input_weights;
+    std::map<COutPoint, PreselectedInput> m_selected;
 };
 } // namespace wallet