1#![cfg(any(feature = "runtime-benchmarks", test))]
22#![cfg_attr(test, allow(dead_code))]
27
28use crate::{
29 address::AddressMapper,
30 exec::{ExportedFunction, Key, PrecompileExt, Stack},
31 limits,
32 storage::meter::Meter,
33 transient_storage::MeterEntry,
34 vm::pvm::{PreparedCall, Runtime},
35 AccountInfo, BalanceOf, BalanceWithDust, BumpNonce, Code, CodeInfoOf, Config, ContractBlob,
36 ContractInfo, DepositLimit, Error, GasMeter, MomentOf, Origin, Pallet as Contracts,
37 PristineCode, Weight,
38};
39use alloc::{vec, vec::Vec};
40use frame_support::{storage::child, traits::fungible::Mutate};
41use frame_system::RawOrigin;
42use pallet_revive_fixtures::bench as bench_fixtures;
43use sp_core::{H160, H256, U256};
44use sp_io::hashing::keccak_256;
45use sp_runtime::traits::{Bounded, Hash};
46
47type StackExt<'a, T> = Stack<'a, T, ContractBlob<T>>;
48
49pub struct CallSetup<T: Config> {
51 contract: Contract<T>,
52 dest: T::AccountId,
53 origin: Origin<T>,
54 gas_meter: GasMeter<T>,
55 storage_meter: Meter<T>,
56 value: BalanceOf<T>,
57 data: Vec<u8>,
58 transient_storage_size: u32,
59}
60
61impl<T> Default for CallSetup<T>
62where
63 T: Config,
64 BalanceOf<T>: Into<U256> + TryFrom<U256>,
65 MomentOf<T>: Into<U256>,
66 T::Hash: frame_support::traits::IsType<H256>,
67{
68 fn default() -> Self {
69 Self::new(VmBinaryModule::dummy())
70 }
71}
72
73impl<T> CallSetup<T>
74where
75 T: Config,
76 BalanceOf<T>: Into<U256> + TryFrom<U256>,
77 MomentOf<T>: Into<U256>,
78 T::Hash: frame_support::traits::IsType<H256>,
79{
80 pub fn new(module: VmBinaryModule) -> Self {
82 let contract = Contract::<T>::new(module, vec![]).unwrap();
83 let dest = contract.account_id.clone();
84 let origin = Origin::from_account_id(contract.caller.clone());
85
86 let storage_meter = Meter::new(default_deposit_limit::<T>());
87
88 #[cfg(feature = "runtime-benchmarks")]
89 {
90 frame_benchmarking::benchmarking::add_to_whitelist(
92 frame_system::Account::<T>::hashed_key_for(&contract.account_id).into(),
93 );
94
95 frame_benchmarking::benchmarking::add_to_whitelist(
98 crate::AccountInfoOf::<T>::hashed_key_for(&T::AddressMapper::to_address(
99 &contract.account_id,
100 ))
101 .into(),
102 );
103 }
104
105 Self {
106 contract,
107 dest,
108 origin,
109 gas_meter: GasMeter::new(Weight::MAX),
110 storage_meter,
111 value: 0u32.into(),
112 data: vec![],
113 transient_storage_size: 0,
114 }
115 }
116
117 pub fn set_storage_deposit_limit(&mut self, balance: BalanceOf<T>) {
119 self.storage_meter = Meter::new(balance);
120 }
121
122 pub fn set_origin(&mut self, origin: Origin<T>) {
124 self.origin = origin;
125 }
126
127 pub fn set_balance(&mut self, value: impl Into<BalanceWithDust<BalanceOf<T>>>) {
129 self.contract.set_balance(value);
130 }
131
132 pub fn set_data(&mut self, value: Vec<u8>) {
134 self.data = value;
135 }
136
137 pub fn set_transient_storage_size(&mut self, size: u32) {
139 self.transient_storage_size = size;
140 }
141
142 pub fn data(&self) -> Vec<u8> {
144 self.data.clone()
145 }
146
147 pub fn contract(&self) -> Contract<T> {
149 self.contract.clone()
150 }
151
152 pub fn ext(&mut self) -> (StackExt<'_, T>, ContractBlob<T>) {
154 let mut ext = StackExt::bench_new_call(
155 T::AddressMapper::to_address(&self.dest),
156 self.origin.clone(),
157 &mut self.gas_meter,
158 &mut self.storage_meter,
159 self.value,
160 );
161 if self.transient_storage_size > 0 {
162 Self::with_transient_storage(&mut ext.0, self.transient_storage_size).unwrap();
163 }
164 ext
165 }
166
167 pub fn prepare_call<'a>(
169 ext: &'a mut StackExt<'a, T>,
170 module: ContractBlob<T>,
171 input: Vec<u8>,
172 aux_data_size: u32,
173 ) -> PreparedCall<'a, StackExt<'a, T>> {
174 module
175 .prepare_call(Runtime::new(ext, input), ExportedFunction::Call, aux_data_size)
176 .unwrap()
177 }
178
179 fn with_transient_storage(ext: &mut StackExt<T>, size: u32) -> Result<(), &'static str> {
181 let &MeterEntry { amount, limit } = ext.transient_storage().meter().current();
182 ext.transient_storage().meter().current_mut().limit = size;
183 for i in 1u32.. {
184 let mut key_data = i.to_le_bytes().to_vec();
185 while key_data.last() == Some(&0) {
186 key_data.pop();
187 }
188 let key = Key::try_from_var(key_data).unwrap();
189 if let Err(e) = ext.set_transient_storage(&key, Some(Vec::new()), false) {
190 ext.transient_storage().meter().current_mut().limit = limit;
192 ext.transient_storage().meter().current_mut().amount = amount;
193 if e == Error::<T>::OutOfTransientStorage.into() {
194 break;
195 } else {
196 return Err("Initialization of the transient storage failed");
197 }
198 }
199 }
200 Ok(())
201 }
202}
203
204pub fn default_deposit_limit<T: Config>() -> BalanceOf<T> {
206 <BalanceOf<T>>::max_value()
207}
208
209pub fn caller_funding<T: Config>() -> BalanceOf<T> {
211 BalanceOf::<T>::max_value() / 10_000u32.into()
214}
215
216#[derive(Clone)]
218pub struct Contract<T: Config> {
219 pub caller: T::AccountId,
220 pub account_id: T::AccountId,
221 pub address: H160,
222}
223
224impl<T> Contract<T>
225where
226 T: Config,
227 BalanceOf<T>: Into<U256> + TryFrom<U256>,
228 MomentOf<T>: Into<U256>,
229 T::Hash: frame_support::traits::IsType<H256>,
230{
231 pub fn new(module: VmBinaryModule, data: Vec<u8>) -> Result<Contract<T>, &'static str> {
233 let caller = T::AddressMapper::to_fallback_account_id(&crate::test_utils::ALICE_ADDR);
234 Self::with_caller(caller, module, data)
235 }
236
237 #[cfg(feature = "runtime-benchmarks")]
239 pub fn with_index(
240 index: u32,
241 module: VmBinaryModule,
242 data: Vec<u8>,
243 ) -> Result<Contract<T>, &'static str> {
244 Self::with_caller(frame_benchmarking::account("instantiator", index, 0), module, data)
245 }
246
247 pub fn with_caller(
249 caller: T::AccountId,
250 module: VmBinaryModule,
251 data: Vec<u8>,
252 ) -> Result<Contract<T>, &'static str> {
253 T::Currency::set_balance(&caller, caller_funding::<T>());
254 let salt = Some([0xffu8; 32]);
255 let origin: T::RuntimeOrigin = RawOrigin::Signed(caller.clone()).into();
256
257 Contracts::<T>::map_account(origin.clone()).ok();
259
260 #[cfg(feature = "runtime-benchmarks")]
261 frame_benchmarking::benchmarking::add_to_whitelist(
262 frame_system::Account::<T>::hashed_key_for(&caller).into(),
263 );
264
265 let outcome = Contracts::<T>::bare_instantiate(
266 origin,
267 U256::zero(),
268 Weight::MAX,
269 DepositLimit::Balance(default_deposit_limit::<T>()),
270 Code::Upload(module.code),
271 data,
272 salt,
273 BumpNonce::Yes,
274 );
275
276 let address = outcome.result?.addr;
277 let account_id = T::AddressMapper::to_fallback_account_id(&address);
278 let result = Contract { caller, address, account_id };
279
280 AccountInfo::<T>::insert_contract(&address, result.info()?);
281 Ok(result)
282 }
283
284 pub fn with_storage(
286 code: VmBinaryModule,
287 stor_num: u32,
288 stor_size: u32,
289 ) -> Result<Self, &'static str> {
290 let contract = Contract::<T>::new(code, vec![])?;
291 let storage_items = (0..stor_num)
292 .map(|i| {
293 let hash = T::Hashing::hash_of(&i)
294 .as_ref()
295 .try_into()
296 .map_err(|_| "Hash too big for storage key")?;
297 Ok((hash, vec![42u8; stor_size as usize]))
298 })
299 .collect::<Result<Vec<_>, &'static str>>()?;
300 contract.store(&storage_items)?;
301 Ok(contract)
302 }
303
304 pub fn store(&self, items: &Vec<([u8; 32], Vec<u8>)>) -> Result<(), &'static str> {
306 let info = self.info()?;
307 for item in items {
308 info.write(&Key::Fix(item.0), Some(item.1.clone()), None, false)
309 .map_err(|_| "Failed to write storage to restoration dest")?;
310 }
311
312 AccountInfo::<T>::insert_contract(&self.address, info);
313 Ok(())
314 }
315
316 pub fn with_unbalanced_storage_trie(
318 code: VmBinaryModule,
319 key: &[u8],
320 ) -> Result<Self, &'static str> {
321 const UNBALANCED_TRIE_LAYERS: u32 = 20;
323
324 if (key.len() as u32) < UNBALANCED_TRIE_LAYERS.div_ceil(2) {
325 return Err("Key size too small to create the specified trie");
326 }
327
328 let value = vec![16u8; limits::PAYLOAD_BYTES as usize];
329 let contract = Contract::<T>::new(code, vec![])?;
330 let info = contract.info()?;
331 let child_trie_info = info.child_trie_info();
332 child::put_raw(&child_trie_info, &key, &value);
333 for l in 0..UNBALANCED_TRIE_LAYERS {
334 let pos = l as usize / 2;
335 let mut key_new = key.to_vec();
336 for i in 0u8..16 {
337 key_new[pos] = if l % 2 == 0 {
338 (key_new[pos] & 0xF0) | i
339 } else {
340 (key_new[pos] & 0x0F) | (i << 4)
341 };
342
343 if key == &key_new {
344 continue;
345 }
346 child::put_raw(&child_trie_info, &key_new, &value);
347 }
348 }
349 Ok(contract)
350 }
351
352 pub fn address_info(addr: &T::AccountId) -> Result<ContractInfo<T>, &'static str> {
354 <AccountInfo<T>>::load_contract(&T::AddressMapper::to_address(addr))
355 .ok_or("Expected contract to exist at this point.")
356 }
357
358 pub fn info(&self) -> Result<ContractInfo<T>, &'static str> {
360 Self::address_info(&self.account_id)
361 }
362
363 pub fn set_balance(&self, value: impl Into<BalanceWithDust<BalanceOf<T>>>) {
365 let (value, dust) = value.into().deconstruct();
366 T::Currency::set_balance(&self.account_id, value);
367 crate::AccountInfoOf::<T>::mutate(&self.address, |account| {
368 account.as_mut().map(|a| a.dust = dust);
369 });
370 }
371
372 pub fn code_exists(hash: &sp_core::H256) -> bool {
374 <PristineCode<T>>::contains_key(hash) && <CodeInfoOf<T>>::contains_key(&hash)
375 }
376
377 pub fn code_removed(hash: &sp_core::H256) -> bool {
379 !<PristineCode<T>>::contains_key(hash) && !<CodeInfoOf<T>>::contains_key(&hash)
380 }
381}
382
383#[derive(Clone)]
385pub struct VmBinaryModule {
386 pub code: Vec<u8>,
387 pub hash: H256,
388}
389
390impl VmBinaryModule {
391 pub fn dummy() -> Self {
393 Self::new(bench_fixtures::DUMMY.to_vec())
394 }
395
396 fn new(code: Vec<u8>) -> Self {
397 let hash = keccak_256(&code);
398 Self { code, hash: H256(hash) }
399 }
400}
401
402#[cfg(feature = "runtime-benchmarks")]
403impl VmBinaryModule {
404 pub fn dummy_unique(replace_with: u32) -> Self {
406 Self::new(bench_fixtures::dummy_unique(replace_with))
407 }
408
409 pub fn sized(size: u32) -> Self {
414 Self::with_num_instructions(size / 3)
417 }
418
419 pub fn evm_sized(size: u32) -> Self {
421 use revm::bytecode::opcode::STOP;
422 let code = vec![STOP; size as usize];
423 Self::new(code)
424 }
425
426 pub fn with_num_instructions(num_instructions: u32) -> Self {
437 use alloc::{fmt::Write, string::ToString};
438 let mut text = "
439 pub @deploy:
440 ret
441 pub @call:
442 "
443 .to_string();
444 for i in 0..num_instructions {
445 match i {
446 0 => writeln!(text, "ecalli {}", crate::SENTINEL).unwrap(),
449 i if i % (limits::code::BASIC_BLOCK_SIZE - 1) == 0 =>
450 text.push_str("fallthrough\n"),
451 _ => text.push_str("a0 = a1 + a2\n"),
452 }
453 }
454 text.push_str("ret\n");
455 let code = polkavm_common::assembler::assemble(&text).unwrap();
456 Self::new(code)
457 }
458
459 pub fn noop() -> Self {
461 Self::new(bench_fixtures::NOOP.to_vec())
462 }
463
464 pub fn instr(do_load: bool) -> Self {
466 let load = match do_load {
467 false => "",
468 true => "a0 = u64 [a0]",
469 };
470 let text = alloc::format!(
471 "
472 pub @deploy:
473 ret
474 pub @call:
475 @loop:
476 jump @done if t0 == a1
477 {load}
478 t0 = t0 + 1
479 jump @loop
480 @done:
481 ret
482 "
483 );
484 let code = polkavm_common::assembler::assemble(&text).unwrap();
485 Self::new(code)
486 }
487
488 pub fn evm_noop(size: u32) -> Self {
490 use revm::bytecode::opcode::JUMPDEST;
491
492 let code = vec![JUMPDEST; size as usize];
493 Self::new(code)
494 }
495}