bench(wallet): Add benchmark of wallet state update

Add a benchmark that measures the time it takes to update a wallet
monitoring 100 UTXOs with a block in which an additional 100 UTXOs are
received.

This benchmark is added because power users report slow state-updating
times, and the *very* slow (up to 5 seconds) state updates have been
pinpointed to being in the state-updater by block for the wallet.

Author: Sword-Smith

Cargo.toml

@@ -296,3 +296,7 @@ required-features = ["arbitrary-impls"]
 [[bench]]
 name = "fast_kernel_mast_hash"
 harness = false
+
+[[bench]]
+name = "wallet_state"
+harness = false

benches/wallet_state.rs

@@ -0,0 +1,168 @@
+use divan::Bencher;
+use neptune_cash::api::export::KeyType;
+use neptune_cash::api::export::NativeCurrencyAmount;
+use neptune_cash::api::export::NeptuneProof;
+use neptune_cash::api::export::Network;
+use neptune_cash::api::export::ReceivingAddress;
+use neptune_cash::api::export::Timestamp;
+use neptune_cash::api::export::Transaction;
+use neptune_cash::api::export::TransactionDetails;
+use neptune_cash::api::export::TransactionProof;
+use neptune_cash::api::export::TxInput;
+use neptune_cash::models::blockchain::block::validity::block_primitive_witness::BlockPrimitiveWitness;
+use neptune_cash::models::blockchain::block::validity::block_proof_witness::BlockProofWitness;
+use neptune_cash::models::blockchain::block::Block;
+use neptune_cash::models::blockchain::block::BlockProof;
+use neptune_cash::models::blockchain::transaction::primitive_witness::PrimitiveWitness;
+use neptune_cash::models::state::wallet::transaction_output::TxOutput;
+use neptune_cash::models::state::wallet::wallet_state::WalletState;
+use num_traits::CheckedSub;
+use rand::Rng;
+
+fn main() {
+    divan::main();
+}
+
+mod maintain_membership_proofs {
+    use super::*;
+
+    /// Maintain 100 membership proofs, while receiving an additional 100 UTXOs.
+    mod maintain_100_100 {
+
+        use neptune_cash::api::benchmarks::devops_wallet_state_genesis;
+
+        use super::*;
+        use crate::helper::next_block_incoming_utxos;
+
+        fn update_wallet_with_block2(bencher: Bencher) {
+            let rt = tokio::runtime::Runtime::new().unwrap();
+            let mut wallet_state = rt.block_on(devops_wallet_state_genesis(Network::Main));
+
+            let genesis = Block::genesis(Network::Main);
+            let own_address = rt
+                .block_on(wallet_state.next_unused_spending_key(KeyType::Generation))
+                .to_address();
+            let block1_time = Network::Main.launch_date() + Timestamp::months(7);
+            let block1 = rt.block_on(next_block_incoming_utxos(
+                &genesis,
+                own_address.clone(),
+                100,
+                &wallet_state,
+                block1_time,
+            ));
+
+            rt.block_on(async {
+                wallet_state
+                    .update_wallet_state_with_new_block(
+                        &genesis.mutator_set_accumulator_after(),
+                        &block1,
+                    )
+                    .await
+                    .unwrap()
+            });
+
+            let block2_time = block1_time + Timestamp::hours(1);
+            let block2 = rt.block_on(next_block_incoming_utxos(
+                &block1,
+                own_address,
+                100,
+                &wallet_state,
+                block2_time,
+            ));
+
+            // Benchmark the receival of 100 UTXOs while already managing 100
+            // UTXOs in the wallet.
+            bencher.bench_local(|| {
+                rt.block_on(async {
+                    wallet_state
+                        .update_wallet_state_with_new_block(
+                            &block1.mutator_set_accumulator_after(),
+                            &block2,
+                        )
+                        .await
+                        .unwrap()
+                });
+            });
+        }
+
+        #[divan::bench(sample_count = 10)]
+        fn apply_block2(bencher: Bencher) {
+            update_wallet_with_block2(bencher);
+        }
+    }
+}
+
+mod helper {
+    use super::*;
+
+    /// Sends the wallet's entire balance to the provided address. Divides the
+    /// transaction up into `N` outputs, guaranteeing that the entire available
+    /// balance is being spent.
+    pub async fn next_block_incoming_utxos(
+        parent: &Block,
+        recipient: ReceivingAddress,
+        num_outputs: usize,
+        sender: &WalletState,
+        timestamp: Timestamp,
+    ) -> Block {
+        let one_nau = NativeCurrencyAmount::from_nau(1);
+
+        let fee = one_nau;
+
+        // create N outputs of 1 nau each
+
+        let mut outputs = vec![(recipient.clone(), one_nau); num_outputs - 1];
+
+        // calc remaining amount and add it to outputs
+        let intermediate_spend = outputs
+            .iter()
+            .map(|(_, amt)| *amt)
+            .sum::<NativeCurrencyAmount>()
+            + fee;
+
+        let msa = parent.mutator_set_accumulator_after();
+        let wallet_status = sender.get_wallet_status(parent.hash(), &msa).await;
+        let change_amt = wallet_status
+            .synced_unspent_available_amount(timestamp)
+            .checked_sub(&intermediate_spend)
+            .unwrap();
+
+        outputs.push((recipient.clone(), change_amt));
+
+        let mut input_funds: Vec<TxInput> = vec![];
+        for input in sender.spendable_inputs(wallet_status, timestamp) {
+            input_funds.push(input);
+        }
+
+        let mut rng = rand::rng();
+        let outputs = outputs.into_iter().map(|(recipient, amount)| {
+            TxOutput::onchain_native_currency_as_change(amount, rng.random(), recipient)
+        });
+        let tx_details = TransactionDetails::new_without_coinbase(
+            input_funds,
+            outputs,
+            fee,
+            timestamp,
+            msa,
+            Network::Main,
+        );
+
+        let kernel = PrimitiveWitness::from_transaction_details(&tx_details).kernel;
+
+        let tx = Transaction {
+            kernel,
+            proof: TransactionProof::SingleProof(NeptuneProof::invalid()),
+        };
+
+        let block_primitive_witness = BlockPrimitiveWitness::new(parent.to_owned(), tx);
+        let body = block_primitive_witness.body().to_owned();
+        let header = block_primitive_witness.header(timestamp, None);
+        let (appendix, proof) = {
+            let block_proof_witness = BlockProofWitness::produce(block_primitive_witness);
+            let appendix = block_proof_witness.appendix();
+            (appendix, BlockProof::SingleProof(NeptuneProof::invalid()))
+        };
+
+        Block::new(header, body, appendix, proof)
+    }
+}

src/api/benchmarks.rs

@@ -0,0 +1,41 @@
+use std::env;
+use std::path::Path;
+use std::path::PathBuf;
+
+use anyhow::Result;
+use rand::distr::Alphanumeric;
+use rand::distr::SampleString;
+
+use crate::api::export::Network;
+use crate::config_models::cli_args;
+use crate::config_models::data_directory::DataDirectory;
+use crate::models::state::wallet::wallet_configuration::WalletConfiguration;
+use crate::models::state::wallet::wallet_entropy::WalletEntropy;
+use crate::models::state::wallet::wallet_state::WalletState;
+
+fn benchmark_data_directory(network: Network) -> Result<DataDirectory> {
+    let mut rng = rand::rng();
+    let user = env::var("USER").unwrap_or_else(|_| "default".to_string());
+    let tmp_root: PathBuf = env::temp_dir()
+        .join(format!("neptune-benchmark-runs-{}", user))
+        .join(Path::new(&Alphanumeric.sample_string(&mut rng, 16)));
+
+    DataDirectory::get(Some(tmp_root), network)
+}
+
+pub async fn devops_wallet_state_genesis(network: Network) -> WalletState {
+    let data_directory = benchmark_data_directory(network).unwrap();
+    DataDirectory::create_dir_if_not_exists(&data_directory.root_dir_path())
+        .await
+        .unwrap();
+
+    DataDirectory::create_dir_if_not_exists(&data_directory.wallet_directory_path())
+        .await
+        .unwrap();
+    let cli_args = cli_args::Args::default();
+    let configuration = WalletConfiguration::new(&data_directory).absorb_options(&cli_args);
+
+    WalletState::try_new(configuration, WalletEntropy::devnet_wallet())
+        .await
+        .unwrap()
+}

src/api/mod.rs

@@ -45,6 +45,7 @@
 //! Please read the [GlobalStateLock](crate::GlobalStateLock) docs carefully because it is critical
 //! not to hold the lock too long or cause a deadlock situation.
 mod api_impl;
+pub mod benchmarks;
 pub mod export;
 pub mod regtest;
 pub mod tx_initiation;

src/config_models/network.rs

@@ -40,7 +40,9 @@ pub enum Network {
 }
 
 impl Network {
-    pub(crate) fn launch_date(&self) -> Timestamp {
+    /// Launch date of the network in question. Defined as the timestamp of the
+    /// genesis block.
+    pub fn launch_date(&self) -> Timestamp {
         match self {
             Network::RegTest => {
                 const SEVEN_DAYS: u64 = 1000 * 60 * 60 * 24 * 7;

src/models/blockchain/block/block_appendix.rs

@@ -27,7 +27,7 @@ pub(crate) const MAX_NUM_CLAIMS: usize = 500;
 /// The appendix can softly be extended with new claims.
 #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, BFieldCodec, GetSize, Default)]
 #[cfg_attr(any(test, feature = "arbitrary-impls"), derive(Arbitrary))]
-pub(crate) struct BlockAppendix {
+pub struct BlockAppendix {
     claims: Vec<Claim>,
 }
 

src/models/blockchain/block/mod.rs

@@ -379,7 +379,7 @@ impl Block {
     ///
     /// Includes the guesser-fee UTXOs which are not included by the
     /// `mutator_set_accumulator` field on the block body.
-    pub(crate) fn mutator_set_accumulator_after(&self) -> MutatorSetAccumulator {
+    pub fn mutator_set_accumulator_after(&self) -> MutatorSetAccumulator {
         let mut msa = self.kernel.body.mutator_set_accumulator.clone();
         let mutator_set_update = MutatorSetUpdate::new(vec![], self.guesser_fee_addition_records());
         mutator_set_update.apply_to_accumulator(&mut msa)
@@ -621,7 +621,7 @@ impl Block {
         utxos
     }
 
-    pub(crate) fn new(
+    pub fn new(
         header: BlockHeader,
         body: BlockBody,
         appendix: BlockAppendix,

src/models/blockchain/block/validity/block_primitive_witness.rs

@@ -39,15 +39,15 @@ use crate::models::proof_abstractions::timestamp::Timestamp;
 // note: I removed unused PartialEq, Eq derive.   If we ever need one
 // PartialEq should be derived manually to ignore maybe_body.
 #[derive(Clone, Debug)]
-pub(crate) struct BlockPrimitiveWitness {
+pub struct BlockPrimitiveWitness {
     predecessor_block: Block,
     transaction: Transaction,
 
     maybe_body: OnceLock<BlockBody>,
 }
 
 impl BlockPrimitiveWitness {
-    pub(crate) fn new(predecessor_block: Block, transaction: Transaction) -> Self {
+    pub fn new(predecessor_block: Block, transaction: Transaction) -> Self {
         Self {
             predecessor_block,
             transaction,
@@ -59,7 +59,7 @@ impl BlockPrimitiveWitness {
         &self.transaction
     }
 
-    pub(crate) fn header(
+    pub fn header(
         &self,
         timestamp: Timestamp,
         target_block_interval: Option<Timestamp>,
@@ -74,7 +74,7 @@ impl BlockPrimitiveWitness {
         )
     }
 
-    pub(crate) fn body(&self) -> &BlockBody {
+    pub fn body(&self) -> &BlockBody {
         self.maybe_body.get_or_init(|| {
             let predecessor_msa_digest = self.predecessor_block
             .mutator_set_accumulator_after()

src/models/blockchain/block/validity/block_proof_witness.rs

@@ -31,7 +31,7 @@ use crate::models::proof_abstractions::SecretWitness;
 ///
 /// This is the witness for the [`BlockProgram`].
 #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, GetSize, BFieldCodec, TasmObject)]
-pub(crate) struct BlockProofWitness {
+pub struct BlockProofWitness {
     pub(super) block_body: BlockBody,
     pub(crate) claims: Vec<Claim>,
     pub(crate) proofs: Vec<Proof>,
@@ -63,11 +63,11 @@ impl BlockProofWitness {
         self.claims.clone()
     }
 
-    pub(crate) fn appendix(&self) -> BlockAppendix {
+    pub fn appendix(&self) -> BlockAppendix {
         BlockAppendix::new(self.claims())
     }
 
-    pub(crate) fn produce(block_primitive_witness: BlockPrimitiveWitness) -> BlockProofWitness {
+    pub fn produce(block_primitive_witness: BlockPrimitiveWitness) -> BlockProofWitness {
         let txk_mast_hash = block_primitive_witness
             .body()
             .transaction_kernel

src/models/blockchain/transaction/primitive_witness.rs

@@ -241,7 +241,7 @@ impl PrimitiveWitness {
     }
 
     /// Create a [`PrimitiveWitness`] from [`TransactionDetails`].
-    pub(crate) fn from_transaction_details(transaction_details: &TransactionDetails) -> Self {
+    pub fn from_transaction_details(transaction_details: &TransactionDetails) -> Self {
         let TransactionDetails {
             tx_inputs,
             tx_outputs,

src/models/state/transaction_details.rs

@@ -215,7 +215,7 @@ impl TransactionDetails {
     /// - mutator set membership proofs, must be valid wrt. supplied mutator set
     ///
     /// See also: [Self::new_with_coinbase].
-    pub(crate) fn new_without_coinbase(
+    pub fn new_without_coinbase(
         tx_inputs: impl Into<TxInputList>,
         tx_outputs: impl Into<TxOutputList>,
         fee: NativeCurrencyAmount,

src/models/state/wallet/mod.rs

@@ -17,7 +17,7 @@ pub(crate) mod wallet_configuration;
 pub(crate) mod wallet_db_tables;
 pub(crate) mod wallet_entropy;
 pub mod wallet_file;
-pub(crate) mod wallet_state;
+pub mod wallet_state;
 pub mod wallet_status;
 
 #[cfg(test)]

src/models/state/wallet/transaction_output.rs

@@ -292,7 +292,7 @@ impl TxOutput {
 
     /// Instantiate a [TxOutput] for native currency intended fro on-chain UTXO
     /// notification.
-    pub(crate) fn onchain_native_currency(
+    pub fn onchain_native_currency(
         amount: NativeCurrencyAmount,
         sender_randomness: Digest,
         receiving_address: ReceivingAddress,
@@ -311,7 +311,7 @@ impl TxOutput {
 
     /// Instantiate a [TxOutput] for native currency intended fro on-chain UTXO
     /// notification.
-    pub(crate) fn onchain_native_currency_as_change(
+    pub fn onchain_native_currency_as_change(
         amount: NativeCurrencyAmount,
         sender_randomness: Digest,
         receiving_address: ReceivingAddress,

src/models/state/wallet/wallet_state.rs

@@ -1345,7 +1345,7 @@ impl WalletState {
     ///
     /// Assume the given block is valid and that the wallet state is not synced
     /// with the new block yet but is synced with the previous block (if any).
-    pub(crate) async fn update_wallet_state_with_new_block(
+    pub async fn update_wallet_state_with_new_block(
         &mut self,
         previous_mutator_set_accumulator: &MutatorSetAccumulator,
         new_block: &Block,
@@ -1424,6 +1424,7 @@ impl WalletState {
                 all_existing_mutxos,
             )
         }
+
         let tx_kernel = new_block.kernel.body.transaction_kernel.clone();
 
         let spent_inputs: Vec<(Utxo, AbsoluteIndexSet, u64)> =
@@ -1837,7 +1838,7 @@ impl WalletState {
     ///     + that are timelocked in the future
     ///     + that are unspendable (no spending key)
     ///     + that are already spent in the mempool
-    pub(crate) fn spendable_inputs(
+    pub fn spendable_inputs(
         &self,
         wallet_status: WalletStatus,
         timestamp: Timestamp,