Refactor: Replace KeyHash with StakeAddress for improved clarity and consistency across modules

lowhung · lowhung · commit 8c08f4440d20 · 2025-10-30T15:06:33.000-07:00
diff --git a/codec/src/block.rs b/codec/src/block.rs
@@ -16,7 +16,7 @@ pub fn map_to_block_issuer(
                 let digest = keyhash_224(vkey);
                 if let Some(issuer) = shelley_genesis_delegates
                     .values()
-                    .find(|v| v.delegate == digest.to_vec())
+                    .find(|v| v.delegate == digest)
                     .map(|i| BlockIssuer::GenesisDelegate(i.clone()))
                 {
                     Some(issuer)
diff --git a/common/src/queries/accounts.rs b/common/src/queries/accounts.rs
@@ -57,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>),
@@ -77,7 +77,7 @@ pub enum AccountsStateQueryResponse {
 
     // DReps-related responses
     DrepDelegators(DrepDelegators),
-    AccountsDrepDelegationsMap(HashMap<KeyHash, Option<DRepChoice>>),
+    AccountsDrepDelegationsMap(HashMap<StakeAddress, Option<DRepChoice>>),
 
     NotFound,
     Error(String),
@@ -173,7 +173,7 @@ 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)]
diff --git a/common/src/queries/pools.rs b/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)>,
 }
diff --git a/common/src/stake_addresses.rs b/common/src/stake_addresses.rs
@@ -166,15 +166,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
                     }
@@ -212,14 +212,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)
@@ -230,14 +229,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)
@@ -248,14 +246,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)
@@ -326,7 +323,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()
@@ -335,9 +332,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
     }
@@ -1244,8 +1241,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]
@@ -1358,8 +1355,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]
@@ -1467,14 +1464,14 @@ mod tests {
 
             assert_eq!(map.len(), 3);
             assert_eq!(
-                map.get(&addr1.get_hash()).unwrap(),
+                map.get(&addr1).unwrap(),
                 &Some(DRepChoice::Abstain)
             );
             assert_eq!(
-                map.get(&addr2.get_hash()).unwrap(),
+                map.get(&addr2).unwrap(),
                 &Some(DRepChoice::Key(DREP_HASH))
             );
-            assert_eq!(map.get(&addr3.get_hash()).unwrap(), &None);
+            assert_eq!(map.get(&addr3).unwrap(), &None);
         }
 
         #[test]
diff --git a/modules/accounts_state/src/spo_distribution_store.rs b/modules/accounts_state/src/spo_distribution_store.rs
@@ -1,7 +1,7 @@
 use std::collections::HashMap;
-
+use std::ops::Add;
 use acropolis_common::types::AddrKeyhash;
-use acropolis_common::PoolId;
+use acropolis_common::{PoolId, StakeAddress};
 use anyhow::Result;
 use fjall::{Config, Keyspace, PartitionCreateOptions};
 
diff --git a/modules/accounts_state/src/state.rs b/modules/accounts_state/src/state.rs
@@ -172,7 +172,7 @@ impl State {
     }
 
     /// 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)> {
         self.stake_addresses.lock().unwrap().get_pool_delegators(pool_operator)
     }
 
@@ -185,7 +185,7 @@ impl State {
     pub fn get_accounts_utxo_values_map(
         &self,
         stake_keys: &[StakeAddress],
-    ) -> Option<HashMap<KeyHash, u64>> {
+    ) -> Option<HashMap<StakeAddress, u64>> {
         let stake_addresses = self.stake_addresses.lock().ok()?; // If lock fails, return None
         stake_addresses.get_accounts_utxo_values_map(stake_keys)
     }
@@ -200,7 +200,7 @@ impl State {
     pub fn get_accounts_balances_map(
         &self,
         stake_keys: &[StakeAddress],
-    ) -> Option<HashMap<KeyHash, u64>> {
+    ) -> Option<HashMap<StakeAddress, u64>> {
         let stake_addresses = self.stake_addresses.lock().ok()?; // If lock fails, return None
         stake_addresses.get_accounts_balances_map(stake_keys)
     }
@@ -220,7 +220,7 @@ impl State {
     pub fn get_drep_delegations_map(
         &self,
         stake_keys: &[StakeAddress],
-    ) -> Option<HashMap<KeyHash, Option<DRepChoice>>> {
+    ) -> Option<HashMap<StakeAddress, Option<DRepChoice>>> {
         let stake_addresses = self.stake_addresses.lock().ok()?; // If lock fails, return None
         stake_addresses.get_drep_delegations_map(stake_keys)
     }
@@ -572,7 +572,7 @@ impl State {
         stake_addresses.generate_spdd()
     }
 
-    pub fn dump_spdd_state(&self) -> HashMap<PoolId, Vec<(KeyHash, u64)>> {
+    pub fn dump_spdd_state(&self) -> HashMap<PoolId, Vec<(StakeAddress, u64)>> {
         let stake_addresses = self.stake_addresses.lock().unwrap();
         stake_addresses.dump_spdd_state()
     }
diff --git a/modules/drep_state/src/state.rs b/modules/drep_state/src/state.rs
@@ -474,12 +474,12 @@ impl State {
         context: &Arc<Context<Message>>,
         delegators: &[(&StakeAddress, &DRepChoice)],
     ) -> Result<()> {
-        let mut stake_key_to_input = HashMap::with_capacity(delegators.len());
+        let mut stake_address_to_input = HashMap::with_capacity(delegators.len());
         let mut stake_addresses = Vec::with_capacity(delegators.len());
 
-        for &(sc, drep) in delegators {
-            stake_addresses.push(sc.clone());
-            stake_key_to_input.insert(sc.get_credential().get_hash(), (sc, drep));
+        for &(stake_address, drep) in delegators {
+            stake_addresses.push(stake_address.clone());
+            stake_address_to_input.insert(stake_address, (stake_address, drep));
         }
 
         let msg = Arc::new(Message::StateQuery(StateQuery::Accounts(
@@ -500,8 +500,8 @@ impl State {
             }
         };
 
-        for (stake_key, old_drep_opt) in result_map {
-            let &(delegator, new_drep_choice) = match stake_key_to_input.get(&stake_key) {
+        for (stake_address, old_drep_opt) in result_map {
+            let &(delegator, new_drep_choice) = match stake_address_to_input.get(&stake_address) {
                 Some(pair) => pair,
                 None => continue,
             };
diff --git a/modules/rest_blockfrost/src/handlers/epochs.rs b/modules/rest_blockfrost/src/handlers/epochs.rs
@@ -15,7 +15,7 @@ use acropolis_common::{
         spdd::{SPDDStateQuery, SPDDStateQueryResponse},
         utils::query_state,
     },
-    NetworkId, PoolId, StakeAddress, StakeCredential,
+    NetworkId, PoolId,
 };
 use anyhow::{anyhow, Result};
 use caryatid_sdk::Context;
@@ -486,16 +486,10 @@ pub async fn handle_epoch_total_stakes_blockfrost(
     .await?;
     let spdd_response = spdd
         .into_iter()
-        .map(|(pool_id, stake_key_hash, amount)| {
-            let stake_address = StakeAddress {
-                network: network.clone(),
-                credential: StakeCredential::AddrKeyHash(stake_key_hash),
-            }
-            .to_string()
-            .map_err(|e| anyhow::anyhow!("Failed to convert stake address to string: {e}"))?;
+        .map(|(pool_id, stake_address, amount)| {
             Ok(SPDDByEpochItemRest {
-                pool_id: pool_id.to_vec(),
-                stake_address,
+                pool_id,
+                stake_address: stake_address.to_string().unwrap(),
                 amount,
             })
         })
@@ -623,16 +617,9 @@ pub async fn handle_epoch_pool_stakes_blockfrost(
     .await?;
     let spdd_response = spdd
         .into_iter()
-        .map(|(key_hash, amount)| {
-            let stake_address = StakeAddress {
-                network: network.clone(),
-                credential: StakeCredential::AddrKeyHash(key_hash),
-            }
-            .to_string()
-            .map_err(|e| anyhow::anyhow!("Failed to convert stake address to string: {e}"))?;
-
+        .map(|(stake_address, amount)| {
             Ok(SPDDByEpochAndPoolItemRest {
-                stake_address,
+                stake_address: stake_address.to_string().unwrap(),
                 amount,
             })
         })
diff --git a/modules/rest_blockfrost/src/handlers/governance.rs b/modules/rest_blockfrost/src/handlers/governance.rs
@@ -4,14 +4,10 @@ use crate::types::{
     DRepInfoREST, DRepMetadataREST, DRepUpdateREST, DRepVoteREST, DRepsListREST, ProposalVoteREST,
     VoterRoleREST,
 };
-use acropolis_common::{
-    messages::{Message, RESTResponse, StateQuery, StateQueryResponse},
-    queries::{
-        accounts::{AccountsStateQuery, AccountsStateQueryResponse},
-        governance::{GovernanceStateQuery, GovernanceStateQueryResponse},
-    },
-    Credential, GovActionId, KeyHash, TxHash, Voter,
-};
+use acropolis_common::{messages::{Message, RESTResponse, StateQuery, StateQueryResponse}, queries::{
+    accounts::{AccountsStateQuery, AccountsStateQueryResponse},
+    governance::{GovernanceStateQuery, GovernanceStateQueryResponse},
+}, Credential, GovActionId, StakeAddress, TxHash, Voter};
 use anyhow::{anyhow, Result};
 use caryatid_sdk::Context;
 use reqwest::Client;
@@ -191,16 +187,15 @@ pub async fn handle_drep_delegators_blockfrost(
         Message::StateQueryResponse(StateQueryResponse::Governance(
             GovernanceStateQueryResponse::DRepDelegators(delegators),
         )) => {
-            let stake_key_to_bech32: HashMap<KeyHash, String> = match delegators
+            let stake_address_to_bech32: HashMap<StakeAddress, String> = match delegators
                 .addresses
                 .iter()
                 .map(|addr| {
                     let credential = addr.get_credential();
                     let bech32 = credential
                         .to_stake_bech32()
                         .map_err(|_| anyhow!("Failed to encode stake address"))?;
-                    let key_hash = credential.get_hash();
-                    Ok((key_hash, bech32))
+                    Ok((addr.clone(), bech32))
                 })
                 .collect::<Result<HashMap<_, _>>>()
             {
@@ -229,8 +224,8 @@ pub async fn handle_drep_delegators_blockfrost(
                 )) => {
                     let mut response = Vec::new();
 
-                    for (key, amount) in map {
-                        let Some(bech32) = stake_key_to_bech32.get(&key) else {
+                    for (stake_address, amount) in map {
+                        let Some(bech32) = stake_address_to_bech32.get(&stake_address) else {
                             return Ok(RESTResponse::with_text(
                                 500,
                                 "Internal error: missing Bech32 for stake key",
diff --git a/modules/rest_blockfrost/src/handlers/pools.rs b/modules/rest_blockfrost/src/handlers/pools.rs
@@ -17,7 +17,7 @@ use acropolis_common::{
         utils::query_state,
     },
     rest_helper::ToCheckedF64,
-    PoolId, PoolRetirement, PoolUpdateAction, StakeCredential, TxIdentifier,
+    PoolId, PoolRetirement, PoolUpdateAction, TxIdentifier,
 };
 use anyhow::Result;
 use caryatid_sdk::Context;
@@ -1004,9 +1004,9 @@ pub async fn handle_pool_delegators_blockfrost(
     };
 
     let mut delegators_rest = Vec::<PoolDelegatorRest>::new();
-    for (d, l) in pool_delegators {
-        let bech32 = StakeCredential::AddrKeyHash(d.clone())
-            .to_stake_bech32()
+    for (stake_address, l) in pool_delegators {
+        let bech32 = stake_address
+            .to_string()
             .map_err(|e| anyhow::anyhow!("Invalid stake address in pool delegators: {e}"))?;
         delegators_rest.push(PoolDelegatorRest {
             address: bech32,
diff --git a/modules/rest_blockfrost/src/types.rs b/modules/rest_blockfrost/src/types.rs
diff --git a/modules/spo_state/src/state.rs b/modules/spo_state/src/state.rs
diff --git a/modules/tx_unpacker/src/tx_unpacker.rs b/modules/tx_unpacker/src/tx_unpacker.rs

Original file line number	Diff line number	Diff line change
`@@ -16,7 +16,7 @@ pub fn map_to_block_issuer(`
`16`	`16`	`let digest = keyhash_224(vkey);`
`17`	`17`	`if let Some(issuer) = shelley_genesis_delegates`
`18`	`18`	`.values()`
`19`		`- .find(\|v\| v.delegate == digest.to_vec())`
	`19`	`+ .find(\|v\| v.delegate == digest)`
`20`	`20`	`.map(\|i\| BlockIssuer::GenesisDelegate(i.clone()))`
`21`	`21`	`{`
`22`	`22`	`Some(issuer)`
Original file line number	Diff line number	Diff line change
`@@ -172,7 +172,7 @@ impl State {`
`172`	`172`	`}`
`173`	`173`
`174`	`174`	`/// Get Pool Delegators with live_stakes`
`175`		`- pub fn get_pool_delegators(&self, pool_operator: &PoolId) -> Vec<(KeyHash, u64)> {`
	`175`	`+ pub fn get_pool_delegators(&self, pool_operator: &PoolId) -> Vec<(StakeAddress, u64)> {`
`176`	`176`	`self.stake_addresses.lock().unwrap().get_pool_delegators(pool_operator)`
`177`	`177`	`}`
`178`	`178`
`@@ -185,7 +185,7 @@ impl State {`
`185`	`185`	`pub fn get_accounts_utxo_values_map(`
`186`	`186`	`&self,`
`187`	`187`	`stake_keys: &[StakeAddress],`
`188`		`- ) -> Option<HashMap<KeyHash, u64>> {`
	`188`	`+ ) -> Option<HashMap<StakeAddress, u64>> {`
`189`	`189`	`let stake_addresses = self.stake_addresses.lock().ok()?; // If lock fails, return None`
`190`	`190`	`stake_addresses.get_accounts_utxo_values_map(stake_keys)`
`191`	`191`	`}`
`@@ -200,7 +200,7 @@ impl State {`
`200`	`200`	`pub fn get_accounts_balances_map(`
`201`	`201`	`&self,`
`202`	`202`	`stake_keys: &[StakeAddress],`
`203`		`- ) -> Option<HashMap<KeyHash, u64>> {`
	`203`	`+ ) -> Option<HashMap<StakeAddress, u64>> {`
`204`	`204`	`let stake_addresses = self.stake_addresses.lock().ok()?; // If lock fails, return None`
`205`	`205`	`stake_addresses.get_accounts_balances_map(stake_keys)`
`206`	`206`	`}`
`@@ -220,7 +220,7 @@ impl State {`
`220`	`220`	`pub fn get_drep_delegations_map(`
`221`	`221`	`&self,`
`222`	`222`	`stake_keys: &[StakeAddress],`
`223`		`- ) -> Option<HashMap<KeyHash, Option<DRepChoice>>> {`
	`223`	`+ ) -> Option<HashMap<StakeAddress, Option<DRepChoice>>> {`
`224`	`224`	`let stake_addresses = self.stake_addresses.lock().ok()?; // If lock fails, return None`
`225`	`225`	`stake_addresses.get_drep_delegations_map(stake_keys)`
`226`	`226`	`}`
`@@ -572,7 +572,7 @@ impl State {`
`572`	`572`	`stake_addresses.generate_spdd()`
`573`	`573`	`}`
`574`	`574`
`575`		`- pub fn dump_spdd_state(&self) -> HashMap<PoolId, Vec<(KeyHash, u64)>> {`
	`575`	`+ pub fn dump_spdd_state(&self) -> HashMap<PoolId, Vec<(StakeAddress, u64)>> {`
`576`	`576`	`let stake_addresses = self.stake_addresses.lock().unwrap();`
`577`	`577`	`stake_addresses.dump_spdd_state()`
`578`	`578`	`}`