ghost-node/pallets/governance/src/benchmarking.rs

#![cfg(feature = "runtime-benchmarks")]

use super::*;
use frame_benchmarking::v2::*;

use frame_support::dispatch::RawOrigin;
use ghost_helpers::merkle_tree::{generate_tree, generate_proof};

#[benchmarks(where T: Config)]
mod benchmarks {
    use super::*;

    #[benchmark]
    fn claim() {
        let network_id = Default::default();
        let current_account: T::AccountId = frame_benchmarking::whitelisted_caller();

        let max_proof_depth = <<T as Config>::MaxProofDepth as Get<u32>>::get();
        let minimum_donation_u128 = <<T as Config>::MinimumDonation as Get<u128>>::get();
        let minimum_donation: BalanceOf<T> = minimum_donation_u128.unique_saturated_into();

        let dummy_shares_u128 = 420u128;
        let dummy_shares: BalanceOf<T> = dummy_shares_u128.unique_saturated_into();

        let secret_key = libsecp256k1::SecretKey::parse(&[69u8; 32]).unwrap();
        let public_key = libsecp256k1::PublicKey::from_secret_key(&secret_key);
        let pubkey_raw = &public_key.serialize()[1..65];
        let pubkey_hash = SubstrateKeccakHasher::hash(pubkey_raw);

        let mut evm_addr_bytes = [0u8; 20];
        evm_addr_bytes.copy_from_slice(&pubkey_hash.as_ref()[12..32]);
        let dummy_evm_address = EvmAddress::from(evm_addr_bytes);

        let total_entries = 1u64 << max_proof_depth;
        let max_index = total_entries - 1;
        let raw_values = (0..total_entries)
            .map(|i| i as TokenId)
            .collect::<Vec<TokenId>>();

        let merkle_tree = generate_tree::<SubstrateKeccakHasher, _, _, _, _>(
            max_index,
            raw_values,
            |index: TokenId| -> Result<(usize, Vec<u8>), ()> {
                let claim_package = ClaimPackage {
                    shares: dummy_shares,
                    merkle_proof: Default::default(),
                    token_id: index,
                    index: index as u32,
                };

                let claim_preimage = claim_package.get_preimage(
                    &dummy_evm_address,
                    network_id,
                );
                Ok((index as usize, claim_preimage.to_vec()))
            }
        ).unwrap();

        let merkle_root = *merkle_tree.last().unwrap();

        let member_share = MemberShare {
            initiated_network: None,
            activated_shares: Default::default(),
            locked_balance: minimum_donation,
        };
        MemberShares::<T>::insert(&current_account, member_share);

        let network_state = NetworkShare::new(
            dummy_shares,
            merkle_root,
            max_proof_depth,
        );
        NetworkShares::<T>::insert(&network_id, network_state);

        let global_state = ShareState::new(dummy_shares);
        GlobalShares::<T>::put(global_state);

        let merkle_proof = generate_proof::<SubstrateKeccakHasher, _>(
            &merkle_tree,
            total_entries - 1,
            0,
        );

        let claim_package = ClaimPackage {
            index: 0,
            merkle_proof: BoundedVec::try_from(merkle_proof).unwrap(),
            shares: dummy_shares,
            token_id: 0
        };

        let preimage = claim_package.get_preimage(&dummy_evm_address, network_id);
        let preimage_hash = SubstrateKeccakHasher::hash(&preimage);

        let message_string_bytes = Pallet::<T>::ethereum_signable_message(
            &current_account,
            &network_id,
            &preimage_hash
        );

        let message_hash = SubstrateKeccakHasher::hash(&message_string_bytes);
        let mut message_bytes = [0u8; 32];
        message_bytes.copy_from_slice(message_hash.as_ref());

        let msg_to_sign = libsecp256k1::Message::parse(&message_bytes);

        let (sig, recovery_id) = libsecp256k1::sign(&msg_to_sign, &secret_key);

        let mut signature_raw = [0u8; 65];
        signature_raw[0..32].copy_from_slice(&sig.r.b32());
        signature_raw[32..64].copy_from_slice(&sig.s.b32());
        signature_raw[64] = recovery_id.serialize();

        let evm_signature = EvmSignature::from(signature_raw);

        #[extrinsic_call]
        claim(
            RawOrigin::Signed(current_account.clone()),
            network_id,
            claim_package,
            dummy_evm_address,
            evm_signature,
        );

        assert_eq!(dummy_shares, GlobalShares::<T>::get().claimed_shares);
        assert_eq!(
            dummy_shares,
            MemberShares::<T>::get(&current_account).activated_shares
        );
        assert_eq!(
            dummy_shares,
            NetworkShares::<T>::get(&network_id).claimed_shares(),
        );
    }

    impl_benchmark_test_suite!(
        Pallet,
        crate::mock::new_test_ext(),
        crate::mock::Test,
    );
}