1use crate::{
25 balance, setup_inputs, BalancingConfig, CandidatePtr, ElectionResult, ExtendedBalance,
26 IdentifierT, PerThing128, VoteWeight, Voter,
27};
28use alloc::{rc::Rc, vec, vec::Vec};
29use sp_arithmetic::{traits::Bounded, PerThing, Rational128};
30
31pub fn phragmms<AccountId: IdentifierT, P: PerThing128>(
45 to_elect: usize,
46 candidates: Vec<AccountId>,
47 voters: Vec<(AccountId, VoteWeight, impl IntoIterator<Item = AccountId>)>,
48 balancing: Option<BalancingConfig>,
49) -> Result<ElectionResult<AccountId, P>, crate::Error> {
50 let (candidates, mut voters) = setup_inputs(candidates, voters);
51
52 let mut winners = vec![];
53 for round in 0..to_elect {
54 if let Some(round_winner) = calculate_max_score::<AccountId, P>(&candidates, &voters) {
55 apply_elected::<AccountId>(&mut voters, Rc::clone(&round_winner));
56
57 round_winner.borrow_mut().round = round;
58 round_winner.borrow_mut().elected = true;
59 winners.push(round_winner);
60
61 if let Some(ref config) = balancing {
62 balance(&mut voters, config);
63 }
64 } else {
65 break
66 }
67 }
68
69 let mut assignments =
70 voters.into_iter().filter_map(|v| v.into_assignment()).collect::<Vec<_>>();
71 assignments
72 .iter_mut()
73 .try_for_each(|a| a.try_normalize())
74 .map_err(|_| crate::Error::ArithmeticError)?;
75 let winners = winners
76 .into_iter()
77 .map(|w_ptr| (w_ptr.borrow().who.clone(), w_ptr.borrow().backed_stake))
78 .collect();
79
80 Ok(ElectionResult { winners, assignments })
81}
82
83pub(crate) fn calculate_max_score<AccountId: IdentifierT, P: PerThing>(
92 candidates: &[CandidatePtr<AccountId>],
93 voters: &[Voter<AccountId>],
94) -> Option<CandidatePtr<AccountId>> {
95 for c_ptr in candidates.iter() {
96 let mut candidate = c_ptr.borrow_mut();
97 if !candidate.elected {
98 candidate.score = Rational128::from(1, P::ACCURACY.into());
99 }
100 }
101
102 for voter in voters.iter() {
103 let mut denominator_contribution: ExtendedBalance = 0;
104
105 for edge in voter.edges.iter() {
107 let edge_candidate = edge.candidate.borrow();
108 if edge_candidate.elected {
109 let edge_contribution: ExtendedBalance =
110 P::from_rational(edge.weight, edge_candidate.backed_stake).deconstruct().into();
111 denominator_contribution += edge_contribution;
112 }
113 }
114
115 for edge in voter.edges.iter() {
117 let mut edge_candidate = edge.candidate.borrow_mut();
118 if !edge_candidate.elected {
119 let prev_d = edge_candidate.score.d();
120 edge_candidate.score = Rational128::from(1, denominator_contribution + prev_d);
121 }
122 }
123 }
124
125 let mut best_score = Rational128::zero();
127 let mut best_candidate = None;
128
129 for c_ptr in candidates.iter() {
130 let mut candidate = c_ptr.borrow_mut();
131 if candidate.approval_stake > 0 {
132 let score_d = candidate.score.d();
134 let one: ExtendedBalance = P::ACCURACY.into();
135 let score_n =
160 candidate.approval_stake.checked_mul(one).unwrap_or_else(Bounded::max_value);
161 candidate.score = Rational128::from(score_n, score_d);
162
163 if !candidate.elected && candidate.score > best_score {
165 best_score = candidate.score;
166 best_candidate = Some(Rc::clone(c_ptr));
167 }
168 } else {
169 candidate.score = Rational128::zero();
170 }
171 }
172
173 best_candidate
174}
175
176pub(crate) fn apply_elected<AccountId: IdentifierT>(
183 voters: &mut Vec<Voter<AccountId>>,
184 elected_ptr: CandidatePtr<AccountId>,
185) {
186 let elected_who = elected_ptr.borrow().who.clone();
187 let cutoff = elected_ptr
188 .borrow()
189 .score
190 .to_den(1)
191 .expect("(n / d) < u128::MAX and (n' / 1) == (n / d), thus n' < u128::MAX'; qed.")
192 .n();
193
194 let mut elected_backed_stake = elected_ptr.borrow().backed_stake;
195 for voter in voters {
196 if let Some(new_edge_index) = voter.edges.iter().position(|e| e.who == elected_who) {
197 let used_budget: ExtendedBalance = voter.edges.iter().map(|e| e.weight).sum();
198
199 let mut new_edge_weight = voter.budget.saturating_sub(used_budget);
200 elected_backed_stake = elected_backed_stake.saturating_add(new_edge_weight);
201
202 for (_, edge) in
204 voter.edges.iter_mut().enumerate().filter(|(edge_index, edge_inner)| {
205 *edge_index != new_edge_index && edge_inner.weight > 0
206 }) {
207 let mut edge_candidate = edge.candidate.borrow_mut();
208 if edge_candidate.backed_stake > cutoff {
209 let stake_to_take =
210 edge.weight.saturating_mul(cutoff) / edge_candidate.backed_stake.max(1);
211
212 edge.weight = edge.weight.saturating_sub(stake_to_take);
214 edge_candidate.backed_stake =
215 edge_candidate.backed_stake.saturating_sub(stake_to_take);
216
217 elected_backed_stake = elected_backed_stake.saturating_add(stake_to_take);
219 new_edge_weight = new_edge_weight.saturating_add(stake_to_take);
220 }
221 }
222
223 voter.edges[new_edge_index].weight = new_edge_weight;
224 }
225 }
226
227 elected_ptr.borrow_mut().backed_stake = elected_backed_stake;
229}
230
231#[cfg(test)]
232mod tests {
233 use super::*;
234 use crate::{Assignment, ElectionResult};
235 use alloc::rc::Rc;
236 use sp_runtime::{Perbill, Percent};
237
238 #[test]
239 fn basic_election_manual_works() {
240 let candidates = vec![1, 2, 3];
244 let voters = vec![(10, 10, vec![1, 2]), (20, 20, vec![1, 3]), (30, 30, vec![2, 3])];
245
246 let (candidates, mut voters) = setup_inputs(candidates, voters);
247
248 let winner =
250 calculate_max_score::<u32, Percent>(candidates.as_ref(), voters.as_ref()).unwrap();
251 assert_eq!(winner.borrow().who, 3);
252 assert_eq!(winner.borrow().score, 50u32.into());
253
254 apply_elected(&mut voters, Rc::clone(&winner));
255 assert_eq!(
256 voters
257 .iter()
258 .find(|x| x.who == 30)
259 .map(|v| (v.who, v.edges.iter().map(|e| (e.who, e.weight)).collect::<Vec<_>>()))
260 .unwrap(),
261 (30, vec![(2, 0), (3, 30)]),
262 );
263 assert_eq!(
264 voters
265 .iter()
266 .find(|x| x.who == 20)
267 .map(|v| (v.who, v.edges.iter().map(|e| (e.who, e.weight)).collect::<Vec<_>>()))
268 .unwrap(),
269 (20, vec![(1, 0), (3, 20)]),
270 );
271
272 winner.borrow_mut().elected = true;
274 winner.borrow_mut().round = 0;
275 drop(winner);
276
277 let config = BalancingConfig { iterations: 10, tolerance: 0 };
279 balance(&mut voters, &config);
280
281 let winner =
283 calculate_max_score::<u32, Percent>(candidates.as_ref(), voters.as_ref()).unwrap();
284 assert_eq!(winner.borrow().who, 2);
285 assert_eq!(winner.borrow().score, 25u32.into());
286
287 apply_elected(&mut voters, Rc::clone(&winner));
288 assert_eq!(
289 voters
290 .iter()
291 .find(|x| x.who == 30)
292 .map(|v| (v.who, v.edges.iter().map(|e| (e.who, e.weight)).collect::<Vec<_>>()))
293 .unwrap(),
294 (30, vec![(2, 15), (3, 15)]),
295 );
296 assert_eq!(
297 voters
298 .iter()
299 .find(|x| x.who == 20)
300 .map(|v| (v.who, v.edges.iter().map(|e| (e.who, e.weight)).collect::<Vec<_>>()))
301 .unwrap(),
302 (20, vec![(1, 0), (3, 20)]),
303 );
304 assert_eq!(
305 voters
306 .iter()
307 .find(|x| x.who == 10)
308 .map(|v| (v.who, v.edges.iter().map(|e| (e.who, e.weight)).collect::<Vec<_>>()))
309 .unwrap(),
310 (10, vec![(1, 0), (2, 10)]),
311 );
312
313 winner.borrow_mut().elected = true;
315 winner.borrow_mut().round = 0;
316 drop(winner);
317
318 balance(&mut voters, &config);
320
321 assert_eq!(
322 voters
323 .iter()
324 .find(|x| x.who == 30)
325 .map(|v| (v.who, v.edges.iter().map(|e| (e.who, e.weight)).collect::<Vec<_>>()))
326 .unwrap(),
327 (30, vec![(2, 20), (3, 10)]),
328 );
329 assert_eq!(
330 voters
331 .iter()
332 .find(|x| x.who == 20)
333 .map(|v| (v.who, v.edges.iter().map(|e| (e.who, e.weight)).collect::<Vec<_>>()))
334 .unwrap(),
335 (20, vec![(1, 0), (3, 20)]),
336 );
337 assert_eq!(
338 voters
339 .iter()
340 .find(|x| x.who == 10)
341 .map(|v| (v.who, v.edges.iter().map(|e| (e.who, e.weight)).collect::<Vec<_>>()))
342 .unwrap(),
343 (10, vec![(1, 0), (2, 10)]),
344 );
345 }
346
347 #[test]
348 fn basic_election_works() {
349 let candidates = vec![1, 2, 3];
350 let voters = vec![(10, 10, vec![1, 2]), (20, 20, vec![1, 3]), (30, 30, vec![2, 3])];
351
352 let config = BalancingConfig { iterations: 2, tolerance: 0 };
353 let ElectionResult::<_, Perbill> { winners, assignments } =
354 phragmms(2, candidates, voters, Some(config)).unwrap();
355 assert_eq!(winners, vec![(3, 30), (2, 30)]);
356 assert_eq!(
357 assignments,
358 vec![
359 Assignment { who: 10u64, distribution: vec![(2, Perbill::one())] },
360 Assignment { who: 20, distribution: vec![(3, Perbill::one())] },
361 Assignment {
362 who: 30,
363 distribution: vec![
364 (2, Perbill::from_parts(666666666)),
365 (3, Perbill::from_parts(333333334)),
366 ],
367 },
368 ]
369 )
370 }
371
372 #[test]
373 fn linear_voting_example_works() {
374 let candidates = vec![11, 21, 31, 41, 51, 61, 71];
375 let voters = vec![
376 (2, 2000, vec![11]),
377 (4, 1000, vec![11, 21]),
378 (6, 1000, vec![21, 31]),
379 (8, 1000, vec![31, 41]),
380 (110, 1000, vec![41, 51]),
381 (120, 1000, vec![51, 61]),
382 (130, 1000, vec![61, 71]),
383 ];
384
385 let config = BalancingConfig { iterations: 2, tolerance: 0 };
386 let ElectionResult::<_, Perbill> { winners, assignments: _ } =
387 phragmms(4, candidates, voters, Some(config)).unwrap();
388 assert_eq!(winners, vec![(11, 3000), (31, 2000), (51, 1500), (61, 1500),]);
389 }
390
391 #[test]
392 fn large_balance_wont_overflow() {
393 let candidates = vec![1u32, 2, 3];
394 let mut voters = (0..1000).map(|i| (10 + i, u64::MAX, vec![1, 2, 3])).collect::<Vec<_>>();
395
396 voters.push((2, u64::MAX, vec![1, 3]));
398
399 let config = BalancingConfig { iterations: 2, tolerance: 0 };
400 let ElectionResult::<_, Perbill> { winners, assignments: _ } =
401 phragmms(2, candidates, voters, Some(config)).unwrap();
402 assert_eq!(winners.into_iter().map(|(w, _)| w).collect::<Vec<_>>(), vec![1u32, 3]);
403 }
404}