feat(data_structures): incorporate stakes tracker into ChainState

Author: aesedepece

data_structures/src/capabilities.rs

@@ -1,3 +1,5 @@
+use serde::{Deserialize, Serialize};
+
 #[repr(u8)]
 #[derive(Clone, Copy, Debug)]
 pub enum Capability {
@@ -7,7 +9,7 @@ pub enum Capability {
     Witnessing = 1,
 }
 
-#[derive(Copy, Clone, Debug, Default, PartialEq)]
+#[derive(Copy, Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
 pub struct CapabilityMap<T>
 where
     T: Default,

data_structures/src/chain/mod.rs

@@ -43,6 +43,7 @@ use crate::{
         versioning::{ProtocolVersion, Versioned},
         ProtobufConvert,
     },
+    staking::prelude::*,
     superblock::SuperBlockState,
     transaction::{
         CommitTransaction, DRTransaction, DRTransactionBody, Memoized, MintTransaction,
@@ -54,6 +55,7 @@ use crate::{
     },
     utxo_pool::{OldUnspentOutputsPool, OwnUnspentOutputsPool, UnspentOutputsPool},
     vrf::{BlockEligibilityClaim, DataRequestEligibilityClaim},
+    wit::Wit,
 };
 
 /// Keeps track of priority being used by transactions included in recent blocks, and provides
@@ -3776,6 +3778,8 @@ pub struct ChainState {
     /// Unspent Outputs Pool
     #[serde(skip)]
     pub unspent_outputs_pool: UnspentOutputsPool,
+    /// Tracks stakes for every validator in the network
+    pub stakes: Stakes<PublicKeyHash, Wit, Epoch, u64>,
 }
 
 impl ChainState {

data_structures/src/staking/aux.rs

@@ -1,17 +1,43 @@
 use std::{rc::Rc, str::FromStr, sync::RwLock};
 
-use super::prelude::*;
+use serde::{Deserialize, Serialize};
 
-/// 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>>>;
+use super::prelude::*;
 
 /// The resulting type for all the fallible functions in this module.
-pub type StakingResult<T, Address, Coins, Epoch> =
-    Result<T, StakesError<Address, Coins, Epoch>>;
+pub type StakingResult<T, Address, Coins, Epoch> = Result<T, StakesError<Address, Coins, Epoch>>;
+
+/// Newtype for a reference-counted and read-write-locked instance of `Stake`.
+///
+/// This newtype is needed for implementing `PartialEq` manually on the locked data, which cannot be done directly
+/// because those are externally owned types.
+#[derive(Clone, Debug, Default, Deserialize, Serialize)]
+pub struct SyncStake<Address, Coins, Epoch, Power>
+where
+    Address: Default,
+    Epoch: Default,
+{
+    /// The lock itself.
+    pub value: Rc<RwLock<Stake<Address, Coins, Epoch, Power>>>,
+}
+
+impl<Address, Coins, Epoch, Power> PartialEq for SyncStake<Address, Coins, Epoch, Power>
+where
+    Address: Default,
+    Epoch: Default + PartialEq,
+    Coins: PartialEq,
+{
+    fn eq(&self, other: &Self) -> bool {
+        let self_stake = self.value.read().unwrap();
+        let other_stake = other.value.read().unwrap();
+
+        self_stake.coins.eq(&other_stake.coins) && other_stake.epochs.eq(&other_stake.epochs)
+    }
+}
 
 /// Couples a validator address with a withdrawer address together. This is meant to be used in `Stakes` as the index
 /// for the `by_key` index.
-#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
+#[derive(Clone, Debug, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)]
 pub struct StakeKey<Address> {
     /// A validator address.
     pub validator: Address,
@@ -46,7 +72,7 @@ where
 
 /// Couples an amount of coins, a validator address and a withdrawer address together. This is meant to be used in
 /// `Stakes` as the index of the `by_coins` index.
-#[derive(Eq, Ord, PartialEq, PartialOrd)]
+#[derive(Clone, Debug, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)]
 pub struct CoinsAndAddresses<Coins, Address> {
     /// An amount of coins.
     pub coins: Coins,

data_structures/src/staking/stake.rs

@@ -1,10 +1,12 @@
 use std::marker::PhantomData;
 
+use serde::{Deserialize, Serialize};
+
 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)]
+#[derive(Copy, Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
 pub struct Stake<Address, Coins, Epoch, Power>
 where
     Address: Default,

data_structures/src/staking/stakes.rs

@@ -1,18 +1,19 @@
-use std::collections::btree_map::Entry;
-use std::collections::BTreeMap;
+use std::collections::{btree_map::Entry, BTreeMap};
 
 use itertools::Itertools;
+use serde::{Deserialize, Serialize};
 
 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)]
+#[derive(Clone, Debug, Deserialize, Default, PartialEq, Serialize)]
 pub struct Stakes<Address, Coins, Epoch, Power>
 where
-    Address: Default,
+    Address: Default + Ord,
+    Coins: Ord,
     Epoch: Default,
 {
     /// A listing of all the stakers, indexed by their address.
@@ -61,10 +62,11 @@ where
         let key = key.into();
 
         // Find or create a matching stake entry
-        let stake_arc = self.by_key.entry(key.clone()).or_default();
+        let stake = self.by_key.entry(key.clone()).or_default();
 
         // Actually increase the number of coins
-        stake_arc
+        stake
+            .value
             .write()?
             .add_stake(coins, epoch, self.minimum_stakeable)?;
 
@@ -75,9 +77,9 @@ where
             addresses: key,
         };
         self.by_coins.remove(&key);
-        self.by_coins.insert(key, stake_arc.clone());
+        self.by_coins.insert(key, stake.clone());
 
-        Ok(stake_arc.read()?.clone())
+        Ok(stake.value.read()?.clone())
     }
 
     /// Obtain a list of stakers, conveniently ordered by one of several strategies.
@@ -126,6 +128,7 @@ where
             .by_key
             .get(&key)
             .ok_or(StakesError::EntryNotFound { key })?
+            .value
             .read()?
             .power(capability, epoch))
     }
@@ -141,6 +144,7 @@ where
             .iter()
             .flat_map(move |(CoinsAndAddresses { addresses, .. }, stake)| {
                 stake
+                    .value
                     .read()
                     .map(move |stake| (addresses.clone(), stake.power(capability, current_epoch)))
             })
@@ -161,7 +165,7 @@ where
 
         if let Entry::Occupied(mut by_address_entry) = self.by_key.entry(key.clone()) {
             let (initial_coins, final_coins) = {
-                let mut stake = by_address_entry.get_mut().write()?;
+                let mut stake = by_address_entry.get_mut().value.write()?;
 
                 // Check the former amount of stake
                 let initial_coins = stake.coins;
@@ -204,6 +208,7 @@ where
             .by_key
             .get_mut(&key)
             .ok_or(StakesError::EntryNotFound { key })?
+            .value
             .write()?;
         stake.epochs.update(capability, current_epoch);
 

node/src/actors/chain_manager/mod.rs

@@ -71,7 +71,6 @@ use witnet_data_structures::{
     data_request::DataRequestPool,
     get_environment,
     radon_report::{RadonReport, ReportContext},
-    staking::prelude::*,
     superblock::{ARSIdentities, AddSuperBlockVote, SuperBlockConsensus},
     transaction::{RevealTransaction, TallyTransaction, Transaction},
     types::{
@@ -80,7 +79,6 @@ use witnet_data_structures::{
     },
     utxo_pool::{Diff, OwnUnspentOutputsPool, UnspentOutputsPool, UtxoWriteBatch},
     vrf::VrfCtx,
-    wit::Wit,
 };
 use witnet_rad::types::RadonTypes;
 use witnet_util::timestamp::seconds_to_human_string;
@@ -249,8 +247,6 @@ pub struct ChainManager {
     import: Force<ChainImport<ImportError>>,
     /// Signals that a chain snapshot export is due.
     export: Force<PathBuf>,
-    /// Tracks stakes for every validator in the network.
-    stakes: Stakes<PublicKeyHash, Wit, Epoch, u64>,
 }
 
 impl ChainManager {

wallet/src/db/encrypted/mod.rs

@@ -4,7 +4,10 @@ use std::sync::Arc;
 use witnet_crypto::cipher;
 
 use super::*;
-use crate::{db::{encrypted::write_batch::EncryptedWriteBatch, GetWith}, types};
+use crate::{
+    db::{encrypted::write_batch::EncryptedWriteBatch, GetWith},
+    types,
+};
 
 mod engine;
 mod prefix;