1use crate::{
26 Candidate, CandidatePtr, Edge, ExtendedBalance, IdentifierT, Support, SupportMap, Supports,
27 VoteWeight, Voter,
28};
29use alloc::{collections::btree_map::BTreeMap, rc::Rc, vec::Vec};
30use sp_arithmetic::{traits::Zero, Perbill};
31type Threshold = ExtendedBalance;
35
36pub fn standard_threshold(
49 committee_size: usize,
50 weights: impl IntoIterator<Item = ExtendedBalance>,
51) -> Threshold {
52 weights
53 .into_iter()
54 .fold(Threshold::zero(), |acc, elem| acc.saturating_add(elem)) /
55 committee_size.max(1) as Threshold
56}
57
58pub fn pjr_check<AccountId: IdentifierT>(
62 supports: &Supports<AccountId>,
63 all_candidates: Vec<AccountId>,
64 all_voters: Vec<(AccountId, VoteWeight, Vec<AccountId>)>,
65) -> Result<(), AccountId> {
66 let t = standard_threshold(
67 supports.len(),
68 all_voters.iter().map(|voter| voter.1 as ExtendedBalance),
69 );
70 t_pjr_check(supports, all_candidates, all_voters, t)
71}
72
73pub fn t_pjr_check<AccountId: IdentifierT>(
108 supports: &Supports<AccountId>,
109 all_candidates: Vec<AccountId>,
110 all_voters: Vec<(AccountId, VoteWeight, Vec<AccountId>)>,
111 t: Threshold,
112) -> Result<(), AccountId> {
113 let (candidates, voters) = prepare_pjr_input(supports, all_candidates, all_voters);
115 pjr_check_core(candidates.as_ref(), voters.as_ref(), t)
117}
118
119pub fn pjr_check_core<AccountId: IdentifierT>(
126 candidates: &[CandidatePtr<AccountId>],
127 voters: &[Voter<AccountId>],
128 t: Threshold,
129) -> Result<(), AccountId> {
130 let unelected = candidates.iter().filter(|c| !c.borrow().elected);
131 let maybe_max_pre_score = unelected
132 .map(|c| (pre_score(Rc::clone(c), voters, t), c.borrow().who.clone()))
133 .max();
134 match maybe_max_pre_score {
136 Some((max_pre_score, counter_example)) if max_pre_score >= t => Err(counter_example),
137 _ => Ok(()),
138 }
139}
140
141pub fn validate_pjr_challenge<AccountId: IdentifierT>(
152 counter_example: AccountId,
153 supports: &Supports<AccountId>,
154 all_candidates: Vec<AccountId>,
155 all_voters: Vec<(AccountId, VoteWeight, Vec<AccountId>)>,
156) -> bool {
157 let threshold = standard_threshold(
158 supports.len(),
159 all_voters.iter().map(|voter| voter.1 as ExtendedBalance),
160 );
161 validate_t_pjr_challenge(counter_example, supports, all_candidates, all_voters, threshold)
162}
163
164pub fn validate_t_pjr_challenge<AccountId: IdentifierT>(
175 counter_example: AccountId,
176 supports: &Supports<AccountId>,
177 all_candidates: Vec<AccountId>,
178 all_voters: Vec<(AccountId, VoteWeight, Vec<AccountId>)>,
179 threshold: Threshold,
180) -> bool {
181 let (candidates, voters) = prepare_pjr_input(supports, all_candidates, all_voters);
182 validate_pjr_challenge_core(counter_example, &candidates, &voters, threshold)
183}
184
185fn validate_pjr_challenge_core<AccountId: IdentifierT>(
194 counter_example: AccountId,
195 candidates: &[CandidatePtr<AccountId>],
196 voters: &[Voter<AccountId>],
197 threshold: Threshold,
198) -> bool {
199 let candidate =
211 match candidates.iter().find(|candidate| candidate.borrow().who == counter_example) {
212 None => return false,
213 Some(candidate) => candidate.clone(),
214 };
215 pre_score(candidate, voters, threshold) >= threshold
216}
217
218fn prepare_pjr_input<AccountId: IdentifierT>(
246 supports: &Supports<AccountId>,
247 all_candidates: Vec<AccountId>,
248 all_voters: Vec<(AccountId, VoteWeight, Vec<AccountId>)>,
249) -> (Vec<CandidatePtr<AccountId>>, Vec<Voter<AccountId>>) {
250 let mut candidates_index: BTreeMap<AccountId, usize> = BTreeMap::new();
251
252 let mut assignment_map: BTreeMap<AccountId, Vec<(AccountId, ExtendedBalance)>> =
254 BTreeMap::new();
255 for (winner_id, Support { voters, .. }) in supports.iter() {
256 for (voter_id, support) in voters.iter() {
257 assignment_map
258 .entry(voter_id.clone())
259 .or_default()
260 .push((winner_id.clone(), *support));
261 }
262 }
263
264 let supports: SupportMap<AccountId> = supports.iter().cloned().collect();
276
277 let candidates = all_candidates
279 .into_iter()
280 .enumerate()
281 .map(|(i, c)| {
282 candidates_index.insert(c.clone(), i);
283
284 let who = c;
286 let maybe_support = supports.get(&who);
287 let elected = maybe_support.is_some();
288 let backed_stake = maybe_support.map(|support| support.total).unwrap_or_default();
289
290 Candidate {
291 who,
292 elected,
293 backed_stake,
294 score: Default::default(),
295 approval_stake: Default::default(),
296 round: Default::default(),
297 }
298 .to_ptr()
299 })
300 .collect::<Vec<_>>();
301
302 let voters = all_voters
304 .into_iter()
305 .map(|(v, w, ts)| {
306 let mut edges: Vec<Edge<AccountId>> = Vec::with_capacity(ts.len());
307 for t in ts {
308 if edges.iter().any(|e| e.who == t) {
309 continue
311 }
312
313 if let Some(idx) = candidates_index.get(&t) {
314 let weight = assignment_map
316 .get(&v)
317 .and_then(|d| {
318 d.iter().find_map(|(x, y)| if x == &t { Some(y) } else { None })
319 })
320 .cloned()
321 .unwrap_or_default();
322 edges.push(Edge {
323 who: t,
324 candidate: Rc::clone(&candidates[*idx]),
325 weight,
326 load: Default::default(),
327 });
328 }
329 }
330
331 let who = v;
332 let budget: ExtendedBalance = w.into();
333 Voter { who, budget, edges, load: Default::default() }
334 })
335 .collect::<Vec<_>>();
336
337 (candidates, voters)
338}
339
340fn pre_score<AccountId: IdentifierT>(
347 unelected: CandidatePtr<AccountId>,
348 voters: &[Voter<AccountId>],
349 t: Threshold,
350) -> ExtendedBalance {
351 debug_assert!(!unelected.borrow().elected);
352 voters
353 .iter()
354 .filter(|v| v.votes_for(&unelected.borrow().who))
355 .fold(Zero::zero(), |acc: ExtendedBalance, voter| acc.saturating_add(slack(voter, t)))
356}
357
358fn slack<AccountId: IdentifierT>(voter: &Voter<AccountId>, t: Threshold) -> ExtendedBalance {
370 let budget = voter.budget;
371 let leftover = voter.edges.iter().fold(Zero::zero(), |acc: ExtendedBalance, edge| {
372 let candidate = edge.candidate.borrow();
373 if candidate.elected {
374 let extra =
375 Perbill::one().min(Perbill::from_rational(t, candidate.backed_stake)) * edge.weight;
376 acc.saturating_add(extra)
377 } else {
378 acc
380 }
381 });
382
383 budget.saturating_sub(leftover)
385}
386
387#[cfg(test)]
388mod tests {
389 use super::*;
390
391 fn setup_voter(who: u32, votes: Vec<(u32, u128, bool)>) -> Voter<u32> {
392 let mut voter = Voter::new(who);
393 let mut budget = 0u128;
394 let candidates = votes
395 .into_iter()
396 .map(|(t, w, e)| {
397 budget += w;
398 Candidate {
399 who: t,
400 elected: e,
401 backed_stake: w,
402 score: Default::default(),
403 approval_stake: Default::default(),
404 round: Default::default(),
405 }
406 })
407 .collect::<Vec<_>>();
408 let edges = candidates
409 .into_iter()
410 .map(|c| Edge {
411 who: c.who,
412 weight: c.backed_stake,
413 candidate: c.to_ptr(),
414 load: Default::default(),
415 })
416 .collect::<Vec<_>>();
417 voter.edges = edges;
418 voter.budget = budget;
419 voter
420 }
421
422 fn assert_core_failure<AccountId: IdentifierT>(
423 candidates: &[CandidatePtr<AccountId>],
424 voters: &[Voter<AccountId>],
425 t: Threshold,
426 ) {
427 let counter_example = pjr_check_core(candidates, voters, t).unwrap_err();
428 assert!(validate_pjr_challenge_core(counter_example, candidates, voters, t));
429 }
430
431 #[test]
432 fn slack_works() {
433 let voter = setup_voter(10, vec![(1, 10, true), (2, 20, true)]);
434
435 assert_eq!(slack(&voter, 15), 5);
436 assert_eq!(slack(&voter, 17), 3);
437 assert_eq!(slack(&voter, 10), 10);
438 assert_eq!(slack(&voter, 5), 20);
439 }
440
441 #[test]
442 fn pre_score_works() {
443 let v1 = setup_voter(10, vec![(1, 10, true), (2, 20, true), (3, 0, false)]);
445 let v2 = setup_voter(20, vec![(1, 5, true), (2, 5, true)]);
447 let v3 = setup_voter(30, vec![(1, 20, true), (2, 20, true), (3, 0, false)]);
449
450 let unelected = Candidate {
451 who: 3u32,
452 elected: false,
453 score: Default::default(),
454 approval_stake: Default::default(),
455 backed_stake: Default::default(),
456 round: Default::default(),
457 }
458 .to_ptr();
459 let score = pre_score(unelected, &vec![v1, v2, v3], 15);
460
461 assert_eq!(score, 15);
462 }
463
464 #[test]
465 fn can_convert_data_from_external_api() {
466 let all_candidates = vec![10, 20, 30, 40];
467 let all_voters = vec![
468 (1, 10, vec![10, 20, 30, 40]),
469 (2, 20, vec![10, 20, 30, 40]),
470 (3, 30, vec![10, 30]),
471 ];
472 let supports: Supports<u32> = vec![
474 (20, Support { total: 15, voters: vec![(1, 5), (2, 10)] }),
475 (40, Support { total: 15, voters: vec![(1, 5), (2, 10)] }),
476 ];
477
478 let (candidates, voters) = prepare_pjr_input(&supports, all_candidates, all_voters);
479
480 assert_eq!(
482 candidates
483 .iter()
484 .map(|c| (c.borrow().who, c.borrow().elected, c.borrow().backed_stake))
485 .collect::<Vec<_>>(),
486 vec![(10, false, 0), (20, true, 15), (30, false, 0), (40, true, 15)],
487 );
488
489 assert_eq!(
491 voters
492 .iter()
493 .map(|v| (
494 v.who,
495 v.budget,
496 v.edges.iter().map(|e| (e.who, e.weight)).collect::<Vec<_>>(),
497 ))
498 .collect::<Vec<_>>(),
499 vec![
500 (1, 10, vec![(10, 0), (20, 5), (30, 0), (40, 5)]),
501 (2, 20, vec![(10, 0), (20, 10), (30, 0), (40, 10)]),
502 (3, 30, vec![(10, 0), (30, 0)]),
503 ],
504 );
505
506 assert_core_failure(&candidates, &voters, 1);
509 assert_core_failure(&candidates, &voters, 10);
510 assert_core_failure(&candidates, &voters, 20);
511 }
512
513 #[test]
520 fn find_upper_bound_for_threshold_scenario_1() {
521 let all_candidates = vec![10, 20, 30, 40];
522 let all_voters = vec![
523 (1, 10, vec![10, 20, 30, 40]),
524 (2, 20, vec![10, 20, 30, 40]),
525 (3, 30, vec![10, 30]),
526 ];
527 let supports: Supports<u32> = vec![
529 (20, Support { total: 15, voters: vec![(1, 5), (2, 10)] }),
530 (40, Support { total: 15, voters: vec![(1, 5), (2, 10)] }),
531 ];
532
533 let (candidates, voters) = prepare_pjr_input(&supports, all_candidates, all_voters);
534
535 find_threshold_phase_change_for_scenario(candidates, voters);
536 }
537
538 #[test]
539 fn find_upper_bound_for_threshold_scenario_2() {
540 let all_candidates = vec![10, 20, 30, 40];
541 let all_voters = vec![
542 (1, 10, vec![10, 20, 30, 40]),
543 (2, 20, vec![10, 20, 30, 40]),
544 (3, 25, vec![10, 30]),
545 ];
546 let supports: Supports<u32> = vec![
548 (20, Support { total: 15, voters: vec![(1, 5), (2, 10)] }),
549 (40, Support { total: 15, voters: vec![(1, 5), (2, 10)] }),
550 ];
551
552 let (candidates, voters) = prepare_pjr_input(&supports, all_candidates, all_voters);
553
554 find_threshold_phase_change_for_scenario(candidates, voters);
555 }
556
557 #[test]
558 fn find_upper_bound_for_threshold_scenario_3() {
559 let all_candidates = vec![10, 20, 30, 40];
560 let all_voters = vec![
561 (1, 10, vec![10, 20, 30, 40]),
562 (2, 20, vec![10, 20, 30, 40]),
563 (3, 35, vec![10, 30]),
564 ];
565 let supports: Supports<u32> = vec![
567 (20, Support { total: 15, voters: vec![(1, 5), (2, 10)] }),
568 (40, Support { total: 15, voters: vec![(1, 5), (2, 10)] }),
569 ];
570
571 let (candidates, voters) = prepare_pjr_input(&supports, all_candidates, all_voters);
572
573 find_threshold_phase_change_for_scenario(candidates, voters);
574 }
575
576 fn find_threshold_phase_change_for_scenario<AccountId: IdentifierT>(
577 candidates: Vec<CandidatePtr<AccountId>>,
578 voters: Vec<Voter<AccountId>>,
579 ) -> Threshold {
580 let mut threshold = 1;
581 let mut prev_threshold = 0;
582
583 while pjr_check_core(&candidates, &voters, threshold).is_err() {
585 prev_threshold = threshold;
586 threshold = threshold
587 .checked_mul(2)
588 .expect("pjr check must fail before we run out of capacity in u128");
589 }
590
591 let mut high_bound = threshold;
593 let mut low_bound = prev_threshold;
594
595 while high_bound - low_bound > 1 {
596 let test = low_bound + ((high_bound - low_bound) / 2);
598 if pjr_check_core(&candidates, &voters, test).is_ok() {
599 high_bound = test;
600 } else {
601 low_bound = test;
602 }
603 }
604
605 println!("highest failing check: {}", low_bound);
606 println!("lowest succeeding check: {}", high_bound);
607
608 let mut unexpected_failures = Vec::new();
610 let mut unexpected_successes = Vec::new();
611 for t in 0..=low_bound {
612 if pjr_check_core(&candidates, &voters, t).is_ok() {
613 unexpected_successes.push(t);
614 }
615 }
616 for t in high_bound..(high_bound * 2) {
617 if pjr_check_core(&candidates, &voters, t).is_err() {
618 unexpected_failures.push(t);
619 }
620 }
621 dbg!(&unexpected_successes, &unexpected_failures);
622 assert!(unexpected_failures.is_empty() && unexpected_successes.is_empty());
623
624 high_bound
625 }
626}