diff --git a/data_structures/Cargo.toml b/data_structures/Cargo.toml
index fa801502b..78baa1c00 100644
--- a/data_structures/Cargo.toml
+++ b/data_structures/Cargo.toml
@@ -51,3 +51,7 @@ rand_distr = "0.4.3"
[[bench]]
name = "sort_active_identities"
harness = false
+
+[[bench]]
+name = "staking"
+harness = false
diff --git a/data_structures/benches/staking.rs b/data_structures/benches/staking.rs
new file mode 100644
index 000000000..8bbee63f8
--- /dev/null
+++ b/data_structures/benches/staking.rs
@@ -0,0 +1,85 @@
+#[macro_use]
+extern crate bencher;
+use bencher::Bencher;
+use rand::Rng;
+use witnet_data_structures::staking::prelude::*;
+
+fn populate(b: &mut Bencher) {
+ let mut stakes = Stakes::<String, u64, u64, u64>::default();
+ let mut i = 1;
+
+ b.iter(|| {
+ let address = format!("{i}");
+ let coins = i;
+ let epoch = i;
+ stakes.add_stake(address, coins, epoch).unwrap();
+
+ i += 1;
+ });
+}
+
+fn rank(b: &mut Bencher) {
+ let mut stakes = Stakes::<String, u64, u64, u64>::default();
+ let mut i = 1;
+
+ let stakers = 100_000;
+ let rf = 10;
+
+ let mut rng = rand::thread_rng();
+
+ loop {
+ let coins = i;
+ let epoch = i;
+ let address = format!("{}", rng.gen::<u64>());
+
+ stakes.add_stake(address, coins, epoch).unwrap();
+
+ i += 1;
+
+ if i == stakers {
+ break;
+ }
+ }
+
+ b.iter(|| {
+ let rank = stakes.rank(Capability::Mining, i);
+ let mut top = rank.take(usize::try_from(stakers / rf).unwrap());
+ let _first = top.next();
+ let _last = top.last();
+
+ i += 1;
+ })
+}
+
+fn query_power(b: &mut Bencher) {
+ let mut stakes = Stakes::<String, u64, u64, u64>::default();
+ let mut i = 1;
+
+ let stakers = 100_000;
+
+ loop {
+ let coins = i;
+ let epoch = i;
+ let address = format!("{i}");
+
+ stakes.add_stake(address, coins, epoch).unwrap();
+
+ i += 1;
+
+ if i == stakers {
+ break;
+ }
+ }
+
+ i = 1;
+
+ b.iter(|| {
+ let address = format!("{i}");
+ let _power = stakes.query_power(&address, Capability::Mining, i);
+
+ i += 1;
+ })
+}
+
+benchmark_main!(benches);
+benchmark_group!(benches, populate, rank, query_power);
diff --git a/data_structures/src/capabilities.rs b/data_structures/src/capabilities.rs
new file mode 100644
index 000000000..80cd8257b
--- /dev/null
+++ b/data_structures/src/capabilities.rs
@@ -0,0 +1,44 @@
+#[repr(u8)]
+#[derive(Clone, Copy, Debug)]
+pub enum Capability {
+ /// The base block mining and superblock voting capability
+ Mining = 0,
+ /// The universal HTTP GET / HTTP POST / WIP-0019 RNG capability
+ Witnessing = 1,
+}
+
+#[derive(Copy, Clone, Debug, Default, PartialEq)]
+pub struct CapabilityMap<T>
+where
+ T: Default,
+{
+ pub mining: T,
+ pub witnessing: T,
+}
+
+impl<T> CapabilityMap<T>
+where
+ T: Copy + Default,
+{
+ #[inline]
+ pub fn get(&self, capability: Capability) -> T {
+ match capability {
+ Capability::Mining => self.mining,
+ Capability::Witnessing => self.witnessing,
+ }
+ }
+
+ #[inline]
+ pub fn update(&mut self, capability: Capability, value: T) {
+ match capability {
+ Capability::Mining => self.mining = value,
+ Capability::Witnessing => self.witnessing = value,
+ }
+ }
+
+ #[inline]
+ pub fn update_all(&mut self, value: T) {
+ self.mining = value;
+ self.witnessing = value;
+ }
+}
diff --git a/data_structures/src/lib.rs b/data_structures/src/lib.rs
index 0b83c5cc9..e14acd6e5 100644
--- a/data_structures/src/lib.rs
+++ b/data_structures/src/lib.rs
@@ -38,6 +38,9 @@ pub mod fee;
/// Module containing data_request structures
pub mod data_request;
+/// Module containing data structures for the staking functionality
+pub mod staking;
+
/// Module containing superblock structures
pub mod superblock;
@@ -69,6 +72,9 @@ mod serialization_helpers;
/// Provides convenient constants, structs and methods for handling values denominated in Wit.
pub mod wit;
+/// Provides support for segmented protocol capabilities.
+pub mod capabilities;
+
lazy_static! {
/// Environment in which we are running: mainnet or testnet.
/// This is used for Bech32 serialization.
diff --git a/data_structures/src/staking/aux.rs b/data_structures/src/staking/aux.rs
new file mode 100644
index 000000000..424158164
--- /dev/null
+++ b/data_structures/src/staking/aux.rs
@@ -0,0 +1,37 @@
+use std::rc::Rc;
+use std::sync::RwLock;
+
+use super::prelude::*;
+
+/// Type alias for a reference-counted and read-write-locked instance of `Stake`.
+pub type SyncStake<Address, Coins, Epoch, Power> = Rc<RwLock<Stake<Address, Coins, Epoch, Power>>>;
+
+/// The resulting type for all the fallible functions in this module.
+pub type Result<T, Address, Coins, Epoch> =
+ std::result::Result<T, StakesError<Address, Coins, Epoch>>;
+
+/// Couples an amount of coins and an address together. This is to be used in `Stakes` as the index
+/// of the `by_coins` index..
+#[derive(Eq, Ord, PartialEq, PartialOrd)]
+pub struct CoinsAndAddress<Coins, Address> {
+ /// An amount of coins.
+ pub coins: Coins,
+ /// The address of a staker.
+ pub address: Address,
+}
+
+/// Allows telling the `census` method in `Stakes` to source addresses from its internal `by_coins`
+/// following different strategies.
+#[repr(u8)]
+#[derive(Clone, Copy, Debug)]
+pub enum CensusStrategy {
+ /// Retrieve all addresses, ordered by decreasing power.
+ All = 0,
+ /// Retrieve every Nth address, ordered by decreasing power.
+ StepBy(usize) = 1,
+ /// Retrieve the most powerful N addresses, ordered by decreasing power.
+ Take(usize) = 2,
+ /// Retrieve a total of N addresses, evenly distributed from the index, ordered by decreasing
+ /// power.
+ Evenly(usize) = 3,
+}
diff --git a/data_structures/src/staking/constants.rs b/data_structures/src/staking/constants.rs
new file mode 100644
index 000000000..d461b0560
--- /dev/null
+++ b/data_structures/src/staking/constants.rs
@@ -0,0 +1,2 @@
+/// A minimum stakeable amount needs to exist to prevent spamming of the tracker.
+pub const MINIMUM_STAKEABLE_AMOUNT_WITS: u64 = 10_000;
diff --git a/data_structures/src/staking/errors.rs b/data_structures/src/staking/errors.rs
new file mode 100644
index 000000000..6169073f4
--- /dev/null
+++ b/data_structures/src/staking/errors.rs
@@ -0,0 +1,41 @@
+use std::sync::PoisonError;
+
+/// All errors related to the staking functionality.
+#[derive(Debug, PartialEq)]
+pub enum StakesError<Address, Coins, Epoch> {
+ /// The amount of coins being staked or the amount that remains after unstaking is below the
+ /// minimum stakeable amount.
+ AmountIsBelowMinimum {
+ /// The number of coins being staked or remaining after staking.
+ amount: Coins,
+ /// The minimum stakeable amount.
+ minimum: Coins,
+ },
+ /// Tried to query `Stakes` for information that belongs to the past.
+ EpochInThePast {
+ /// The Epoch being referred.
+ epoch: Epoch,
+ /// The latest Epoch.
+ latest: Epoch,
+ },
+ /// An operation thrown an Epoch value that overflows.
+ EpochOverflow {
+ /// The computed Epoch value.
+ computed: u64,
+ /// The maximum Epoch.
+ maximum: Epoch,
+ },
+ /// Tried to query `Stakes` for the address of a staker that is not registered in `Stakes`.
+ IdentityNotFound {
+ /// The unknown address.
+ identity: Address,
+ },
+ /// Tried to obtain a lock on a write-locked piece of data that is already locked.
+ PoisonedLock,
+}
+
+impl<T, Address, Coins, Epoch> From<PoisonError<T>> for StakesError<Address, Coins, Epoch> {
+ fn from(_value: PoisonError<T>) -> Self {
+ StakesError::PoisonedLock
+ }
+}
diff --git a/data_structures/src/staking/mod.rs b/data_structures/src/staking/mod.rs
new file mode 100644
index 000000000..1a5b21418
--- /dev/null
+++ b/data_structures/src/staking/mod.rs
@@ -0,0 +1,107 @@
+#![deny(missing_docs)]
+
+/// Auxiliary convenience types and data structures.
+pub mod aux;
+/// Constants related to the staking functionality.
+pub mod constants;
+/// Errors related to the staking functionality.
+pub mod errors;
+/// The data structure and related logic for stake entries.
+pub mod stake;
+/// The data structure and related logic for keeping track of multiple stake entries.
+pub mod stakes;
+
+/// Module re-exporting virtually every submodule on a single level to ease importing of everything
+/// staking-related.
+pub mod prelude {
+ pub use crate::capabilities::*;
+
+ pub use super::aux::*;
+ pub use super::constants::*;
+ pub use super::errors::*;
+ pub use super::stake::*;
+ pub use super::stakes::*;
+}
+
+#[cfg(test)]
+pub mod test {
+ use super::prelude::*;
+
+ #[test]
+ fn test_e2e() {
+ let mut stakes = Stakes::<String, u64, u64, u64>::with_minimum(1);
+
+ // Alpha stakes 2 @ epoch 0
+ stakes.add_stake("Alpha", 2, 0).unwrap();
+
+ // Nobody holds any power just yet
+ let rank = stakes.rank(Capability::Mining, 0).collect::<Vec<_>>();
+ assert_eq!(rank, vec![("Alpha".into(), 0)]);
+
+ // One epoch later, Alpha starts to hold power
+ let rank = stakes.rank(Capability::Mining, 1).collect::<Vec<_>>();
+ assert_eq!(rank, vec![("Alpha".into(), 2)]);
+
+ // Beta stakes 5 @ epoch 10
+ stakes.add_stake("Beta", 5, 10).unwrap();
+
+ // Alpha is still leading, but Beta has scheduled its takeover
+ let rank = stakes.rank(Capability::Mining, 10).collect::<Vec<_>>();
+ assert_eq!(rank, vec![("Alpha".into(), 20), ("Beta".into(), 0)]);
+
+ // Beta eventually takes over after epoch 16
+ let rank = stakes.rank(Capability::Mining, 16).collect::<Vec<_>>();
+ assert_eq!(rank, vec![("Alpha".into(), 32), ("Beta".into(), 30)]);
+ let rank = stakes.rank(Capability::Mining, 17).collect::<Vec<_>>();
+ assert_eq!(rank, vec![("Beta".into(), 35), ("Alpha".into(), 34)]);
+
+ // Gamma should never take over, even in a million epochs, because it has only 1 coin
+ stakes.add_stake("Gamma", 1, 30).unwrap();
+ let rank = stakes
+ .rank(Capability::Mining, 1_000_000)
+ .collect::<Vec<_>>();
+ assert_eq!(
+ rank,
+ vec![
+ ("Beta".into(), 4_999_950),
+ ("Alpha".into(), 2_000_000),
+ ("Gamma".into(), 999_970)
+ ]
+ );
+
+ // But Delta is here to change it all
+ stakes.add_stake("Delta", 1_000, 50).unwrap();
+ let rank = stakes.rank(Capability::Mining, 50).collect::<Vec<_>>();
+ assert_eq!(
+ rank,
+ vec![
+ ("Beta".into(), 200),
+ ("Alpha".into(), 100),
+ ("Gamma".into(), 20),
+ ("Delta".into(), 0)
+ ]
+ );
+ let rank = stakes.rank(Capability::Mining, 51).collect::<Vec<_>>();
+ assert_eq!(
+ rank,
+ vec![
+ ("Delta".into(), 1_000),
+ ("Beta".into(), 205),
+ ("Alpha".into(), 102),
+ ("Gamma".into(), 21)
+ ]
+ );
+
+ // If Alpha removes all of its stake, it should immediately disappear
+ stakes.remove_stake("Alpha", 2).unwrap();
+ let rank = stakes.rank(Capability::Mining, 51).collect::<Vec<_>>();
+ assert_eq!(
+ rank,
+ vec![
+ ("Delta".into(), 1_000),
+ ("Beta".into(), 205),
+ ("Gamma".into(), 21),
+ ]
+ );
+ }
+}
diff --git a/data_structures/src/staking/simple.rs b/data_structures/src/staking/simple.rs
new file mode 100644
index 000000000..e69de29bb
diff --git a/data_structures/src/staking/stake.rs b/data_structures/src/staking/stake.rs
new file mode 100644
index 000000000..38fff6ae4
--- /dev/null
+++ b/data_structures/src/staking/stake.rs
@@ -0,0 +1,122 @@
+use std::marker::PhantomData;
+
+use super::prelude::*;
+
+/// A data structure that keeps track of a staker's staked coins and the epochs for different
+/// capabilities.
+#[derive(Copy, Clone, Debug, Default, PartialEq)]
+pub struct Stake<Address, Coins, Epoch, Power>
+where
+ Address: Default,
+ Epoch: Default,
+{
+ /// An amount of staked coins.
+ pub coins: Coins,
+ /// The average epoch used to derive coin age for different capabilities.
+ pub epochs: CapabilityMap<Epoch>,
+ // These two phantom fields are here just for the sake of specifying generics.
+ phantom_address: PhantomData<Address>,
+ phantom_power: PhantomData<Power>,
+}
+
+impl<Address, Coins, Epoch, Power> Stake<Address, Coins, Epoch, Power>
+where
+ Address: Default,
+ Coins: Copy
+ + From<u64>
+ + PartialOrd
+ + num_traits::Zero
+ + std::ops::Add<Output = Coins>
+ + std::ops::Sub<Output = Coins>
+ + std::ops::Mul
+ + std::ops::Mul<Epoch, Output = Power>,
+ Epoch: Copy + Default + num_traits::Saturating + std::ops::Sub<Output = Epoch>,
+ Power: std::ops::Add<Output = Power>
+ + std::ops::Div<Output = Power>
+ + std::ops::Div<Coins, Output = Epoch>,
+{
+ /// Increase the amount of coins staked by a certain staker.
+ ///
+ /// When adding stake:
+ /// - Amounts are added together.
+ /// - Epochs are weight-averaged, using the amounts as the weight.
+ ///
+ /// This type of averaging makes the entry equivalent to an unbounded record of all stake
+ /// additions and removals, without the overhead in memory and computation.
+ pub fn add_stake(
+ &mut self,
+ coins: Coins,
+ epoch: Epoch,
+ minimum_stakeable: Option<Coins>,
+ ) -> Result<Coins, Address, Coins, Epoch> {
+ // Make sure that the amount to be staked is equal or greater than the minimum
+ let minimum = minimum_stakeable.unwrap_or(Coins::from(MINIMUM_STAKEABLE_AMOUNT_WITS));
+ if coins < minimum {
+ Err(StakesError::AmountIsBelowMinimum {
+ amount: coins,
+ minimum,
+ })?;
+ }
+
+ let coins_before = self.coins;
+ let epoch_before = self.epochs.get(Capability::Mining);
+
+ let product_before = coins_before * epoch_before;
+ let product_added = coins * epoch;
+
+ let coins_after = coins_before + coins;
+ let epoch_after = (product_before + product_added) / coins_after;
+
+ self.coins = coins_after;
+ self.epochs.update_all(epoch_after);
+
+ Ok(coins_after)
+ }
+
+ /// Construct a Stake entry from a number of coins and a capability map. This is only useful for
+ /// tests.
+ #[cfg(test)]
+ pub fn from_parts(coins: Coins, epochs: CapabilityMap<Epoch>) -> Self {
+ Self {
+ coins,
+ epochs,
+ phantom_address: Default::default(),
+ phantom_power: Default::default(),
+ }
+ }
+
+ /// Derives the power of an identity in the network on a certain epoch from an entry. Most
+ /// normally, the epoch is the current epoch.
+ pub fn power(&self, capability: Capability, current_epoch: Epoch) -> Power {
+ self.coins * (current_epoch.saturating_sub(self.epochs.get(capability)))
+ }
+
+ /// Remove a certain amount of staked coins.
+ pub fn remove_stake(
+ &mut self,
+ coins: Coins,
+ minimum_stakeable: Option<Coins>,
+ ) -> Result<Coins, Address, Coins, Epoch> {
+ let coins_after = self.coins.sub(coins);
+
+ if coins_after > Coins::zero() {
+ let minimum = minimum_stakeable.unwrap_or(Coins::from(MINIMUM_STAKEABLE_AMOUNT_WITS));
+
+ if coins_after < minimum {
+ Err(StakesError::AmountIsBelowMinimum {
+ amount: coins_after,
+ minimum,
+ })?;
+ }
+ }
+
+ self.coins = coins_after;
+
+ Ok(self.coins)
+ }
+
+ /// Set the epoch for a certain capability. Most normally, the epoch is the current epoch.
+ pub fn reset_age(&mut self, capability: Capability, current_epoch: Epoch) {
+ self.epochs.update(capability, current_epoch);
+ }
+}
diff --git a/data_structures/src/staking/stakes.rs b/data_structures/src/staking/stakes.rs
new file mode 100644
index 000000000..13d0896d2
--- /dev/null
+++ b/data_structures/src/staking/stakes.rs
@@ -0,0 +1,466 @@
+use std::collections::btree_map::Entry;
+use std::collections::BTreeMap;
+
+use itertools::Itertools;
+
+use super::prelude::*;
+
+/// The main data structure that provides the "stakes tracker" functionality.
+///
+/// This structure holds indexes of stake entries. Because the entries themselves are reference
+/// counted and write-locked, we can have as many indexes here as we need at a negligible cost.
+#[derive(Default)]
+pub struct Stakes<Address, Coins, Epoch, Power>
+where
+ Address: Default,
+ Epoch: Default,
+{
+ /// A listing of all the stakers, indexed by their address.
+ by_address: BTreeMap<Address, SyncStake<Address, Coins, Epoch, Power>>,
+ /// A listing of all the stakers, indexed by their coins and address.
+ ///
+ /// Because this uses a compound key to prevent duplicates, if we want to know which addresses
+ /// have staked a particular amount, we just need to run a range lookup on the tree.
+ by_coins: BTreeMap<CoinsAndAddress<Coins, Address>, SyncStake<Address, Coins, Epoch, Power>>,
+ /// The amount of coins that can be staked or can be left staked after unstaking.
+ minimum_stakeable: Option<Coins>,
+}
+
+impl<Address, Coins, Epoch, Power> Stakes<Address, Coins, Epoch, Power>
+where
+ Address: Default,
+ Coins: Copy
+ + Default
+ + Ord
+ + From<u64>
+ + num_traits::Zero
+ + std::ops::Add<Output = Coins>
+ + std::ops::Sub<Output = Coins>
+ + std::ops::Mul
+ + std::ops::Mul<Epoch, Output = Power>,
+ Address: Clone + Ord + 'static,
+ Epoch: Copy + Default + num_traits::Saturating + std::ops::Sub<Output = Epoch>,
+ Power: Copy
+ + Default
+ + Ord
+ + std::ops::Add<Output = Power>
+ + std::ops::Div<Output = Power>
+ + std::ops::Div<Coins, Output = Epoch>
+ + 'static,
+{
+ /// Register a certain amount of additional stake for a certain address and epoch.
+ pub fn add_stake<IA>(
+ &mut self,
+ address: IA,
+ coins: Coins,
+ epoch: Epoch,
+ ) -> Result<Stake<Address, Coins, Epoch, Power>, Address, Coins, Epoch>
+ where
+ IA: Into<Address>,
+ {
+ let address = address.into();
+ let stake_arc = self.by_address.entry(address.clone()).or_default();
+
+ // Actually increase the number of coins
+ stake_arc
+ .write()?
+ .add_stake(coins, epoch, self.minimum_stakeable)?;
+
+ // Update the position of the staker in the `by_coins` index
+ // If this staker was not indexed by coins, this will index it now
+ let key = CoinsAndAddress {
+ coins,
+ address: address.clone(),
+ };
+ self.by_coins.remove(&key);
+ self.by_coins.insert(key, stake_arc.clone());
+
+ Ok(stake_arc.read()?.clone())
+ }
+
+ /// Obtain a list of stakers, conveniently ordered by one of several strategies.
+ ///
+ /// ## Strategies
+ ///
+ /// - `All`: retrieve all addresses, ordered by decreasing power.
+ /// - `StepBy`: retrieve every Nth address, ordered by decreasing power.
+ /// - `Take`: retrieve the most powerful N addresses, ordered by decreasing power.
+ /// - `Evenly`: retrieve a total of N addresses, evenly distributed from the index, ordered by
+ /// decreasing power.
+ pub fn census(
+ &self,
+ capability: Capability,
+ epoch: Epoch,
+ strategy: CensusStrategy,
+ ) -> Box<dyn Iterator<Item = Address>> {
+ let iterator = self.rank(capability, epoch).map(|(address, _)| address);
+
+ match strategy {
+ CensusStrategy::All => Box::new(iterator),
+ CensusStrategy::StepBy(step) => Box::new(iterator.step_by(step)),
+ CensusStrategy::Take(head) => Box::new(iterator.take(head)),
+ CensusStrategy::Evenly(count) => {
+ let collected = iterator.collect::<Vec<_>>();
+ let step = collected.len() / count;
+
+ Box::new(collected.into_iter().step_by(step).take(count))
+ }
+ }
+ }
+
+ /// Tells what is the power of an identity in the network on a certain epoch.
+ pub fn query_power(
+ &self,
+ address: &Address,
+ capability: Capability,
+ epoch: Epoch,
+ ) -> Result<Power, Address, Coins, Epoch> {
+ Ok(self
+ .by_address
+ .get(address)
+ .ok_or(StakesError::IdentityNotFound {
+ identity: address.clone(),
+ })?
+ .read()?
+ .power(capability, epoch))
+ }
+
+ /// For a given capability, obtain the full list of stakers ordered by their power in that
+ /// capability.
+ pub fn rank(
+ &self,
+ capability: Capability,
+ current_epoch: Epoch,
+ ) -> impl Iterator<Item = (Address, Power)> + 'static {
+ self.by_coins
+ .iter()
+ .flat_map(move |(CoinsAndAddress { address, .. }, stake)| {
+ stake
+ .read()
+ .map(move |stake| (address.clone(), stake.power(capability, current_epoch)))
+ })
+ .sorted_by_key(|(_, power)| *power)
+ .rev()
+ }
+
+ /// Remove a certain amount of staked coins from a given identity at a given epoch.
+ pub fn remove_stake<IA>(
+ &mut self,
+ address: IA,
+ coins: Coins,
+ ) -> Result<Coins, Address, Coins, Epoch>
+ where
+ IA: Into<Address>,
+ {
+ let address = address.into();
+ if let Entry::Occupied(mut by_address_entry) = self.by_address.entry(address.clone()) {
+ let (initial_coins, final_coins) = {
+ let mut stake = by_address_entry.get_mut().write()?;
+
+ // Check the former amount of stake
+ let initial_coins = stake.coins;
+
+ // Reduce the amount of stake
+ let final_coins = stake.remove_stake(coins, self.minimum_stakeable)?;
+
+ (initial_coins, final_coins)
+ };
+
+ // No need to keep the entry if the stake has gone to zero
+ if final_coins.is_zero() {
+ by_address_entry.remove();
+ self.by_coins.remove(&CoinsAndAddress {
+ coins: initial_coins,
+ address,
+ });
+ }
+
+ Ok(final_coins)
+ } else {
+ Err(StakesError::IdentityNotFound { identity: address })
+ }
+ }
+
+ /// Set the epoch for a certain address and capability. Most normally, the epoch is the current
+ /// epoch.
+ pub fn reset_age<IA>(
+ &mut self,
+ address: IA,
+ capability: Capability,
+ current_epoch: Epoch,
+ ) -> Result<(), Address, Coins, Epoch>
+ where
+ IA: Into<Address>,
+ {
+ let address = address.into();
+ let mut stake = self
+ .by_address
+ .get_mut(&address)
+ .ok_or(StakesError::IdentityNotFound { identity: address })?
+ .write()?;
+ stake.epochs.update(capability, current_epoch);
+
+ Ok(())
+ }
+
+ /// Creates an instance of `Stakes` with a custom minimum stakeable amount.
+ pub fn with_minimum(minimum: Coins) -> Self {
+ Stakes {
+ minimum_stakeable: Some(minimum),
+ ..Default::default()
+ }
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ #[test]
+ fn test_stakes_initialization() {
+ let stakes = Stakes::<String, u64, u64, u64>::default();
+ let ranking = stakes.rank(Capability::Mining, 0).collect::<Vec<_>>();
+ assert_eq!(ranking, Vec::default());
+ }
+
+ #[test]
+ fn test_add_stake() {
+ let mut stakes = Stakes::<String, u64, u64, u64>::with_minimum(5);
+ let alice = "Alice".into();
+ let bob = "Bob".into();
+
+ // Let's check default power
+ assert_eq!(
+ stakes.query_power(&alice, Capability::Mining, 0),
+ Err(StakesError::IdentityNotFound {
+ identity: alice.clone()
+ })
+ );
+ assert_eq!(
+ stakes.query_power(&alice, Capability::Mining, 1_000),
+ Err(StakesError::IdentityNotFound {
+ identity: alice.clone()
+ })
+ );
+
+ // Let's make Alice stake 100 Wit at epoch 100
+ assert_eq!(
+ stakes.add_stake(&alice, 100, 100).unwrap(),
+ Stake::from_parts(
+ 100,
+ CapabilityMap {
+ mining: 100,
+ witnessing: 100
+ }
+ )
+ );
+
+ // Let's see how Alice's stake accrues power over time
+ assert_eq!(stakes.query_power(&alice, Capability::Mining, 99), Ok(0));
+ assert_eq!(stakes.query_power(&alice, Capability::Mining, 100), Ok(0));
+ assert_eq!(stakes.query_power(&alice, Capability::Mining, 101), Ok(100));
+ assert_eq!(
+ stakes.query_power(&alice, Capability::Mining, 200),
+ Ok(10_000)
+ );
+
+ // Let's make Alice stake 50 Wits at epoch 150 this time
+ assert_eq!(
+ stakes.add_stake(&alice, 50, 300).unwrap(),
+ Stake::from_parts(
+ 150,
+ CapabilityMap {
+ mining: 166,
+ witnessing: 166
+ }
+ )
+ );
+ assert_eq!(
+ stakes.query_power(&alice, Capability::Mining, 299),
+ Ok(19_950)
+ );
+ assert_eq!(
+ stakes.query_power(&alice, Capability::Mining, 300),
+ Ok(20_100)
+ );
+ assert_eq!(
+ stakes.query_power(&alice, Capability::Mining, 301),
+ Ok(20_250)
+ );
+ assert_eq!(
+ stakes.query_power(&alice, Capability::Mining, 400),
+ Ok(35_100)
+ );
+
+ // Now let's make Bob stake 500 Wits at epoch 1000 this time
+ assert_eq!(
+ stakes.add_stake(&bob, 500, 1_000).unwrap(),
+ Stake::from_parts(
+ 500,
+ CapabilityMap {
+ mining: 1_000,
+ witnessing: 1_000
+ }
+ )
+ );
+
+ // Before Bob stakes, Alice has all the power
+ assert_eq!(
+ stakes.query_power(&alice, Capability::Mining, 999),
+ Ok(124950)
+ );
+ assert_eq!(stakes.query_power(&bob, Capability::Mining, 999), Ok(0));
+
+ // New stakes don't change power in the same epoch
+ assert_eq!(
+ stakes.query_power(&alice, Capability::Mining, 1_000),
+ Ok(125100)
+ );
+ assert_eq!(stakes.query_power(&bob, Capability::Mining, 1_000), Ok(0));
+
+ // Shortly after, Bob's stake starts to gain power
+ assert_eq!(
+ stakes.query_power(&alice, Capability::Mining, 1_001),
+ Ok(125250)
+ );
+ assert_eq!(stakes.query_power(&bob, Capability::Mining, 1_001), Ok(500));
+
+ // After enough time, Bob overpowers Alice
+ assert_eq!(
+ stakes.query_power(&alice, Capability::Mining, 2_000),
+ Ok(275_100)
+ );
+ assert_eq!(
+ stakes.query_power(&bob, Capability::Mining, 2_000),
+ Ok(500_000)
+ );
+ }
+
+ #[test]
+ fn test_coin_age_resets() {
+ // First, lets create a setup with a few stakers
+ let mut stakes = Stakes::<String, u64, u64, u64>::with_minimum(5);
+ let alice = "Alice".into();
+ let bob = "Bob".into();
+ let charlie = "Charlie".into();
+
+ stakes.add_stake(&alice, 10, 0).unwrap();
+ stakes.add_stake(&bob, 20, 20).unwrap();
+ stakes.add_stake(&charlie, 30, 30).unwrap();
+
+ // Let's really start our test at epoch 100
+ assert_eq!(
+ stakes.query_power(&alice, Capability::Mining, 100),
+ Ok(1_000)
+ );
+ assert_eq!(stakes.query_power(&bob, Capability::Mining, 100), Ok(1_600));
+ assert_eq!(
+ stakes.query_power(&charlie, Capability::Mining, 100),
+ Ok(2_100)
+ );
+ assert_eq!(
+ stakes.query_power(&alice, Capability::Witnessing, 100),
+ Ok(1_000)
+ );
+ assert_eq!(
+ stakes.query_power(&bob, Capability::Witnessing, 100),
+ Ok(1_600)
+ );
+ assert_eq!(
+ stakes.query_power(&charlie, Capability::Witnessing, 100),
+ Ok(2_100)
+ );
+ assert_eq!(
+ stakes.rank(Capability::Mining, 100).collect::<Vec<_>>(),
+ [
+ (charlie.clone(), 2100),
+ (bob.clone(), 1600),
+ (alice.clone(), 1000)
+ ]
+ );
+ assert_eq!(
+ stakes.rank(Capability::Witnessing, 100).collect::<Vec<_>>(),
+ [
+ (charlie.clone(), 2100),
+ (bob.clone(), 1600),
+ (alice.clone(), 1000)
+ ]
+ );
+
+ // Now let's slash Charlie's mining coin age right after
+ stakes.reset_age(&charlie, Capability::Mining, 101).unwrap();
+ assert_eq!(
+ stakes.query_power(&alice, Capability::Mining, 101),
+ Ok(1_010)
+ );
+ assert_eq!(stakes.query_power(&bob, Capability::Mining, 101), Ok(1_620));
+ assert_eq!(stakes.query_power(&charlie, Capability::Mining, 101), Ok(0));
+ assert_eq!(
+ stakes.query_power(&alice, Capability::Witnessing, 101),
+ Ok(1_010)
+ );
+ assert_eq!(
+ stakes.query_power(&bob, Capability::Witnessing, 101),
+ Ok(1_620)
+ );
+ assert_eq!(
+ stakes.query_power(&charlie, Capability::Witnessing, 101),
+ Ok(2_130)
+ );
+ assert_eq!(
+ stakes.rank(Capability::Mining, 101).collect::<Vec<_>>(),
+ [
+ (bob.clone(), 1_620),
+ (alice.clone(), 1_010),
+ (charlie.clone(), 0)
+ ]
+ );
+ assert_eq!(
+ stakes.rank(Capability::Witnessing, 101).collect::<Vec<_>>(),
+ [
+ (charlie.clone(), 2_130),
+ (bob.clone(), 1_620),
+ (alice.clone(), 1_010)
+ ]
+ );
+
+ // Don't panic, Charlie! After enough time, you can take over again ;)
+ assert_eq!(
+ stakes.query_power(&alice, Capability::Mining, 300),
+ Ok(3_000)
+ );
+ assert_eq!(stakes.query_power(&bob, Capability::Mining, 300), Ok(5_600));
+ assert_eq!(
+ stakes.query_power(&charlie, Capability::Mining, 300),
+ Ok(5_970)
+ );
+ assert_eq!(
+ stakes.query_power(&alice, Capability::Witnessing, 300),
+ Ok(3_000)
+ );
+ assert_eq!(
+ stakes.query_power(&bob, Capability::Witnessing, 300),
+ Ok(5_600)
+ );
+ assert_eq!(
+ stakes.query_power(&charlie, Capability::Witnessing, 300),
+ Ok(8_100)
+ );
+ assert_eq!(
+ stakes.rank(Capability::Mining, 300).collect::<Vec<_>>(),
+ [
+ (charlie.clone(), 5_970),
+ (bob.clone(), 5_600),
+ (alice.clone(), 3_000)
+ ]
+ );
+ assert_eq!(
+ stakes.rank(Capability::Witnessing, 300).collect::<Vec<_>>(),
+ [
+ (charlie.clone(), 8_100),
+ (bob.clone(), 5_600),
+ (alice.clone(), 3_000)
+ ]
+ );
+ }
+}