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