1#![cfg(feature = "runtime-benchmarks")]
21
22use frame_benchmarking::{v2::*, BenchmarkError};
23use frame_support::assert_ok;
24use frame_system::{pallet_prelude::BlockNumberFor, RawOrigin};
25use sp_runtime::traits::{Bounded, CheckedDiv, CheckedMul};
26
27use crate::*;
28
29const SEED: u32 = 0;
30
31type BalanceOf<T> =
32 <<T as Config>::Currency as Currency<<T as frame_system::Config>::AccountId>>::Balance;
33
34fn add_locks<T: Config>(who: &T::AccountId, n: u8) {
35 for id in 0..n {
36 let lock_id = [id; 8];
37 let locked = 256_u32;
38 let reasons = WithdrawReasons::TRANSFER | WithdrawReasons::RESERVE;
39 T::Currency::set_lock(lock_id, who, locked.into(), reasons);
40 }
41}
42
43fn add_vesting_schedules<T: Config>(
44 target: &T::AccountId,
45 n: u32,
46) -> Result<BalanceOf<T>, &'static str> {
47 let min_transfer = T::MinVestedTransfer::get();
48 let locked = min_transfer.checked_mul(&20_u32.into()).unwrap();
49 let per_block = min_transfer;
51 let starting_block = 1_u32;
52
53 let source = account("source", 0, SEED);
54 T::Currency::make_free_balance_be(&source, BalanceOf::<T>::max_value());
55
56 T::BlockNumberProvider::set_block_number(BlockNumberFor::<T>::zero());
57
58 let mut total_locked: BalanceOf<T> = Zero::zero();
59 for _ in 0..n {
60 total_locked += locked;
61
62 let schedule = VestingInfo::new(locked, per_block, starting_block.into());
63 assert_ok!(Pallet::<T>::do_vested_transfer(&source, target, schedule));
64
65 T::Currency::make_free_balance_be(&source, BalanceOf::<T>::max_value());
67 }
68
69 Ok(total_locked)
70}
71
72#[benchmarks]
73mod benchmarks {
74 use super::*;
75
76 #[benchmark]
77 fn vest_locked(
78 l: Linear<0, { MaxLocksOf::<T>::get() - 1 }>,
79 s: Linear<1, T::MAX_VESTING_SCHEDULES>,
80 ) -> Result<(), BenchmarkError> {
81 let caller = whitelisted_caller();
82 T::Currency::make_free_balance_be(&caller, T::Currency::minimum_balance());
83
84 add_locks::<T>(&caller, l as u8);
85 let expected_balance = add_vesting_schedules::<T>(&caller, s)?;
86
87 assert_eq!(frame_system::Pallet::<T>::block_number(), BlockNumberFor::<T>::zero());
89 assert_eq!(
90 Pallet::<T>::vesting_balance(&caller),
91 Some(expected_balance),
92 "Vesting schedule not added",
93 );
94
95 #[extrinsic_call]
96 vest(RawOrigin::Signed(caller.clone()));
97
98 assert_eq!(
100 Pallet::<T>::vesting_balance(&caller),
101 Some(expected_balance),
102 "Vesting schedule was removed",
103 );
104
105 Ok(())
106 }
107
108 #[benchmark]
109 fn vest_unlocked(
110 l: Linear<0, { MaxLocksOf::<T>::get() - 1 }>,
111 s: Linear<1, T::MAX_VESTING_SCHEDULES>,
112 ) -> Result<(), BenchmarkError> {
113 let caller = whitelisted_caller();
114 T::Currency::make_free_balance_be(&caller, T::Currency::minimum_balance());
115
116 add_locks::<T>(&caller, l as u8);
117 add_vesting_schedules::<T>(&caller, s)?;
118
119 T::BlockNumberProvider::set_block_number(21_u32.into());
121 assert_eq!(
122 Pallet::<T>::vesting_balance(&caller),
123 Some(BalanceOf::<T>::zero()),
124 "Vesting schedule still active",
125 );
126
127 #[extrinsic_call]
128 vest(RawOrigin::Signed(caller.clone()));
129
130 assert_eq!(Pallet::<T>::vesting_balance(&caller), None, "Vesting schedule was not removed",);
132
133 Ok(())
134 }
135
136 #[benchmark]
137 fn vest_other_locked(
138 l: Linear<0, { MaxLocksOf::<T>::get() - 1 }>,
139 s: Linear<1, T::MAX_VESTING_SCHEDULES>,
140 ) -> Result<(), BenchmarkError> {
141 let other = account::<T::AccountId>("other", 0, SEED);
142 let other_lookup = T::Lookup::unlookup(other.clone());
143
144 T::Currency::make_free_balance_be(&other, T::Currency::minimum_balance());
145 add_locks::<T>(&other, l as u8);
146 let expected_balance = add_vesting_schedules::<T>(&other, s)?;
147
148 assert_eq!(frame_system::Pallet::<T>::block_number(), BlockNumberFor::<T>::zero());
150 assert_eq!(
151 Pallet::<T>::vesting_balance(&other),
152 Some(expected_balance),
153 "Vesting schedule not added",
154 );
155
156 let caller = whitelisted_caller::<T::AccountId>();
157
158 #[extrinsic_call]
159 vest_other(RawOrigin::Signed(caller.clone()), other_lookup);
160
161 assert_eq!(
163 Pallet::<T>::vesting_balance(&other),
164 Some(expected_balance),
165 "Vesting schedule was removed",
166 );
167
168 Ok(())
169 }
170
171 #[benchmark]
172 fn vest_other_unlocked(
173 l: Linear<0, { MaxLocksOf::<T>::get() - 1 }>,
174 s: Linear<1, { T::MAX_VESTING_SCHEDULES }>,
175 ) -> Result<(), BenchmarkError> {
176 let other = account::<T::AccountId>("other", 0, SEED);
177 let other_lookup = T::Lookup::unlookup(other.clone());
178
179 T::Currency::make_free_balance_be(&other, T::Currency::minimum_balance());
180 add_locks::<T>(&other, l as u8);
181 add_vesting_schedules::<T>(&other, s)?;
182 T::BlockNumberProvider::set_block_number(21_u32.into());
184
185 assert_eq!(
186 Pallet::<T>::vesting_balance(&other),
187 Some(BalanceOf::<T>::zero()),
188 "Vesting schedule still active",
189 );
190
191 let caller = whitelisted_caller::<T::AccountId>();
192
193 #[extrinsic_call]
194 vest_other(RawOrigin::Signed(caller.clone()), other_lookup);
195
196 assert_eq!(Pallet::<T>::vesting_balance(&other), None, "Vesting schedule was not removed",);
198
199 Ok(())
200 }
201
202 #[benchmark]
203 fn vested_transfer(
204 l: Linear<0, { MaxLocksOf::<T>::get() - 1 }>,
205 s: Linear<0, { T::MAX_VESTING_SCHEDULES - 1 }>,
206 ) -> Result<(), BenchmarkError> {
207 let caller = whitelisted_caller();
208 T::Currency::make_free_balance_be(&caller, BalanceOf::<T>::max_value());
209
210 let target = account::<T::AccountId>("target", 0, SEED);
211 let target_lookup = T::Lookup::unlookup(target.clone());
212 T::Currency::make_free_balance_be(&target, T::Currency::minimum_balance());
214 add_locks::<T>(&target, l as u8);
215 let orig_balance = T::Currency::free_balance(&target);
217 let mut expected_balance = add_vesting_schedules::<T>(&target, s)?;
218
219 let transfer_amount = T::MinVestedTransfer::get();
220 let per_block = transfer_amount.checked_div(&20_u32.into()).unwrap();
221 expected_balance += transfer_amount;
222
223 let vesting_schedule = VestingInfo::new(transfer_amount, per_block, 1_u32.into());
224
225 #[extrinsic_call]
226 _(RawOrigin::Signed(caller.clone()), target_lookup, vesting_schedule);
227
228 assert_eq!(
229 orig_balance + expected_balance,
230 T::Currency::free_balance(&target),
231 "Transfer didn't happen",
232 );
233 assert_eq!(
234 Pallet::<T>::vesting_balance(&target),
235 Some(expected_balance),
236 "Lock not correctly updated",
237 );
238
239 Ok(())
240 }
241
242 #[benchmark]
243 fn force_vested_transfer(
244 l: Linear<0, { MaxLocksOf::<T>::get() - 1 }>,
245 s: Linear<0, { T::MAX_VESTING_SCHEDULES - 1 }>,
246 ) -> Result<(), BenchmarkError> {
247 let source = account::<T::AccountId>("source", 0, SEED);
248 let source_lookup = T::Lookup::unlookup(source.clone());
249 T::Currency::make_free_balance_be(&source, BalanceOf::<T>::max_value());
250
251 let target = account::<T::AccountId>("target", 0, SEED);
252 let target_lookup = T::Lookup::unlookup(target.clone());
253 T::Currency::make_free_balance_be(&target, T::Currency::minimum_balance());
255 add_locks::<T>(&target, l as u8);
256 let orig_balance = T::Currency::free_balance(&target);
258 let mut expected_balance = add_vesting_schedules::<T>(&target, s)?;
259
260 let transfer_amount = T::MinVestedTransfer::get();
261 let per_block = transfer_amount.checked_div(&20_u32.into()).unwrap();
262 expected_balance += transfer_amount;
263
264 let vesting_schedule = VestingInfo::new(transfer_amount, per_block, 1_u32.into());
265
266 #[extrinsic_call]
267 _(RawOrigin::Root, source_lookup, target_lookup, vesting_schedule);
268
269 assert_eq!(
270 orig_balance + expected_balance,
271 T::Currency::free_balance(&target),
272 "Transfer didn't happen",
273 );
274 assert_eq!(
275 Pallet::<T>::vesting_balance(&target),
276 Some(expected_balance),
277 "Lock not correctly updated",
278 );
279
280 Ok(())
281 }
282
283 #[benchmark]
284 fn not_unlocking_merge_schedules(
285 l: Linear<0, { MaxLocksOf::<T>::get() - 1 }>,
286 s: Linear<2, { T::MAX_VESTING_SCHEDULES }>,
287 ) -> Result<(), BenchmarkError> {
288 let caller = whitelisted_caller::<T::AccountId>();
289 T::Currency::make_free_balance_be(&caller, T::Currency::minimum_balance());
291 add_locks::<T>(&caller, l as u8);
292 let expected_balance = add_vesting_schedules::<T>(&caller, s)?;
294
295 assert_eq!(frame_system::Pallet::<T>::block_number(), BlockNumberFor::<T>::zero());
297 assert_eq!(
298 Pallet::<T>::vesting_balance(&caller),
299 Some(expected_balance),
300 "Vesting balance should equal sum locked of all schedules",
301 );
302 assert_eq!(
303 Vesting::<T>::get(&caller).unwrap().len(),
304 s as usize,
305 "There should be exactly max vesting schedules"
306 );
307
308 #[extrinsic_call]
309 merge_schedules(RawOrigin::Signed(caller.clone()), 0, s - 1);
310
311 let expected_schedule = VestingInfo::new(
312 T::MinVestedTransfer::get() * 20_u32.into() * 2_u32.into(),
313 T::MinVestedTransfer::get() * 2_u32.into(),
314 1_u32.into(),
315 );
316 let expected_index = (s - 2) as usize;
317 assert_eq!(Vesting::<T>::get(&caller).unwrap()[expected_index], expected_schedule);
318 assert_eq!(
319 Pallet::<T>::vesting_balance(&caller),
320 Some(expected_balance),
321 "Vesting balance should equal total locked of all schedules",
322 );
323 assert_eq!(
324 Vesting::<T>::get(&caller).unwrap().len(),
325 (s - 1) as usize,
326 "Schedule count should reduce by 1"
327 );
328
329 Ok(())
330 }
331
332 #[benchmark]
333 fn unlocking_merge_schedules(
334 l: Linear<0, { MaxLocksOf::<T>::get() - 1 }>,
335 s: Linear<2, { T::MAX_VESTING_SCHEDULES }>,
336 ) -> Result<(), BenchmarkError> {
337 let test_dest: T::AccountId = account("test_dest", 0, SEED);
339
340 let caller = whitelisted_caller::<T::AccountId>();
341 T::Currency::make_free_balance_be(&caller, T::Currency::minimum_balance());
343 add_locks::<T>(&caller, l as u8);
344 let total_transferred = add_vesting_schedules::<T>(&caller, s)?;
346
347 T::BlockNumberProvider::set_block_number(11_u32.into());
350 let expected_balance = total_transferred / 2_u32.into();
353 assert_eq!(
354 Pallet::<T>::vesting_balance(&caller),
355 Some(expected_balance),
356 "Vesting balance should reflect that we are half way through all schedules duration",
357 );
358 assert_eq!(
359 Vesting::<T>::get(&caller).unwrap().len(),
360 s as usize,
361 "There should be exactly max vesting schedules"
362 );
363 assert!(T::Currency::transfer(
365 &caller,
366 &test_dest,
367 expected_balance,
368 ExistenceRequirement::AllowDeath
369 )
370 .is_err());
371
372 #[extrinsic_call]
373 merge_schedules(RawOrigin::Signed(caller.clone()), 0, s - 1);
374
375 let expected_schedule = VestingInfo::new(
376 T::MinVestedTransfer::get() * 2_u32.into() * 10_u32.into(),
377 T::MinVestedTransfer::get() * 2_u32.into(),
378 11_u32.into(),
379 );
380 let expected_index = (s - 2) as usize;
381 assert_eq!(
382 Vesting::<T>::get(&caller).unwrap()[expected_index],
383 expected_schedule,
384 "New schedule is properly created and placed"
385 );
386 assert_eq!(
387 Pallet::<T>::vesting_balance(&caller),
388 Some(expected_balance),
389 "Vesting balance should equal half total locked of all schedules",
390 );
391 assert_eq!(
392 Vesting::<T>::get(&caller).unwrap().len(),
393 (s - 1) as usize,
394 "Schedule count should reduce by 1"
395 );
396 assert_ok!(T::Currency::transfer(
398 &caller,
399 &test_dest,
400 expected_balance,
401 ExistenceRequirement::AllowDeath
402 ));
403
404 Ok(())
405 }
406
407 #[benchmark]
408 fn force_remove_vesting_schedule(
409 l: Linear<0, { MaxLocksOf::<T>::get() - 1 }>,
410 s: Linear<2, { T::MAX_VESTING_SCHEDULES }>,
411 ) -> Result<(), BenchmarkError> {
412 let source = account::<T::AccountId>("source", 0, SEED);
413 T::Currency::make_free_balance_be(&source, BalanceOf::<T>::max_value());
414
415 let target = account::<T::AccountId>("target", 0, SEED);
416 let target_lookup = T::Lookup::unlookup(target.clone());
417 T::Currency::make_free_balance_be(&target, T::Currency::minimum_balance());
418
419 add_locks::<T>(&target, l as u8);
421 add_vesting_schedules::<T>(&target, s)?;
422
423 let schedule_index = s - 1;
425
426 #[extrinsic_call]
427 _(RawOrigin::Root, target_lookup, schedule_index);
428
429 assert_eq!(
430 Vesting::<T>::get(&target).unwrap().len(),
431 schedule_index as usize,
432 "Schedule count should reduce by 1"
433 );
434
435 Ok(())
436 }
437
438 impl_benchmark_test_suite! {
439 Pallet,
440 mock::ExtBuilder::default().existential_deposit(256).build(),
441 mock::Test
442 }
443}