wip: add bench tests

Author: ValuedMammal

bench/benches/bench.rs

@@ -3,5 +3,8 @@ extern crate criterion;
 
 use criterion::{criterion_group, criterion_main};
 
-criterion_group!(benches, bdk_chain::tx_graph::bench::filter_chain_unspents);
+criterion_group!(benches,
+    bdk_chain::tx_graph::bench::filter_chain_unspents,
+    bdk_chain::tx_graph::bench::list_canonical_txs,
+);
 criterion_main!(benches);

crates/chain/src/tx_graph.rs

@@ -1554,67 +1554,55 @@ fn tx_outpoint_range(txid: Txid) -> RangeInclusive<OutPoint> {
     OutPoint::new(txid, u32::MIN)..=OutPoint::new(txid, u32::MAX)
 }
 
-/// Bench
-#[allow(unused)]
-#[allow(missing_docs)]
 #[cfg(bdk_bench)]
-pub mod bench {
-    use std::str::FromStr;
-
+mod bench_util {
     use bdk_core::{CheckPoint, ConfirmationBlockTime};
-    use bitcoin::absolute;
-    use bitcoin::hashes::Hash;
-    use bitcoin::transaction;
-    use bitcoin::{Address, BlockHash, Network, TxIn};
-    use criterion::Criterion;
-    use miniscript::Descriptor;
-    use miniscript::DescriptorPublicKey;
+    use bitcoin::{
+        absolute, constants, hashes::Hash, secp256k1::Secp256k1, transaction, BlockHash, Network,
+        TxIn,
+    };
+    use miniscript::{Descriptor, DescriptorPublicKey};
 
     use super::*;
     use crate::keychain_txout::KeychainTxOutIndex;
     use crate::local_chain::LocalChain;
     use crate::IndexedTxGraph;
 
     #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
-    enum Keychain {
+    pub enum Keychain {
         External,
     }
 
     const EXTERNAL: &str = "tr([ab28dc00/86h/1h/0h]tpubDCdDtzAMZZrkwKBxwNcGCqe4FRydeD9rfMisoi7qLdraG79YohRfPW4YgdKQhpgASdvh612xXNY5xYzoqnyCgPbkpK4LSVcH5Xv4cK7johH/0/*)";
 
-    fn get_params() -> (
-        IndexedTxGraph<ConfirmationBlockTime, KeychainTxOutIndex<Keychain>>,
-        LocalChain,
-    ) {
-        let genesis = bitcoin::constants::genesis_block(Network::Regtest).block_hash();
+    type KeychainTxGraph = IndexedTxGraph<ConfirmationBlockTime, KeychainTxOutIndex<Keychain>>;
+
+    pub fn init_graph_chain() -> (KeychainTxGraph, LocalChain) {
+        let genesis = constants::genesis_block(Network::Regtest).block_hash();
         let block_0 = BlockId {
             height: 0,
             hash: genesis,
         };
+        // chain tip 100
         let mut cp = CheckPoint::new(block_0);
-        let block_100 = BlockId {
+        let chain_tip = BlockId {
             height: 100,
             hash: BlockHash::all_zeros(),
         };
-        cp = cp.push(block_100).unwrap();
+        cp = cp.push(chain_tip).unwrap();
         let chain = LocalChain::from_tip(cp).unwrap();
 
-        let mut graph = IndexedTxGraph::new({
-            let mut index = KeychainTxOutIndex::new(10);
-            index
-                .insert_descriptor(Keychain::External, parse_descriptor(EXTERNAL))
-                .unwrap();
-            index
-        });
+        let desc = parse_descriptor(EXTERNAL);
+        let spk_0 = desc.at_derivation_index(0).unwrap().script_pubkey();
+        let mut index = KeychainTxOutIndex::new(10);
+        index.insert_descriptor(Keychain::External, desc).unwrap();
+        let mut graph = IndexedTxGraph::new(index);
 
-        // insert funding tx (coinbase)
-        let addr_0 =
-            Address::from_str("bcrt1plhmjhj75nut38qwwm5w7xqysy25xhd4ckuv7zu5tey3nkmcwh3cqvan5mz")
-                .unwrap()
-                .assume_checked();
+        // insert funding tx (coinbase) confirmed at chain tip
         let tx_0 = Transaction {
+            input: vec![TxIn::default()],
             output: vec![TxOut {
-                script_pubkey: addr_0.script_pubkey(),
+                script_pubkey: spk_0.clone(),
                 value: Amount::ONE_BTC,
             }],
             ..new_tx(0)
@@ -1624,43 +1612,140 @@ pub mod bench {
         let _ = graph.insert_anchor(
             txid_0,
             ConfirmationBlockTime {
-                block_id: block_100,
+                block_id: chain_tip,
                 confirmation_time: 100,
             },
         );
 
         (graph, chain)
     }
 
-    fn parse_descriptor(s: &str) -> miniscript::Descriptor<DescriptorPublicKey> {
-        <Descriptor<DescriptorPublicKey>>::parse_descriptor(
-            &bitcoin::secp256k1::Secp256k1::new(),
-            s,
-        )
-        .unwrap()
-        .0
+    /// Add `n` conflicts to `graph` that spend the given `previous_output`, incrementing
+    /// the tx last-seen on each iteration.
+    pub fn add_conflicts(n: u32, graph: &mut KeychainTxGraph, previous_output: OutPoint) {
+        let spk_1 = parse_descriptor(EXTERNAL)
+            .at_derivation_index(1)
+            .unwrap()
+            .script_pubkey();
+        for i in 1..n + 1 {
+            let value = Amount::ONE_BTC - Amount::from_sat(i as u64 * 10);
+            let tx = Transaction {
+                input: vec![TxIn {
+                    previous_output,
+                    ..Default::default()
+                }],
+                output: vec![TxOut {
+                    value,
+                    script_pubkey: spk_1.clone(),
+                }],
+                ..new_tx(i)
+            };
+            let txid = tx.compute_txid();
+            let _ = graph.insert_tx(tx);
+            let _ = graph.insert_seen_at(txid, i as u64);
+        }
     }
 
-    fn new_tx(lt: u32) -> Transaction {
+    /// Apply a chain of `n` unconfirmed txs where each subsequent
+    /// tx spends the output of the previous one.
+    pub fn chain_unconfirmed(n: u32, graph: &mut KeychainTxGraph) {
+        let (mut previous_output, txout) = graph
+            .graph()
+            .all_txouts()
+            .find(|(_, txout)| txout.value == Amount::ONE_BTC)
+            .expect("initial graph should have txout");
+        let init_amount = txout.value;
+        let fee = Amount::from_sat(1000);
+
+        let desc = graph.index.get_descriptor(Keychain::External).unwrap();
+        let spks: Vec<_> = (0..n)
+            .map(|i| desc.at_derivation_index(i).unwrap().script_pubkey())
+            .collect();
+
+        for i in 0..n {
+            // create tx
+            let tx = Transaction {
+                input: vec![TxIn {
+                    previous_output,
+                    ..Default::default()
+                }],
+                output: vec![TxOut {
+                    value: init_amount - (fee * (i as u64 + 1)),
+                    script_pubkey: spks[i as usize].clone(),
+                }],
+                ..new_tx(i)
+            };
+            let txid = tx.compute_txid();
+            let update = TxUpdate {
+                txs: vec![Arc::new(tx)],
+                ..Default::default()
+            };
+            let _ = graph.apply_update_at(update, Some(21));
+            // store the next prevout
+            previous_output = OutPoint::new(txid, 0);
+        }
+    }
+
+    pub fn parse_descriptor(s: &str) -> Descriptor<DescriptorPublicKey> {
+        <Descriptor<DescriptorPublicKey>>::parse_descriptor(&Secp256k1::new(), s)
+            .unwrap()
+            .0
+    }
+
+    pub fn new_tx(lt: u32) -> Transaction {
         Transaction {
             version: transaction::Version::TWO,
             lock_time: absolute::LockTime::from_consensus(lt),
             input: vec![],
             output: vec![],
         }
     }
+}
 
+/// Bench
+#[cfg(bdk_bench)]
+#[allow(missing_docs)]
+pub mod bench {
+    use std::hint::black_box;
+
+    use criterion::Criterion;
+
+    use super::*;
+    use crate::tx_graph::bench_util::*;
+
+    #[inline(never)]
     pub fn filter_chain_unspents(bench: &mut Criterion) {
-        let (graph, chain) = get_params();
-        // TODO: insert conflicts
+        let (mut graph, chain) = black_box(init_graph_chain());
+        let txouts: Vec<_> = graph.graph().all_txouts().collect();
+        assert_eq!(txouts.len(), 1);
+        let prevout = txouts.iter().next().unwrap().0;
+        black_box(add_conflicts(1000, &mut graph, prevout));
+
         bench.bench_function("filter_chain_unspents", |b| {
             b.iter(|| {
-                TxGraph::filter_chain_unspents(
+                let unspent = TxGraph::filter_chain_unspents(
                     graph.graph(),
                     &chain,
                     chain.tip().block_id(),
                     graph.index.outpoints().clone(),
-                )
+                );
+                assert_eq!(unspent.count(), 1);
+            })
+        });
+    }
+
+    #[inline(never)]
+    pub fn list_canonical_txs(bench: &mut Criterion) {
+        let (mut graph, chain) = black_box(init_graph_chain());
+        assert_eq!(graph.graph().txs.len(), 1);
+        black_box(chain_unconfirmed(100, &mut graph));
+
+        bench.bench_function("list_canonical_txs", |b| {
+            b.iter(|| {
+                let txs = graph
+                    .graph()
+                    .list_canonical_txs(&chain, chain.tip().block_id());
+                assert_eq!(txs.count(), 100 + 1);
             })
         });
     }