1use super::*;
21use crate::list::List;
22use alloc::{vec, vec::Vec};
23use frame_benchmarking::v1::{
24 account, benchmarks_instance_pallet, whitelist_account, whitelisted_caller,
25};
26use frame_election_provider_support::ScoreProvider;
27use frame_support::{assert_ok, traits::Get};
28use frame_system::RawOrigin as SystemOrigin;
29use sp_runtime::traits::One;
30
31benchmarks_instance_pallet! {
32 #[extra]
34 iter {
35 let n = 100;
36
37 List::<T, _>::unsafe_clear();
39
40 let bag_thresh = T::BagThresholds::get()[0];
42 let second_bag_thresh = T::BagThresholds::get()[1];
43
44
45 for i in 0..n/2 {
46 let node: T::AccountId = account("node", i, 0);
47 assert_ok!(List::<T, _>::insert(node.clone(), bag_thresh - One::one()));
48 }
49 for i in 0..n/2 {
50 let node: T::AccountId = account("node", i, 1);
51 assert_ok!(List::<T, _>::insert(node.clone(), bag_thresh + One::one()));
52 }
53 assert_eq!(
54 List::<T, _>::get_bags().into_iter().map(|(bag, nodes)| (bag, nodes.len())).collect::<Vec<_>>(),
55 vec![
56 (bag_thresh, (n / 2) as usize),
57 (second_bag_thresh, (n / 2) as usize),
58 ]
59 );
60 }: {
61 let voters = <Pallet<T, _> as SortedListProvider<T::AccountId>>::iter();
62 let len = voters.collect::<Vec<_>>().len();
63 assert_eq!(len as u32, n,"len is {}, expected {}", len, n);
64 }
65
66 #[extra]
68 iter_take {
69 let n = 100;
70
71 List::<T, _>::unsafe_clear();
73
74 let bag_thresh = T::BagThresholds::get()[0];
76 let second_bag_thresh = T::BagThresholds::get()[1];
77
78
79 for i in 0..n/2 {
80 let node: T::AccountId = account("node", i, 0);
81 assert_ok!(List::<T, _>::insert(node.clone(), bag_thresh - One::one()));
82 }
83 for i in 0..n/2 {
84 let node: T::AccountId = account("node", i, 1);
85 assert_ok!(List::<T, _>::insert(node.clone(), bag_thresh + One::one()));
86 }
87 assert_eq!(
88 List::<T, _>::get_bags().into_iter().map(|(bag, nodes)| (bag, nodes.len())).collect::<Vec<_>>(),
89 vec![
90 (bag_thresh, (n / 2) as usize),
91 (second_bag_thresh, (n / 2) as usize),
92 ]
93 );
94 }: {
95 let voters = <Pallet<T, _> as SortedListProvider<T::AccountId>>::iter().take(n as usize / 4 );
97 let len = voters.collect::<Vec<_>>().len();
98 assert_eq!(len as u32, n / 4,"len is {}, expected {}", len, n / 4);
99 }
100
101 #[extra]
102 iter_next {
103 let n = 100;
104
105 List::<T, _>::unsafe_clear();
107
108 let bag_thresh = T::BagThresholds::get()[0];
110 let second_bag_thresh = T::BagThresholds::get()[1];
111
112
113 for i in 0..n/2 {
114 let node: T::AccountId = account("node", i, 0);
115 assert_ok!(List::<T, _>::insert(node.clone(), bag_thresh - One::one()));
116 }
117 for i in 0..n/2 {
118 let node: T::AccountId = account("node", i, 1);
119 assert_ok!(List::<T, _>::insert(node.clone(), bag_thresh + One::one()));
120 }
121 assert_eq!(
122 List::<T, _>::get_bags().into_iter().map(|(bag, nodes)| (bag, nodes.len())).collect::<Vec<_>>(),
123 vec![
124 (bag_thresh, (n / 2) as usize),
125 (second_bag_thresh, (n / 2) as usize),
126 ]
127 );
128 }: {
129 let mut iter_var = <Pallet<T, _> as SortedListProvider<T::AccountId>>::iter();
131 let mut voters = Vec::<T::AccountId>::with_capacity((n/4) as usize);
132 for _ in 0..(n/4) {
133 let next = iter_var.next().unwrap();
134 voters.push(next);
135 }
136
137 let len = voters.len();
138 assert_eq!(len as u32, n / 4,"len is {}, expected {}", len, n / 4);
139 }
140
141 #[extra]
142 iter_from {
143 let n = 100;
144
145 List::<T, _>::unsafe_clear();
147
148 let bag_thresh = T::BagThresholds::get()[0];
150
151 for i in 0..n/4 {
152 let node: T::AccountId = account("node", i, 0);
153 assert_ok!(List::<T, _>::insert(node.clone(), bag_thresh - One::one()));
154 }
155 for i in 0..n/4 {
156 let node: T::AccountId = account("node", i, 1);
157 assert_ok!(List::<T, _>::insert(node.clone(), bag_thresh + One::one()));
158 }
159
160 let bag_thresh = T::BagThresholds::get()[2];
161
162 for i in 0..n/4 {
163 let node: T::AccountId = account("node", i, 2);
164 assert_ok!(List::<T, _>::insert(node.clone(), bag_thresh - One::one()));
165 }
166
167 for i in 0..n/4 {
168 let node: T::AccountId = account("node", i, 3);
169 assert_ok!(List::<T, _>::insert(node.clone(), bag_thresh + One::one()));
170 }
171
172 assert_eq!(
173 List::<T, _>::get_bags().into_iter().map(|(bag, nodes)| (bag, nodes.len())).collect::<Vec<_>>(),
174 vec![
175 (T::BagThresholds::get()[0], (n / 4) as usize),
176 (T::BagThresholds::get()[1], (n / 4) as usize),
177 (T::BagThresholds::get()[2], (n / 4) as usize),
178 (T::BagThresholds::get()[3], (n / 4) as usize),
179 ]
180 );
181
182 let from: T::AccountId = account("node", 0, 2);
184 }: {
185 let voters = <Pallet<T, _> as SortedListProvider<T::AccountId>>::iter_from(&from).unwrap();
186 let len = voters.collect::<Vec<_>>().len();
187 assert_eq!(len as u32, 74,"len is {}, expected {}", len, 74);
188 }
189
190
191 rebag_non_terminal {
192 List::<T, _>::unsafe_clear();
203
204 let origin_bag_thresh = T::BagThresholds::get()[0];
206 let dest_bag_thresh = T::BagThresholds::get()[1];
207
208 let origin_head: T::AccountId = account("origin_head", 0, 0);
210 assert_ok!(List::<T, _>::insert(origin_head.clone(), origin_bag_thresh));
211
212 let origin_middle: T::AccountId = account("origin_middle", 0, 0); assert_ok!(List::<T, _>::insert(origin_middle.clone(), origin_bag_thresh));
214
215 let origin_tail: T::AccountId = account("origin_tail", 0, 0);
216 assert_ok!(List::<T, _>::insert(origin_tail.clone(), origin_bag_thresh));
217
218 let dest_head: T::AccountId = account("dest_head", 0, 0);
220 assert_ok!(List::<T, _>::insert(dest_head.clone(), dest_bag_thresh));
221
222 let origin_middle_lookup = T::Lookup::unlookup(origin_middle.clone());
223
224 assert_eq!(
226 List::<T, _>::get_bags(),
227 vec![
228 (origin_bag_thresh, vec![origin_head.clone(), origin_middle.clone(), origin_tail.clone()]),
229 (dest_bag_thresh, vec![dest_head.clone()])
230 ]
231 );
232
233 let caller = whitelisted_caller();
234 T::ScoreProvider::set_score_of(&origin_middle, dest_bag_thresh);
236 }: rebag(SystemOrigin::Signed(caller), origin_middle_lookup.clone())
237 verify {
238 assert_eq!(
240 List::<T, _>::get_bags(),
241 vec![
242 (
243 origin_bag_thresh,
244 vec![origin_head, origin_tail],
245 ),
246 (
247 dest_bag_thresh,
248 vec![dest_head, origin_middle],
249 )
250 ]
251 );
252 }
253
254 rebag_terminal {
255 List::<T, I>::unsafe_clear();
265
266 let origin_bag_thresh = T::BagThresholds::get()[0];
268 let dest_bag_thresh = T::BagThresholds::get()[1];
269
270 let origin_head: T::AccountId = account("origin_head", 0, 0);
272 assert_ok!(List::<T, _>::insert(origin_head.clone(), origin_bag_thresh));
273
274 let origin_tail: T::AccountId = account("origin_tail", 0, 0); assert_ok!(List::<T, _>::insert(origin_tail.clone(), origin_bag_thresh));
276
277 let dest_head: T::AccountId = account("dest_head", 0, 0);
279 assert_ok!(List::<T, _>::insert(dest_head.clone(), dest_bag_thresh));
280
281 let origin_tail_lookup = T::Lookup::unlookup(origin_tail.clone());
282
283 assert_eq!(
285 List::<T, _>::get_bags(),
286 vec![
287 (origin_bag_thresh, vec![origin_head.clone(), origin_tail.clone()]),
288 (dest_bag_thresh, vec![dest_head.clone()])
289 ]
290 );
291
292 let caller = whitelisted_caller();
293 T::ScoreProvider::set_score_of(&origin_tail, dest_bag_thresh);
295 }: rebag(SystemOrigin::Signed(caller), origin_tail_lookup.clone())
296 verify {
297 assert_eq!(
299 List::<T, _>::get_bags(),
300 vec![
301 (origin_bag_thresh, vec![origin_head.clone()]),
302 (dest_bag_thresh, vec![dest_head.clone(), origin_tail])
303 ]
304 );
305 }
306
307 put_in_front_of {
308 List::<T, I>::unsafe_clear();
316
317 let bag_thresh = T::BagThresholds::get()[0];
318
319 let lighter: T::AccountId = account("lighter", 0, 0);
321 assert_ok!(List::<T, _>::insert(lighter.clone(), bag_thresh));
322
323 let heavier_prev: T::AccountId = account("heavier_prev", 0, 0);
324 assert_ok!(List::<T, _>::insert(heavier_prev.clone(), bag_thresh));
325
326 let heavier: T::AccountId = account("heavier", 0, 0);
327 assert_ok!(List::<T, _>::insert(heavier.clone(), bag_thresh));
328
329 let heavier_next: T::AccountId = account("heavier_next", 0, 0);
330 assert_ok!(List::<T, _>::insert(heavier_next.clone(), bag_thresh));
331
332 T::ScoreProvider::set_score_of(&lighter, bag_thresh - One::one());
333 T::ScoreProvider::set_score_of(&heavier, bag_thresh);
334
335 let lighter_lookup = T::Lookup::unlookup(lighter.clone());
336
337 assert_eq!(
338 List::<T, _>::iter().map(|n| n.id().clone()).collect::<Vec<_>>(),
339 vec![lighter.clone(), heavier_prev.clone(), heavier.clone(), heavier_next.clone()]
340 );
341
342 whitelist_account!(heavier);
343 }: _(SystemOrigin::Signed(heavier.clone()), lighter_lookup.clone())
344 verify {
345 assert_eq!(
346 List::<T, _>::iter().map(|n| n.id().clone()).collect::<Vec<_>>(),
347 vec![heavier, lighter, heavier_prev, heavier_next]
348 )
349 }
350
351 on_idle {
352 List::<T, _>::unsafe_clear();
359
360 let bag_thresh = T::BagThresholds::get();
361 let low = bag_thresh[0];
362 let mid = bag_thresh[1];
363 let high = bag_thresh[2];
364
365 let rebag_budget = <T as Config<I>>::MaxAutoRebagPerBlock::get();
366
367 let pending_count = rebag_budget / 3; let regular_count = rebag_budget + 5;
370
371 for i in 0..regular_count {
373 let node: T::AccountId = account("regular_node", i, 0);
374 let score = match i % 3 {
375 0 => low - One::one(),
376 1 => mid - One::one(),
377 _ => high - One::one(),
378 };
379 assert_ok!(List::<T, _>::insert(node.clone(), score));
380 }
381
382 for i in (0..regular_count).step_by(4) {
384 let node: T::AccountId = account("regular_node", i, 0);
385 let _ = List::<T, _>::remove(&node); }
387
388 <Pallet<T, I>>::lock();
390
391 for i in 0..pending_count {
393 let pending_node: T::AccountId = account("pending_node", i, 0);
394 let pending_score = match i % 3 {
395 0 => mid,
396 1 => high,
397 _ => high + high,
398 };
399
400 if i % 7 != 0 {
402 T::ScoreProvider::set_score_of(&pending_node, pending_score);
403 }
404
405 let _ = <Pallet<T, I> as SortedListProvider<T::AccountId>>::on_insert(
406 pending_node, pending_score
407 );
408 }
409
410 <Pallet<T, I>>::unlock();
411
412 for i in 0..regular_count {
414 let node: T::AccountId = account("regular_node", i, 0);
415 let new_score = match i % 3 {
416 0 => mid,
417 1 => high,
418 _ => high + high, };
420 T::ScoreProvider::set_score_of(&node, new_score);
421 }
422
423 assert_eq!(
424 PendingRebag::<T, I>::count(),
425 pending_count,
426 "Expected exactly {} pending rebag entries",
427 pending_count
428 );
429 assert!(List::<T, _>::get_bags().len() >= 2);
431 }
432 : {
433 use frame_support::traits::Hooks;
434 <Pallet<T, I> as Hooks<_>>::on_idle(Default::default(), Weight::MAX);
435 }
436 verify {
437 assert_eq!(PendingRebag::<T, I>::count(), 0, "All pending rebag entries should be processed");
441
442 let total_rebagged: usize = List::<T, _>::get_bags()
444 .iter()
445 .filter(|(b, _)| *b > T::BagThresholds::get()[0])
446 .map(|(_, nodes)| nodes.len())
447 .sum();
448
449 let expected = <T as Config<I>>::MaxAutoRebagPerBlock::get() as usize;
450 assert_eq!(total_rebagged, expected, "Expected exactly {:?} rebagged nodes, found {:?}", expected, total_rebagged);
451 }
452
453 impl_benchmark_test_suite!(
454 Pallet,
455 mock::ExtBuilder::default().skip_genesis_ids().build(),
456 mock::Runtime
457 );
458}