Refactor: SPO distribution store refactored to use StakeAddress

common/src/address.rs

@@ -274,6 +274,7 @@ impl ShelleyAddress {
             ShelleyAddressPaymentPart::ScriptHash(data) => (data, 1),
         };
 
+        // TODO: MH - make delegation hash an Option<Hash<28>> type
         let (delegation_hash, delegation_bits): (Vec<u8>, u8) = match &self.delegation {
             ShelleyAddressDelegationPart::None => (Vec::new(), 3),
             ShelleyAddressDelegationPart::StakeKeyHash(hash) => (hash.to_vec(), 0),

common/src/queries/accounts.rs

@@ -1,8 +1,6 @@
 use std::collections::HashMap;
 
-use crate::{
-    DRepChoice, KeyHash, PoolId, PoolLiveStakeInfo, RewardType, StakeAddress, TxIdentifier,
-};
+use crate::{DRepChoice, PoolId, PoolLiveStakeInfo, RewardType, StakeAddress, TxIdentifier};
 
 pub const DEFAULT_ACCOUNTS_QUERY_TOPIC: (&str, &str) =
     ("accounts-state-query-topic", "cardano.query.accounts");
@@ -59,17 +57,17 @@ pub enum AccountsStateQueryResponse {
     AccountAssets(AccountAssets),
     AccountAssetsTotals(AccountAssetsTotals),
     AccountUTxOs(AccountUTxOs),
-    AccountsUtxoValuesMap(HashMap<KeyHash, u64>),
+    AccountsUtxoValuesMap(HashMap<StakeAddress, u64>),
     AccountsUtxoValuesSum(u64),
-    AccountsBalancesMap(HashMap<KeyHash, u64>),
+    AccountsBalancesMap(HashMap<StakeAddress, u64>),
     AccountsBalancesSum(u64),
 
     // Epochs-related responses
     ActiveStakes(u64),
-    /// Vec<(PoolId, StakeKey, ActiveStakeAmount)>
-    SPDDByEpoch(Vec<(PoolId, KeyHash, u64)>),
-    /// Vec<(StakeKey, ActiveStakeAmount)>
-    SPDDByEpochAndPool(Vec<(KeyHash, u64)>),
+    /// Vec<(PoolId, StakeAddress, ActiveStakeAmount)>
+    SPDDByEpoch(Vec<(PoolId, StakeAddress, u64)>),
+    /// Vec<(StakeAddress, ActiveStakeAmount)>
+    SPDDByEpochAndPool(Vec<(StakeAddress, u64)>),
 
     // Pools-related responses
     OptimalPoolSizing(Option<OptimalPoolSizing>),
@@ -79,7 +77,7 @@ pub enum AccountsStateQueryResponse {
 
     // DReps-related responses
     DrepDelegators(DrepDelegators),
-    AccountsDrepDelegationsMap(HashMap<KeyHash, Option<DRepChoice>>),
+    AccountsDrepDelegationsMap(HashMap<StakeAddress, Option<DRepChoice>>),
 
     NotFound,
     Error(String),
@@ -186,10 +184,10 @@ pub struct OptimalPoolSizing {
 
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
 pub struct PoolDelegators {
-    pub delegators: Vec<(KeyHash, u64)>,
+    pub delegators: Vec<(StakeAddress, u64)>,
 }
 
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
 pub struct DrepDelegators {
-    pub delegators: Vec<(KeyHash, u64)>,
+    pub delegators: Vec<(StakeAddress, u64)>,
 }

common/src/queries/pools.rs

@@ -1,6 +1,6 @@
 use crate::{
-    queries::governance::VoteRecord, rational_number::RationalNumber, KeyHash, PoolEpochState,
-    PoolId, PoolMetadata, PoolRegistration, PoolRetirement, PoolUpdateEvent, Relay,
+    queries::governance::VoteRecord, rational_number::RationalNumber, PoolEpochState, PoolId,
+    PoolMetadata, PoolRegistration, PoolRetirement, PoolUpdateEvent, Relay, StakeAddress,
 };
 
 pub const DEFAULT_POOLS_QUERY_TOPIC: (&str, &str) =
@@ -94,5 +94,5 @@ pub struct PoolActiveStakeInfo {
 
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
 pub struct PoolDelegators {
-    pub delegators: Vec<(KeyHash, u64)>,
+    pub delegators: Vec<(StakeAddress, u64)>,
 }

common/src/stake_addresses.rs

@@ -1,6 +1,6 @@
 use crate::{
     math::update_value_with_delta, messages::DRepDelegationDistribution, DRepChoice,
-    DRepCredential, DelegatedStake, KeyHash, Lovelace, PoolId, PoolLiveStakeInfo, StakeAddress,
+    DRepCredential, DelegatedStake, Lovelace, PoolId, PoolLiveStakeInfo, StakeAddress,
     StakeAddressDelta, Withdrawal,
 };
 use anyhow::Result;
@@ -163,15 +163,15 @@ impl StakeAddressMap {
     }
 
     /// Get Pool Delegators with live_stakes
-    pub fn get_pool_delegators(&self, pool_operator: &PoolId) -> Vec<(KeyHash, u64)> {
+    pub fn get_pool_delegators(&self, pool_operator: &PoolId) -> Vec<(StakeAddress, u64)> {
         // Find stake addresses delegated to pool_operator
-        let delegators: Vec<(KeyHash, u64)> = self
+        let delegators: Vec<(StakeAddress, u64)> = self
             .inner
             .iter()
             .filter_map(|(stake_address, sas)| match sas.delegated_spo.as_ref() {
                 Some(delegated_spo) => {
                     if delegated_spo.eq(pool_operator) {
-                        Some((*stake_address.get_hash(), sas.utxo_value + sas.rewards))
+                        Some((stake_address.clone(), sas.utxo_value + sas.rewards))
                     } else {
                         None
                     }
@@ -184,15 +184,15 @@ impl StakeAddressMap {
     }
 
     /// Get DRep Delegators with live_stakes
-    pub fn get_drep_delegators(&self, drep: &DRepChoice) -> Vec<(KeyHash, u64)> {
+    pub fn get_drep_delegators(&self, drep: &DRepChoice) -> Vec<(StakeAddress, u64)> {
         // Find stake addresses delegated to drep
-        let delegators: Vec<(KeyHash, u64)> = self
+        let delegators: Vec<(StakeAddress, u64)> = self
             .inner
             .iter()
             .filter_map(|(stake_address, sas)| match sas.delegated_drep.as_ref() {
                 Some(delegated_drep) => {
                     if delegated_drep.eq(drep) {
-                        Some((*stake_address.get_hash(), sas.utxo_value))
+                        Some((stake_address.clone(), sas.utxo_value))
                     } else {
                         None
                     }
@@ -209,14 +209,13 @@ impl StakeAddressMap {
     pub fn get_accounts_utxo_values_map(
         &self,
         stake_addresses: &[StakeAddress],
-    ) -> Option<HashMap<KeyHash, u64>> {
+    ) -> Option<HashMap<StakeAddress, u64>> {
         let mut map = HashMap::new();
 
         for stake_address in stake_addresses {
             let account = self.get(stake_address)?;
             let utxo_value = account.utxo_value;
-            let key_hash = stake_address.get_hash();
-            map.insert(*key_hash, utxo_value);
+            map.insert(stake_address.clone(), utxo_value);
         }
 
         Some(map)
@@ -227,14 +226,13 @@ impl StakeAddressMap {
     pub fn get_accounts_balances_map(
         &self,
         stake_addresses: &[StakeAddress],
-    ) -> Option<HashMap<KeyHash, u64>> {
+    ) -> Option<HashMap<StakeAddress, u64>> {
         let mut map = HashMap::new();
 
         for stake_address in stake_addresses {
             let account = self.get(stake_address)?;
             let balance = account.utxo_value + account.rewards;
-            let key_hash = stake_address.get_hash();
-            map.insert(*key_hash, balance);
+            map.insert(stake_address.clone(), balance);
         }
 
         Some(map)
@@ -245,14 +243,13 @@ impl StakeAddressMap {
     pub fn get_drep_delegations_map(
         &self,
         stake_addresses: &[StakeAddress],
-    ) -> Option<HashMap<KeyHash, Option<DRepChoice>>> {
+    ) -> Option<HashMap<StakeAddress, Option<DRepChoice>>> {
         let mut map = HashMap::new();
 
         for stake_address in stake_addresses {
             let account = self.get(stake_address)?;
             let maybe_drep = account.delegated_drep.clone();
-            let key_hash = stake_address.get_hash();
-            map.insert(*key_hash, maybe_drep);
+            map.insert(stake_address.clone(), maybe_drep);
         }
 
         Some(map)
@@ -323,7 +320,7 @@ impl StakeAddressMap {
 
     /// Dump current Stake Pool Delegation Distribution State
     /// <PoolId -> (Stake Key, Active Stakes Amount)>
-    pub fn dump_spdd_state(&self) -> HashMap<PoolId, Vec<(KeyHash, u64)>> {
+    pub fn dump_spdd_state(&self) -> HashMap<PoolId, Vec<(StakeAddress, u64)>> {
         let entries: Vec<_> = self
             .inner
             .par_iter()
@@ -332,9 +329,9 @@ impl StakeAddressMap {
             })
             .collect();
 
-        let mut result: HashMap<PoolId, Vec<(KeyHash, u64)>> = HashMap::new();
+        let mut result: HashMap<PoolId, Vec<(StakeAddress, u64)>> = HashMap::new();
         for (spo, entry) in entries {
-            result.entry(spo).or_default().push((entry.0.get_credential().get_hash(), entry.1));
+            result.entry(spo).or_default().push((entry.0, entry.1));
         }
         result
     }
@@ -552,7 +549,7 @@ impl StakeAddressMap {
 mod tests {
     use super::*;
     use crate::hash::Hash;
-    use crate::{NetworkId, StakeAddress, StakeCredential};
+    use crate::{KeyHash, NetworkId, StakeAddress, StakeCredential};
 
     const STAKE_KEY_HASH: KeyHash = KeyHash::new([0x99; 28]);
     const STAKE_KEY_HASH_2: KeyHash = KeyHash::new([0xaa; 28]);
@@ -1248,8 +1245,8 @@ mod tests {
             let map = stake_addresses.get_accounts_utxo_values_map(&keys).unwrap();
 
             assert_eq!(map.len(), 2);
-            assert_eq!(map.get(addr1.get_hash()).copied().unwrap(), 1000);
-            assert_eq!(map.get(addr2.get_hash()).copied().unwrap(), 2000);
+            assert_eq!(map.get(&addr1).copied().unwrap(), 1000);
+            assert_eq!(map.get(&addr2).copied().unwrap(), 2000);
         }
 
         #[test]
@@ -1362,8 +1359,8 @@ mod tests {
             let map = stake_addresses.get_accounts_balances_map(&addresses).unwrap();
 
             assert_eq!(map.len(), 2);
-            assert_eq!(map.get(addr1.get_hash()).copied().unwrap(), 1100);
-            assert_eq!(map.get(addr2.get_hash()).copied().unwrap(), 2000);
+            assert_eq!(map.get(&addr1).copied().unwrap(), 1100);
+            assert_eq!(map.get(&addr2).copied().unwrap(), 2000);
         }
 
         #[test]
@@ -1470,15 +1467,9 @@ mod tests {
             let map = stake_addresses.get_drep_delegations_map(&addresses).unwrap();
 
             assert_eq!(map.len(), 3);
-            assert_eq!(
-                map.get(addr1.get_hash()).unwrap(),
-                &Some(DRepChoice::Abstain)
-            );
-            assert_eq!(
-                map.get(addr2.get_hash()).unwrap(),
-                &Some(DRepChoice::Key(DREP_HASH))
-            );
-            assert_eq!(map.get(addr3.get_hash()).unwrap(), &None);
+            assert_eq!(map.get(&addr1).unwrap(), &Some(DRepChoice::Abstain));
+            assert_eq!(map.get(&addr2).unwrap(), &Some(DRepChoice::Key(DREP_HASH)));
+            assert_eq!(map.get(&addr3).unwrap(), &None);
         }
 
         #[test]

modules/accounts_state/Cargo.toml

@@ -17,18 +17,17 @@ anyhow = { workspace = true }
 bigdecimal = "0.4.8"
 chrono = { workspace = true }
 config = { workspace = true }
-dashmap = { workspace = true }
 hex = { workspace = true }
 imbl = { workspace = true }
 serde = { workspace = true }
-serde_json = { workspace = true }
-serde_with = { workspace = true }
 tokio = { workspace = true }
 tracing = { workspace = true }
 fjall = "2.11.2"
-rayon = "1.10.0"
 csv = "1.3.1"
 itertools = "0.14.0"
 
+[dev-dependencies]
+tempfile = "3"
+
 [lib]
 path = "src/accounts_state.rs"