use crate::arithmetic::{traits::Zero, Normalizable, PerThing, Rational128, ThresholdOrd};
use alloc::{collections::btree_map::BTreeMap, rc::Rc, vec, vec::Vec};
use codec::{Decode, DecodeWithMemTracking, Encode, MaxEncodedLen};
use core::{cell::RefCell, cmp::Ordering, fmt::Debug};
use scale_info::TypeInfo;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
#[cfg(test)]
mod mock;
#[cfg(test)]
mod tests;
mod assignments;
pub mod balancing;
pub mod helpers;
pub mod node;
pub mod phragmen;
pub mod phragmms;
pub mod pjr;
pub mod reduce;
pub mod traits;
pub use assignments::{Assignment, StakedAssignment};
pub use balancing::*;
pub use helpers::*;
pub use phragmen::*;
pub use phragmms::*;
pub use pjr::*;
pub use reduce::reduce;
pub use traits::{IdentifierT, PerThing128};
#[derive(
Eq,
PartialEq,
Debug,
Clone,
codec::Encode,
codec::Decode,
codec::DecodeWithMemTracking,
scale_info::TypeInfo,
)]
pub enum Error {
SolutionWeightOverflow,
SolutionTargetOverflow,
SolutionInvalidIndex,
SolutionInvalidPageIndex,
ArithmeticError,
InvalidSupportEdge,
TooManyVoters,
BoundsExceeded,
DuplicateVoter,
DuplicateTarget,
}
pub type VoteWeight = u64;
pub type ExtendedBalance = u128;
#[derive(
Clone,
Copy,
PartialEq,
Eq,
Encode,
Decode,
DecodeWithMemTracking,
MaxEncodedLen,
TypeInfo,
Debug,
Default,
)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct ElectionScore {
pub minimal_stake: ExtendedBalance,
pub sum_stake: ExtendedBalance,
pub sum_stake_squared: ExtendedBalance,
}
#[cfg(feature = "std")]
impl ElectionScore {
pub fn pretty(&self, token: &str, decimals: u32) -> String {
format!(
"ElectionScore (minimal_stake: {}, sum_stake: {}, sum_stake_squared: {})",
pretty_balance(self.minimal_stake, token, decimals),
pretty_balance(self.sum_stake, token, decimals),
pretty_balance(self.sum_stake_squared, token, decimals),
)
}
}
#[cfg(feature = "std")]
pub fn pretty_balance<B: Into<u128>>(b: B, token: &str, decimals: u32) -> String {
let b: u128 = b.into();
format!("{} {}", b / 10u128.pow(decimals), token)
}
impl ElectionScore {
fn iter_by_significance(self) -> impl Iterator<Item = ExtendedBalance> {
[self.minimal_stake, self.sum_stake, self.sum_stake_squared].into_iter()
}
pub fn strict_threshold_better(self, other: Self, threshold: impl PerThing) -> bool {
match self
.iter_by_significance()
.zip(other.iter_by_significance())
.map(|(this, that)| (this.ge(&that), this.tcmp(&that, threshold.mul_ceil(that))))
.collect::<Vec<(bool, Ordering)>>()
.as_slice()
{
[(x, Ordering::Greater), _, _] => {
debug_assert!(x);
true
},
[(true, Ordering::Equal), (_, Ordering::Greater), _] => true,
[(true, Ordering::Equal), (true, Ordering::Equal), (_, Ordering::Less)] => true,
_ => false,
}
}
pub fn strict_better(self, other: Self) -> bool {
self.strict_threshold_better(other, crate::runtime::Perbill::zero())
}
}
impl core::cmp::Ord for ElectionScore {
fn cmp(&self, other: &Self) -> Ordering {
[self.minimal_stake, self.sum_stake, other.sum_stake_squared].cmp(&[
other.minimal_stake,
other.sum_stake,
self.sum_stake_squared,
])
}
}
impl core::cmp::PartialOrd for ElectionScore {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
#[derive(Clone, Copy)]
pub struct BalancingConfig {
pub iterations: usize,
pub tolerance: ExtendedBalance,
}
pub type CandidatePtr<A> = Rc<RefCell<Candidate<A>>>;
#[derive(Debug, Clone, Default)]
pub struct Candidate<AccountId> {
who: AccountId,
score: Rational128,
approval_stake: ExtendedBalance,
backed_stake: ExtendedBalance,
elected: bool,
round: usize,
}
impl<AccountId> Candidate<AccountId> {
pub fn to_ptr(self) -> CandidatePtr<AccountId> {
Rc::new(RefCell::new(self))
}
}
#[derive(Clone)]
pub struct Edge<AccountId> {
who: AccountId,
load: Rational128,
candidate: CandidatePtr<AccountId>,
weight: ExtendedBalance,
}
#[cfg(test)]
impl<AccountId: Clone> Edge<AccountId> {
fn new(candidate: Candidate<AccountId>, weight: ExtendedBalance) -> Self {
let who = candidate.who.clone();
let candidate = Rc::new(RefCell::new(candidate));
Self { weight, who, candidate, load: Default::default() }
}
}
#[cfg(feature = "std")]
impl<A: IdentifierT> core::fmt::Debug for Edge<A> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
write!(f, "Edge({:?}, weight = {:?})", self.who, self.weight)
}
}
#[derive(Clone, Default)]
pub struct Voter<AccountId> {
who: AccountId,
edges: Vec<Edge<AccountId>>,
budget: ExtendedBalance,
load: Rational128,
}
#[cfg(feature = "std")]
impl<A: IdentifierT> core::fmt::Debug for Voter<A> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
write!(f, "Voter({:?}, budget = {}, edges = {:?})", self.who, self.budget, self.edges)
}
}
impl<AccountId: IdentifierT> Voter<AccountId> {
pub fn new(who: AccountId) -> Self {
Self {
who,
edges: Default::default(),
budget: Default::default(),
load: Default::default(),
}
}
pub fn votes_for(&self, target: &AccountId) -> bool {
self.edges.iter().any(|e| &e.who == target)
}
pub fn into_assignment<P: PerThing>(self) -> Option<Assignment<AccountId, P>> {
let who = self.who;
let budget = self.budget;
let distribution = self
.edges
.into_iter()
.filter_map(|e| {
let per_thing = P::from_rational(e.weight, budget);
if per_thing.is_zero() {
None
} else {
Some((e.who, per_thing))
}
})
.collect::<Vec<_>>();
if distribution.len() > 0 {
Some(Assignment { who, distribution })
} else {
None
}
}
pub fn try_normalize(&mut self) -> Result<(), &'static str> {
let edge_weights = self.edges.iter().map(|e| e.weight).collect::<Vec<_>>();
edge_weights.normalize(self.budget).map(|normalized| {
for (edge, corrected) in self.edges.iter_mut().zip(normalized.into_iter()) {
let mut candidate = edge.candidate.borrow_mut();
candidate.backed_stake = candidate.backed_stake.saturating_sub(edge.weight);
edge.weight = corrected;
candidate.backed_stake = candidate.backed_stake.saturating_add(edge.weight);
}
})
}
pub fn try_normalize_elected(&mut self) -> Result<(), &'static str> {
let elected_edge_weights = self
.edges
.iter()
.filter_map(|e| if e.candidate.borrow().elected { Some(e.weight) } else { None })
.collect::<Vec<_>>();
elected_edge_weights.normalize(self.budget).map(|normalized| {
for (edge, corrected) in self
.edges
.iter_mut()
.filter(|e| e.candidate.borrow().elected)
.zip(normalized.into_iter())
{
let mut candidate = edge.candidate.borrow_mut();
candidate.backed_stake = candidate.backed_stake.saturating_sub(edge.weight);
edge.weight = corrected;
candidate.backed_stake = candidate.backed_stake.saturating_add(edge.weight);
}
})
}
#[inline]
pub fn budget(&self) -> ExtendedBalance {
self.budget
}
}
#[derive(Debug)]
pub struct ElectionResult<AccountId, P: PerThing> {
pub winners: Vec<(AccountId, ExtendedBalance)>,
pub assignments: Vec<Assignment<AccountId, P>>,
}
#[derive(Debug, Encode, Decode, DecodeWithMemTracking, Clone, Eq, PartialEq, TypeInfo)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct Support<AccountId> {
pub total: ExtendedBalance,
pub voters: Vec<(AccountId, ExtendedBalance)>,
}
impl<AccountId> Default for Support<AccountId> {
fn default() -> Self {
Self { total: Default::default(), voters: vec![] }
}
}
impl<AccountId> Backings for &Support<AccountId> {
fn total(&self) -> ExtendedBalance {
self.total
}
}
pub type Supports<A> = Vec<(A, Support<A>)>;
pub type SupportMap<A> = BTreeMap<A, Support<A>>;
pub fn to_support_map<AccountId: IdentifierT>(
assignments: &[StakedAssignment<AccountId>],
) -> SupportMap<AccountId> {
let mut supports = <BTreeMap<AccountId, Support<AccountId>>>::new();
for StakedAssignment { who, distribution } in assignments.iter() {
for (c, weight_extended) in distribution.iter() {
let support = supports.entry(c.clone()).or_default();
support.total = support.total.saturating_add(*weight_extended);
support.voters.push((who.clone(), *weight_extended));
}
}
supports
}
pub fn to_supports<AccountId: IdentifierT>(
assignments: &[StakedAssignment<AccountId>],
) -> Supports<AccountId> {
to_support_map(assignments).into_iter().collect()
}
pub trait EvaluateSupport {
fn evaluate(&self) -> ElectionScore;
}
impl<AccountId: IdentifierT> EvaluateSupport for Supports<AccountId> {
fn evaluate(&self) -> ElectionScore {
evaluate_support(self.iter().map(|(_, s)| s))
}
}
pub trait Backings {
fn total(&self) -> ExtendedBalance;
}
pub fn evaluate_support(backings: impl Iterator<Item = impl Backings>) -> ElectionScore {
let mut minimal_stake = ExtendedBalance::max_value();
let mut sum_stake: ExtendedBalance = Zero::zero();
let mut sum_stake_squared: ExtendedBalance = Zero::zero();
for support in backings {
sum_stake = sum_stake.saturating_add(support.total());
let squared = support.total().saturating_mul(support.total());
sum_stake_squared = sum_stake_squared.saturating_add(squared);
if support.total() < minimal_stake {
minimal_stake = support.total();
}
}
ElectionScore { minimal_stake, sum_stake, sum_stake_squared }
}
pub fn setup_inputs<AccountId: IdentifierT>(
initial_candidates: Vec<AccountId>,
initial_voters: Vec<(AccountId, VoteWeight, impl IntoIterator<Item = AccountId>)>,
) -> (Vec<CandidatePtr<AccountId>>, Vec<Voter<AccountId>>) {
let mut c_idx_cache = BTreeMap::<AccountId, usize>::new();
let candidates = initial_candidates
.into_iter()
.enumerate()
.map(|(idx, who)| {
c_idx_cache.insert(who.clone(), idx);
Candidate {
who,
score: Default::default(),
approval_stake: Default::default(),
backed_stake: Default::default(),
elected: Default::default(),
round: Default::default(),
}
.to_ptr()
})
.collect::<Vec<CandidatePtr<AccountId>>>();
let voters = initial_voters
.into_iter()
.filter_map(|(who, voter_stake, votes)| {
let mut edges: Vec<Edge<AccountId>> = Vec::new();
for v in votes {
if edges.iter().any(|e| e.who == v) {
continue;
}
if let Some(idx) = c_idx_cache.get(&v) {
let mut candidate = candidates[*idx].borrow_mut();
candidate.approval_stake =
candidate.approval_stake.saturating_add(voter_stake.into());
edges.push(Edge {
who: v.clone(),
candidate: Rc::clone(&candidates[*idx]),
load: Default::default(),
weight: Default::default(),
});
} }
if edges.is_empty() {
None
} else {
Some(Voter { who, edges, budget: voter_stake.into(), load: Rational128::zero() })
}
})
.collect::<Vec<_>>();
(candidates, voters)
}