wallet: Replace CWalletTx::SetTx with Update

diff --git a/src/wallet/transaction.cpp b/src/wallet/transaction.cpp
index 9779fe47..10746d09 100644
--- a/src/wallet/transaction.cpp
+++ b/src/wallet/transaction.cpp
@@ -4,6 +4,7 @@
 
 #include <wallet/transaction.h>
 
+#include <consensus/validation.h>
 #include <interfaces/chain.h>
 
 using interfaces::FoundBlock;
@@ -54,5 +55,73 @@ void CWalletTx::updateState(interfaces::Chain& chain)
     } else if (auto* conf = state<TxStateBlockConflicted>()) {
         lookup_block(conf->conflicting_block_hash, conf->conflicting_block_height, m_state);
     }
+
+    // If the above downgraded a previously-confirmed witness variant back to unconfirmed,
+    // the canonical choice is no longer pinned by confirmation. Re-apply the least-weight rule.
+    if (!isConfirmed()) RecomputeCanonical();
+}
+
+bool CWalletTx::Update(CTransactionRef new_tx, const TxState& new_state)
+{
+    Assert(new_tx);
+    if (!Assume(GetHash() == new_tx->GetHash())) {
+        return false;
+    }
+    bool ret = false;
+    const auto& [tx_pair, inserted] = m_txs.emplace(new_tx->GetWitnessHash(), std::move(new_tx));
+    if (inserted) {
+        ret = true;
+    }
+    const auto& [wtxid, tx] = *tx_pair;
+
+    if (new_state.index() != m_state.index()) {
+        m_state = new_state;
+        if (state<TxStateConfirmed>()) {
+            m_canonical_wtxid = wtxid;
+        }
+        ret = true;
+    } else {
+        assert(TxStateSerializedIndex(m_state) == TxStateSerializedIndex(new_state));
+        assert(TxStateSerializedBlockHash(m_state) == TxStateSerializedBlockHash(new_state));
+    }
+
+    // While unconfirmed, derive the canonical variant from all known variants
+    if (!isConfirmed()) {
+        const Wtxid prev_canonical = m_canonical_wtxid;
+        RecomputeCanonical();
+        if (m_canonical_wtxid != prev_canonical) {
+            ret = true;
+        }
+    }
+
+    return ret;
+}
+
+void CWalletTx::RecomputeCanonical()
+{
+    // Recompute the canonical variant among the witness variants. They share
+    // the txid but differ in the wtxid. Prefer variant with witness data and
+    // the least weight.
+    Assert(!m_txs.empty());
+
+    // Returns true if 'a' should be preferred over 'b'
+    auto is_better = [](const CTransactionRef& a, const CTransactionRef& b) {
+        // A witnessed variant always beats a witnessless one
+        if (a->HasWitness() != b->HasWitness()) return a->HasWitness();
+        // Otherwise the lighter one wins
+        return GetTransactionWeight(*a) < GetTransactionWeight(*b);
+    };
+
+    auto it = m_txs.begin();
+    auto best_wtxid = it->first;
+    const CTransactionRef* best = &it->second;
+    it = std::next(it);
+    for (; it != m_txs.end(); it = std::next(it)) {
+        if (is_better(it->second, *best)) {
+            best = &it->second;
+            best_wtxid = it->first;
+        }
+    }
+    m_canonical_wtxid = best_wtxid;
 }
 } // namespace wallet
diff --git a/src/wallet/transaction.h b/src/wallet/transaction.h
index 96ebeb37..d57f24ac 100644
--- a/src/wallet/transaction.h
+++ b/src/wallet/transaction.h
@@ -353,11 +353,11 @@ public:
 
     CTransactionRef GetTx() const { return m_txs.at(m_canonical_wtxid); }
 
-    void SetTx(CTransactionRef arg)
-    {
-        Assert(arg);
-        m_txs.emplace(arg->GetWitnessHash(), std::move(arg));
-    }
+    // Update the state of this wallet transaction along with a transaction that may have a different wtxid.
+    // If the given transaction has a different wtxid, the transaction is stored if it has not been seen before.
+    // The canonical wtxid is also updated. The tx that is confirmed becomes canonical. For unconfirmed txs,
+    // those with witnesses are preferred, followed by least weight.
+    bool Update(CTransactionRef tx, const TxState& arg_state);
 
     //! make sure balances are recalculated
     void MarkDirty()
@@ -408,6 +408,9 @@ private:
     Wtxid m_canonical_wtxid;
     std::map<Wtxid, CTransactionRef> m_txs;
 
+    //! Set m_canonical_wtxid to the best variant under the unconfirmed rule
+    //! (witnessed preferred, then least weight). Ignores state.
+    void RecomputeCanonical();
 };
 
 struct WalletTxOrderComparator {
diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
index 36a50660..eb40aea0 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -1084,22 +1084,7 @@ CWalletTx* CWallet::AddToWallet(CTransactionRef tx, const TxState& state, const
 
     if (!fInsertedNew)
     {
-        if (state.index() != wtx.m_state.index()) {
-            wtx.m_state = state;
-            fUpdated = true;
-        } else {
-            assert(TxStateSerializedIndex(wtx.m_state) == TxStateSerializedIndex(state));
-            assert(TxStateSerializedBlockHash(wtx.m_state) == TxStateSerializedBlockHash(state));
-        }
-        // If we have a witness-stripped version of this transaction, and we
-        // see a new version with a witness, then we must be upgrading a pre-segwit
-        // wallet.  Store the new version of the transaction with the witness,
-        // as the stripped-version must be invalid.
-        // TODO: Store all versions of the transaction, instead of just one.
-        if (tx->HasWitness() && !wtx.GetTx()->HasWitness()) {
-            wtx.SetTx(tx);
-            fUpdated = true;
-        }
+        fUpdated |= wtx.Update(tx, state);
     }
 
     // Mark inactive coinbase transactions and their descendants as abandoned