use alloc::{vec, vec::Vec};
use pezsp_runtime::traits::{One, StaticLookup, TrailingZeroInput};
use codec::Decode;
use pezframe_benchmarking::v2::*;
use pezframe_support::{
assert_ok,
traits::{Get, KeyOwnerProofSystem, OnInitialize},
};
use pezframe_system::{pezpallet_prelude::BlockNumberFor, RawOrigin};
use pezpallet_session::{historical::Pezpallet as Historical, Pezpallet as Session, *};
use pezpallet_staking::{
benchmarking::create_validator_with_nominators, testing_utils::create_validators,
MaxNominationsOf, RewardDestination,
};
const MAX_VALIDATORS: u32 = 1000;
pub struct Pezpallet<T: Config>(pezpallet_session::Pezpallet<T>);
pub trait Config:
pezpallet_session::Config + pezpallet_session::historical::Config + pezpallet_staking::Config
{
}
impl<T: Config> OnInitialize<BlockNumberFor<T>> for Pezpallet<T> {
fn on_initialize(n: BlockNumberFor<T>) -> pezframe_support::weights::Weight {
pezpallet_session::Pezpallet::<T>::on_initialize(n)
}
}
#[benchmarks]
mod benchmarks {
use super::*;
#[benchmark]
fn set_keys() -> Result<(), BenchmarkError> {
let n = MaxNominationsOf::<T>::get();
let (v_stash, _) = create_validator_with_nominators::<T>(
n,
MaxNominationsOf::<T>::get(),
false,
true,
RewardDestination::Staked,
)?;
let v_controller =
pezpallet_staking::Pezpallet::<T>::bonded(&v_stash).ok_or("not stash")?;
let keys = T::Keys::decode(&mut TrailingZeroInput::zeroes()).unwrap();
let proof: Vec<u8> = vec![0, 1, 2, 3];
let v_controller_key = pezframe_system::Account::<T>::hashed_key_for(&v_controller);
pezframe_benchmarking::benchmarking::add_to_whitelist(v_controller_key.into());
assert_ok!(Session::<T>::ensure_can_pay_key_deposit(&v_controller));
#[extrinsic_call]
_(RawOrigin::Signed(v_controller), keys, proof);
Ok(())
}
#[benchmark]
fn purge_keys() -> Result<(), BenchmarkError> {
let n = MaxNominationsOf::<T>::get();
let (v_stash, _) = create_validator_with_nominators::<T>(
n,
MaxNominationsOf::<T>::get(),
false,
true,
RewardDestination::Staked,
)?;
let v_controller =
pezpallet_staking::Pezpallet::<T>::bonded(&v_stash).ok_or("not stash")?;
let keys = T::Keys::decode(&mut TrailingZeroInput::zeroes()).unwrap();
let proof: Vec<u8> = vec![0, 1, 2, 3];
assert_ok!(Session::<T>::ensure_can_pay_key_deposit(&v_controller));
Session::<T>::set_keys(RawOrigin::Signed(v_controller.clone()).into(), keys, proof)?;
let v_controller_key = pezframe_system::Account::<T>::hashed_key_for(&v_controller);
pezframe_benchmarking::benchmarking::add_to_whitelist(v_controller_key.into());
#[extrinsic_call]
_(RawOrigin::Signed(v_controller));
Ok(())
}
#[benchmark(extra)]
fn check_membership_proof_current_session(n: Linear<2, MAX_VALIDATORS>) {
let (key, key_owner_proof1) = check_membership_proof_setup::<T>(n);
let key_owner_proof2 = key_owner_proof1.clone();
#[block]
{
Historical::<T>::check_proof(key, key_owner_proof1);
}
assert!(Historical::<T>::check_proof(key, key_owner_proof2).is_some());
}
#[benchmark(extra)]
fn check_membership_proof_historical_session(n: Linear<2, MAX_VALIDATORS>) {
let (key, key_owner_proof1) = check_membership_proof_setup::<T>(n);
Session::<T>::rotate_session();
let key_owner_proof2 = key_owner_proof1.clone();
#[block]
{
Historical::<T>::check_proof(key, key_owner_proof1);
}
assert!(Historical::<T>::check_proof(key, key_owner_proof2).is_some());
}
impl_benchmark_test_suite!(
Pezpallet,
crate::mock::new_test_ext(),
crate::mock::Test,
extra = false
);
}
fn check_membership_proof_setup<T: Config>(
n: u32,
) -> ((pezsp_runtime::KeyTypeId, &'static [u8; 32]), pezsp_session::MembershipProof) {
pezpallet_staking::ValidatorCount::<T>::put(n);
for (n, who) in create_validators::<T>(n, 1000).unwrap().into_iter().enumerate() {
use rand::{RngCore, SeedableRng};
let validator = T::Lookup::lookup(who).unwrap();
let controller = pezpallet_staking::Pezpallet::<T>::bonded(&validator).unwrap();
let keys = {
let mut keys = [0u8; 128];
if n > 0 {
let mut rng = rand::rngs::StdRng::seed_from_u64(n as u64);
rng.fill_bytes(&mut keys);
}
keys
};
let keys: T::Keys = Decode::decode(&mut &keys[..]).unwrap();
let proof: Vec<u8> = vec![];
Session::<T>::set_keys(RawOrigin::Signed(controller).into(), keys, proof).unwrap();
}
Pezpallet::<T>::on_initialize(pezframe_system::pezpallet_prelude::BlockNumberFor::<T>::one());
while Validators::<T>::get().len() < n as usize {
Session::<T>::rotate_session();
}
let key = (pezsp_runtime::KeyTypeId(*b"babe"), &[0u8; 32]);
(key, Historical::<T>::prove(key).unwrap())
}