From 9a7604fd257c4a549d0f22f28a36fef607c0dcf4 Mon Sep 17 00:00:00 2001
From: Ava Chow <github@achow101.com>
Date: Wed, 29 Apr 2026 14:34:39 -0700
Subject: [PATCH] bench: Use setup() in WalletMigration to prepare the legacy
 wallet

WalletMigration needs a new wallet with legacy records for each run of
the benchmark. This can be done in setup() rather than duplicating the
records of an initial wallet.
---
 src/bench/wallet_migration.cpp | 101 ++++++++++++++++++---------------
 1 file changed, 56 insertions(+), 45 deletions(-)

diff --git a/src/bench/wallet_migration.cpp b/src/bench/wallet_migration.cpp
index 7d7aca113d1..578fdb50d6b 100644
--- a/src/bench/wallet_migration.cpp
+++ b/src/bench/wallet_migration.cpp
@@ -27,53 +27,64 @@ static void WalletMigration(benchmark::Bench& bench)
     // Number of imported watch only addresses
     int NUM_WATCH_ONLY_ADDR = 20;
 
-    // Setup legacy wallet
-    std::unique_ptr<CWallet> wallet = std::make_unique<CWallet>(test_setup->m_node.chain.get(), "", CreateMockableWalletDatabase());
-    {
-        LegacyDataSPKM* legacy_spkm = wallet->GetOrCreateLegacyDataSPKM();
-        WalletBatch batch{wallet->GetDatabase()};
-
-        // Write a best block record as migration expects one to exist
-        CBlockLocator loc;
-        batch.WriteBestBlock(loc);
-
-        // Add watch-only addresses
-        std::vector<CScript> scripts_watch_only;
-        for (int w = 0; w < NUM_WATCH_ONLY_ADDR; ++w) {
-            CKey key = GenerateRandomKey();
-            LOCK(wallet->cs_wallet);
-            const PKHash dest{key.GetPubKey()};
-            const CScript& script = scripts_watch_only.emplace_back(GetScriptForDestination(dest));
-            assert(legacy_spkm->LoadWatchOnly(script));
-            assert(wallet->SetAddressBook(dest, strprintf("watch_%d", w), /*purpose=*/std::nullopt));
-            batch.WriteWatchOnly(script, CKeyMetadata());
-        }
-
-        // Generate transactions and local addresses
-        for (int j = 0; j < 500; ++j) {
-            CKey key = GenerateRandomKey();
-            CPubKey pubkey = key.GetPubKey();
-            // Load key, scripts and create address book record
-            Assert(legacy_spkm->LoadKey(key, pubkey));
-            CTxDestination dest{PKHash(pubkey)};
-            Assert(wallet->SetAddressBook(dest, strprintf("legacy_%d", j), /*purpose=*/std::nullopt));
-
-            CMutableTransaction mtx;
-            mtx.vout.emplace_back(COIN, GetScriptForDestination(dest));
-            mtx.vout.emplace_back(COIN, scripts_watch_only.at(j % NUM_WATCH_ONLY_ADDR));
-            mtx.vin.resize(2);
-            wallet->AddToWallet(MakeTransactionRef(mtx), TxStateInactive{}, /*update_wtx=*/nullptr, /*rescanning_old_block=*/true);
-            batch.WriteKey(pubkey, key.GetPrivKey(), CKeyMetadata());
-        }
+    // Add watch-only addresses
+    std::vector<std::pair<CScript, CTxDestination>> scripts_watch_only;
+    for (int w = 0; w < NUM_WATCH_ONLY_ADDR; ++w) {
+        CKey key = GenerateRandomKey();
+        const PKHash dest{key.GetPubKey()};
+        scripts_watch_only.emplace_back(GetScriptForDestination(dest), dest);
     }
 
-    bench.epochs(/*numEpochs=*/1).epochIterations(/*numIters=*/1) // run the migration exactly once
-         .run([&] {
-             auto res{MigrateLegacyToDescriptor(std::move(wallet), /*passphrase=*/"", *loader->context())};
-             assert(res);
-             assert(res->wallet);
-             assert(res->watchonly_wallet);
-         });
+    // Generate transactions and local addresses
+    std::vector<CKey> keys(500);
+    std::ranges::generate(keys, []{ return GenerateRandomKey(); });
+
+    std::unique_ptr<CWallet> wallet;
+    size_t i = 0;
+    bench.epochs(/*numEpochs=*/1) // run the migration exactly once
+        .setup([&] {
+            // Setup legacy wallet
+            wallet = std::make_unique<CWallet>(test_setup->m_node.chain.get(), std::string(i++, 'A'), CreateMockableWalletDatabase());
+            LegacyDataSPKM* legacy_spkm = wallet->GetOrCreateLegacyDataSPKM();
+            WalletBatch batch{wallet->GetDatabase()};
+
+            LOCK(wallet->cs_wallet);
+
+            // Write a best block record as migration expects one to exist
+            CBlockLocator loc;
+            batch.WriteBestBlock(loc);
+
+            // Add watch-only addresses
+            for (size_t w = 0; w < scripts_watch_only.size(); ++w) {
+                const auto& [script, dest] = scripts_watch_only.at(w);
+                assert(legacy_spkm->LoadWatchOnly(script));
+                assert(wallet->SetAddressBook(dest, strprintf("watch_%d", w), /*purpose=*/std::nullopt));
+                batch.WriteWatchOnly(script, CKeyMetadata());
+            }
+
+            // Generate transactions and local addresses
+            for (size_t j = 0; j < keys.size(); ++j) {
+                const CKey& key = keys.at(j);
+                // Load key, scripts and create address book record
+                CPubKey pubkey = key.GetPubKey();
+                Assert(legacy_spkm->LoadKey(key, pubkey));
+                CTxDestination dest{PKHash(pubkey)};
+                Assert(wallet->SetAddressBook(dest, strprintf("legacy_%d", j), /*purpose=*/std::nullopt));
+
+                CMutableTransaction mtx;
+                mtx.vout.emplace_back(COIN, GetScriptForDestination(dest));
+                mtx.vout.emplace_back(COIN, scripts_watch_only.at(j % NUM_WATCH_ONLY_ADDR).first);
+                mtx.vin.resize(2);
+                wallet->AddToWallet(MakeTransactionRef(mtx), TxStateInactive{}, /*update_wtx=*/nullptr, /*rescanning_old_block=*/true);
+                batch.WriteKey(pubkey, key.GetPrivKey(), CKeyMetadata());
+            }
+        })
+        .run([&] {
+            auto res{MigrateLegacyToDescriptor(std::move(wallet), /*passphrase=*/"", *loader->context())};
+            assert(res);
+            assert(res->wallet);
+            assert(res->watchonly_wallet);
+        });
 }
 
 BENCHMARK(WalletMigration);