1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
use convert_case::{Case, Casing};
use glob::glob;
use serde::Serialize;
use std::{collections::HashMap, ops::Div, path::Path};
use subweight_core::{
    parse::{
        overhead::Weight,
        pallet::{ChromaticExtrinsic, ComponentRange},
        storage::Weights,
    },
    scope::Scope,
    term::Term,
};

// Substrate default maximum weight of a block in nanoseconds.
// Since the maximum block weight is one-third of the execution time,
// it corresponds to a block time of 6 seconds.
const MAX_BLOCK_WEIGHT_NS: f64 = 2_000_000_000_000.;

pub struct MaxBlockWeight(f64);
impl Default for MaxBlockWeight {
    fn default() -> Self {
        MaxBlockWeight(MAX_BLOCK_WEIGHT_NS)
    }
}
impl Div<&MaxBlockWeight> for f64 {
    type Output = Self;

    fn div(self, max_block_weight: &MaxBlockWeight) -> Self::Output {
        self / max_block_weight.0
    }
}
impl MaxBlockWeight {
    pub fn new<T: Into<f64>>(value: T) -> Self {
        MaxBlockWeight(value.into())
    }
}

#[derive(Clone, Debug, Serialize)]
pub struct WeightInfo {
    pub weight: u128,
    pub relative_weight: f64,
}

/// Returns a HashMap <pallet_name, <extrinsic_name, weigh_info>>
/// of the analyzed weights.
///
/// # Arguments
///
///   * `folder_path` - A Path to a folder where the weight files are stored.
///     `paritydb_weights.rs` is mandatory and pallet weights should start by
///     `pallet_`.
///   * `max_block_weight` - The maximal weight of a block.
///
/// # Examples
///
/// ```
///    use weightanalyzer::analyze_weight;
///    use std::path::Path;
///    use weightanalyzer::MaxBlockWeight;
///    let weight_by_pallet = analyze_weight(Path::new("../../runtime/gdev/src/weights/"), &MaxBlockWeight::default());
///    println!("{:?}", weight_by_pallet);
/// ```
pub fn analyze_weight(
    folder_path: &Path,
    max_block_weight: &MaxBlockWeight,
) -> Result<HashMap<String, HashMap<String, WeightInfo>>, String> {
    let pallet_weights = read_pallet_weight(folder_path)?;
    let db_weight = read_db_weight(folder_path)?;
    let overhead_weights = read_overhead_weight(folder_path)?;

    // Initialize scope with db weights
    let mut scope = Scope::from_substrate();
    scope = scope.with_storage_weights(db_weight.weights.read, db_weight.weights.write);

    process(pallet_weights, scope, max_block_weight, &overhead_weights)
}

fn read_pallet_weight(folder_path: &Path) -> Result<Vec<Vec<ChromaticExtrinsic>>, String> {
    let mut parsed_files = Vec::new();
    for path in glob(folder_path.join("*").to_str().expect("Invalid pallet path"))
        .expect("Invalid pallet pattern")
        .filter_map(Result::ok)
    {
        let file = subweight_core::parse::pallet::parse_file(&path);
        if let Ok(file) = file {
            parsed_files.push(file);
        }
    }
    if parsed_files.is_empty() {
        return Err("No pallet found".into());
    }
    Ok(parsed_files)
}

fn read_db_weight(folder_path: &Path) -> Result<Weights, String> {
    subweight_core::parse::storage::parse_file(folder_path.join("paritydb_weights.rs").as_path())
}

fn read_overhead_weight(folder_path: &Path) -> Result<Weight, String> {
    subweight_core::parse::overhead::parse_file(folder_path.join("extrinsic_weights.rs").as_path())
}

fn evaluate_weight(
    extrinsic: ChromaticExtrinsic,
    scope: &mut Scope<Term<u128>>,
    max_block_weight: &MaxBlockWeight,
    overhead: &Weight,
) -> Result<(String, String, WeightInfo), String> {
    // Extend the scope with the maximum value of the complexity parameter.
    if let Some(params) = extrinsic.comp_ranges {
        params
            .iter()
            .for_each(|(key, val): (&String, &ComponentRange)| {
                scope.put_var(key.as_str(), Term::Scalar(val.max.into()));
            });
    }

    // Evaluate the weight
    let mut weight = extrinsic
        .term
        .simplify(subweight_core::Dimension::Time)
        .expect("Can't evaluate")
        .eval(scope)?;

    // Add base extrinsic overhead
    if let Weight::ExtrinsicBase(i) = overhead {
        weight += i
            .simplify(subweight_core::Dimension::Time)
            .expect("Can't evaluate")
            .eval(scope)?;
    }

    let relative_weight = (weight as f64) / max_block_weight * 100.;
    Ok((
        extrinsic
            .pallet
            .to_case(Case::Title)
            .replace("Pallet", "")
            .replace(".rs", "")
            .chars()
            .filter(|c| !c.is_whitespace())
            .collect(),
        extrinsic.name,
        WeightInfo {
            weight,
            relative_weight,
        },
    ))
}

fn process(
    pallet_weights: Vec<Vec<ChromaticExtrinsic>>,
    mut scope: Scope<Term<u128>>,
    max_block_weight: &MaxBlockWeight,
    overhead: &Weight,
) -> Result<HashMap<String, HashMap<String, WeightInfo>>, String> {
    let mut weight_by_pallet: HashMap<String, HashMap<String, WeightInfo>> = HashMap::new();
    for i in pallet_weights {
        for j in i {
            let (pallet, extrinsic, weight) =
                evaluate_weight(j, &mut scope, max_block_weight, overhead)?;
            if let Some(i) = weight_by_pallet.get_mut(&pallet) {
                i.insert(extrinsic, weight);
            } else {
                weight_by_pallet.insert(pallet, HashMap::from([(extrinsic, weight)]));
            }
        }
    }
    Ok(weight_by_pallet)
}

#[cfg(test)]
mod tests {
    use crate::{analyze_weight, MaxBlockWeight};
    use std::path::Path;
    #[test]
    fn should_works() {
        let weight_by_pallet = analyze_weight(
            Path::new("../../runtime/gdev/src/weights/"),
            &MaxBlockWeight::default(),
        );
        assert!(
            weight_by_pallet
                .clone()
                .unwrap()
                .get("Balances")
                .unwrap()
                .len()
                == 10
        ); // 8 extrinsics in pallet
        println!("{:?}", weight_by_pallet); // cargo test  -- --nocapture
    }
    #[test]
    #[should_panic]
    fn should_not_works() {
        let _ = analyze_weight(Path::new(""), &MaxBlockWeight::default()).unwrap();
    }
}