chia_sdk_test/
benchmark.rs

1use std::{collections::HashMap, fs::File};
2
3use chia_bls::{SecretKey, Signature};
4use chia_consensus::spendbundle_conditions::get_conditions_from_spendbundle;
5use chia_protocol::{CoinSpend, SpendBundle};
6use chia_sdk_types::TESTNET11_CONSTANTS;
7use clvmr::Allocator;
8use prettytable::{Table, row};
9
10use crate::Simulator;
11
12#[derive(Debug, Clone)]
13pub struct Benchmark {
14    pub title: String,
15    pub data_keys: Vec<String>,
16    pub data: HashMap<String, Vec<u64>>,
17}
18
19impl Benchmark {
20    pub fn new(title: String) -> Self {
21        Self {
22            title,
23            data_keys: Vec::new(),
24            data: HashMap::new(),
25        }
26    }
27
28    pub fn add_spends(
29        &mut self,
30        ctx: &mut Allocator,
31        sim: &mut Simulator,
32        coin_spends: Vec<CoinSpend>,
33        key: &str,
34        keys: &[SecretKey],
35    ) -> anyhow::Result<()> {
36        let sb = SpendBundle::new(coin_spends, Signature::default());
37        let sb_conds = get_conditions_from_spendbundle(
38            ctx,
39            &sb,
40            u64::MAX,
41            sim.height(),
42            &TESTNET11_CONSTANTS,
43        )?;
44
45        let key = key.to_string();
46        if !self.data_keys.contains(&key) {
47            self.data_keys.push(key.clone());
48        }
49        self.data.entry(key).or_default().push(sb_conds.cost);
50
51        sim.spend_coins(sb.coin_spends, keys)?;
52        Ok(())
53    }
54
55    #[allow(clippy::cast_precision_loss)]
56    pub fn print_summary(&self, filename: Option<&str>) {
57        let mut table = Table::new();
58        table.add_row(row![format!("Cost statistics for {}", self.title)]);
59        table.add_row(row!["label", "avg", "n", "min", "max", "median"]);
60        for key in &self.data_keys {
61            let data = &self.data[key];
62
63            let total = data.iter().sum::<u64>();
64            let avg = format!("{:.1}", total as f64 / data.len() as f64);
65
66            let mut sorted = data.clone();
67            sorted.sort_unstable();
68            let data_min = sorted[0];
69            let data_max = sorted[sorted.len() - 1];
70
71            let data_median = if sorted.len().is_multiple_of(2) {
72                (sorted[sorted.len() / 2] + sorted[sorted.len() / 2 - 1]) as f64 / 2.0
73            } else {
74                sorted[sorted.len() / 2] as f64
75            };
76
77            table.add_row(row![key, avg, data.len(), data_min, data_max, data_median]);
78        }
79
80        table.printstd();
81        if let Some(filename) = filename {
82            let mut file = File::create(filename).unwrap();
83            table.print(&mut file).unwrap();
84        }
85    }
86}