Refactor: Replace manual key hash conversions with type-safe utilities, update imports, and adjust serialization logic across modules for consistency.

Author: lowhung

codec/src/map_parameters.rs

@@ -327,7 +327,7 @@ pub fn map_certificate(
             }),
             alonzo::Certificate::PoolRetirement(pool_key_hash, epoch) => Ok(TxCertificateWithPos {
                 cert: TxCertificate::PoolRetirement(PoolRetirement {
-                    operator: pool_key_hash.to_vec(),
+                    operator: to_pool_id(pool_key_hash),
                     epoch: *epoch,
                 }),
                 tx_identifier,
@@ -395,7 +395,7 @@ pub fn map_certificate(
                     Ok(TxCertificateWithPos {
                         cert: TxCertificate::StakeDelegation(StakeDelegation {
                             stake_address: map_stake_address(cred, network_id),
-                            operator: pool_key_hash.to_vec(),
+                            operator: to_pool_id(pool_key_hash),
                         }),
                         tx_identifier,
                         cert_index: cert_index.try_into().unwrap(),
@@ -415,8 +415,8 @@ pub fn map_certificate(
                     pool_metadata,
                 } => Ok(TxCertificateWithPos {
                     cert: TxCertificate::PoolRegistration(PoolRegistration {
-                        operator: operator.to_vec(),
-                        vrf_key_hash: vrf_keyhash.to_vec(),
+                        operator: to_pool_id(operator),
+                        vrf_key_hash: to_vrf_key(vrf_keyhash),
                         pledge: *pledge,
                         cost: *cost,
                         margin: Ratio {
@@ -435,7 +435,7 @@ pub fn map_certificate(
                             .into_iter()
                             .map(|v| {
                                 StakeAddress::new(
-                                    StakeCredential::AddrKeyHash(v.to_vec()),
+                                    StakeCredential::AddrKeyHash(to_hash(v)),
                                     network_id.clone(),
                                 )
                             })
@@ -455,7 +455,7 @@ pub fn map_certificate(
                 conway::Certificate::PoolRetirement(pool_key_hash, epoch) => {
                     Ok(TxCertificateWithPos {
                         cert: TxCertificate::PoolRetirement(PoolRetirement {
-                            operator: pool_key_hash.to_vec(),
+                            operator: to_pool_id(pool_key_hash),
                             epoch: *epoch,
                         }),
                         tx_identifier,
@@ -494,7 +494,7 @@ pub fn map_certificate(
                     Ok(TxCertificateWithPos {
                         cert: TxCertificate::StakeAndVoteDelegation(StakeAndVoteDelegation {
                             stake_address: map_stake_address(cred, network_id),
-                            operator: pool_key_hash.to_vec(),
+                            operator: to_pool_id(pool_key_hash),
                             drep: map_drep(drep),
                         }),
                         tx_identifier,
@@ -507,7 +507,7 @@ pub fn map_certificate(
                         cert: TxCertificate::StakeRegistrationAndDelegation(
                             StakeRegistrationAndDelegation {
                                 stake_address: map_stake_address(cred, network_id),
-                                operator: pool_key_hash.to_vec(),
+                                operator: to_pool_id(pool_key_hash),
                                 deposit: *coin,
                             },
                         ),
@@ -533,7 +533,7 @@ pub fn map_certificate(
                         cert: TxCertificate::StakeRegistrationAndStakeAndVoteDelegation(
                             StakeRegistrationAndStakeAndVoteDelegation {
                                 stake_address: map_stake_address(cred, network_id),
-                                operator: pool_key_hash.to_vec(),
+                                operator: to_pool_id(pool_key_hash),
                                 drep: map_drep(drep),
                                 deposit: *coin,
                             },

common/src/address.rs

@@ -398,7 +398,7 @@ impl StakeAddress {
             NetworkId::Testnet => bech32::Hrp::parse("stake_test")?,
         };
 
-        let data = self.to_binary()?;
+        let data = self.to_binary();
         Ok(bech32::encode::<bech32::Bech32>(hrp, &data.as_slice())?)
     }
 
@@ -431,20 +431,20 @@ impl StakeAddress {
     }
 
     /// Convert to binary format (29 bytes)
-    pub fn to_binary(&self) -> Result<Vec<u8>> {
+    pub fn to_binary(&self) -> Vec<u8> {
         let network_bits = match self.network {
             NetworkId::Mainnet => 0b1u8,
             NetworkId::Testnet => 0b0u8,
         };
 
-        let (stake_bits, stake_hash): (u8, &KeyHash) = match &self.credential {
-            StakeCredential::AddrKeyHash(data) => (0b1110, data),
-            StakeCredential::ScriptHash(data) => (0b1111, data),
+        let (stake_bits, stake_hash): (u8, &Vec<u8>) = match &self.credential {
+            StakeCredential::AddrKeyHash(data) => (0b1110, &data.to_vec()),
+            StakeCredential::ScriptHash(data) => (0b1111, &data.to_vec()),
         };
 
         let mut data = vec![network_bits | (stake_bits << 4)];
-        data.extend(stake_hash.as_ref());
-        data.try_into().map_err(|_| anyhow!("Invalid hash size for stake address"))
+        data.extend(stake_hash);
+        data
     }
 
     /// Read from binary format (29 bytes)
@@ -505,9 +505,7 @@ impl<C> minicbor::Encode<C> for StakeAddress {
         e: &mut minicbor::Encoder<W>,
         _ctx: &mut C,
     ) -> Result<(), minicbor::encode::Error<W::Error>> {
-        let data = self
-            .to_binary()
-            .map_err(|_| minicbor::encode::Error::message("Failed to convert to binary"))?;
+        let data = self.to_binary();
         e.bytes(&data.as_slice())?;
         Ok(())
     }
@@ -944,7 +942,7 @@ mod tests {
     #[test]
     fn stake_addresses_encode_mainnet_stake() {
         let address = mainnet_stake_address();
-        let binary = address.to_binary().unwrap();
+        let binary = address.to_binary();
 
         // CBOR encoding wraps the raw 29-byte stake address in a byte string:
         // - 0x58: CBOR major type 2 (byte string) with 1-byte length follows
@@ -965,7 +963,7 @@ mod tests {
     fn stake_addresses_decode_mainnet_stake() {
         let binary = {
             let mut v = vec![0x58, 0x1d];
-            v.extend_from_slice(&mainnet_stake_address().to_binary().unwrap().as_slice());
+            v.extend_from_slice(&mainnet_stake_address().to_binary().as_slice());
             v
         };
 

common/src/snapshot/pool_params.rs

@@ -15,8 +15,8 @@
 use super::streaming_snapshot::{
     cbor, Coin, Nullable, PoolMetadata, Relay, RewardAccount, Set, UnitInterval,
 };
-use crate::{PoolId, VrfKeyHash};
 use crate::types::AddrKeyhash;
+use crate::{PoolId, VrfKeyHash};
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct PoolParams {

modules/accounts_state/src/snapshot.rs

@@ -195,15 +195,15 @@ mod tests {
     use super::*;
     use acropolis_common::stake_addresses::StakeAddressState;
     use acropolis_common::NetworkId::Mainnet;
-    use acropolis_common::{PoolId, StakeAddress, StakeCredential};
+    use acropolis_common::{KeyHash, PoolId, StakeAddress, StakeCredential};
 
     // Helper function to create stake addresses for testing
     fn create_test_stake_address(id: u8) -> StakeAddress {
-        let mut hash = vec![0u8; 28];
+        let mut hash = [0u8; 28];
         hash[0] = id;
         StakeAddress {
             network: Mainnet,
-            credential: StakeCredential::AddrKeyHash(hash),
+            credential: StakeCredential::AddrKeyHash(KeyHash::new(hash)),
         }
     }
 

modules/accounts_state/src/state.rs

@@ -993,14 +993,14 @@ impl State {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use acropolis_common::crypto::keyhash_256;
+    use acropolis_common::crypto::{keyhash_224, keyhash_256};
     use acropolis_common::{
         protocol_params::ConwayParams, rational_number::RationalNumber, Anchor, Committee,
         Constitution, CostModel, DRepVotingThresholds, NetworkId, PoolVotingThresholds, Pot,
         PotDelta, Ratio, Registration, StakeAddress, StakeAddressDelta, StakeAndVoteDelegation,
         StakeCredential, StakeRegistrationAndStakeAndVoteDelegation,
         StakeRegistrationAndVoteDelegation, TxCertificateWithPos, TxIdentifier, VoteDelegation,
-        Withdrawal,
+        VrfKeyHash, Withdrawal,
     };
 
     // Helper to create a StakeAddress from a byte slice
@@ -1392,10 +1392,7 @@ mod tests {
 
         // Get the KeyHash once
         let spo1_hash = spo1.get_hash();
-
         let drep_key_hash = test_keyhash_from_bytes(&DREP_HASH);
-        let drep_script_hash = test_keyhash_from_bytes(&DREP_HASH);
-
         let pool_id_1 = PoolId::new(*spo1_hash);
 
         let certificates = vec![
@@ -1419,16 +1416,16 @@ mod tests {
             TxCertificateWithPos {
                 cert: TxCertificate::VoteDelegation(VoteDelegation {
                     stake_address: spo1.clone(),
-                    drep: DRepChoice::Key(DREP_HASH.to_vec()),
+                    drep: DRepChoice::Key(drep_key_hash),
                 }),
                 tx_identifier,
                 cert_index: 0,
             },
             TxCertificateWithPos {
                 cert: TxCertificate::StakeAndVoteDelegation(StakeAndVoteDelegation {
                     stake_address: spo2.clone(),
-                    operator: spo1.get_hash().to_vec(),
-                    drep: DRepChoice::Script(DREP_HASH.to_vec()),
+                    operator: pool_id_1,
+                    drep: DRepChoice::Script(drep_key_hash),
                 }),
                 tx_identifier,
                 cert_index: 0,
@@ -1448,7 +1445,7 @@ mod tests {
                 cert: TxCertificate::StakeRegistrationAndStakeAndVoteDelegation(
                     StakeRegistrationAndStakeAndVoteDelegation {
                         stake_address: spo4.clone(),
-                        operator: spo1.get_hash().to_vec(),
+                        operator: pool_id_1,
                         drep: DRepChoice::NoConfidence,
                         deposit: 1,
                     },

modules/historical_accounts_state/src/immutable_historical_account_store.rs

@@ -212,8 +212,10 @@ impl ImmutableHistoricalAccountStore {
         &self,
         account: &StakeAddress,
     ) -> Result<Option<Vec<DelegationUpdate>>> {
-        let mut immutable_delegations = self
-            .collect_partition::<DelegationUpdate>(&self.delegation_history, &account.get_hash())?;
+        let mut immutable_delegations = self.collect_partition::<DelegationUpdate>(
+            &self.delegation_history,
+            &account.get_hash().into_inner(),
+        )?;
 
         self.merge_pending(
             account,
@@ -229,8 +231,10 @@ impl ImmutableHistoricalAccountStore {
         &self,
         account: &StakeAddress,
     ) -> Result<Option<Vec<AccountWithdrawal>>> {
-        let mut immutable_mirs =
-            self.collect_partition::<AccountWithdrawal>(&self.mir_history, &account.get_hash())?;
+        let mut immutable_mirs = self.collect_partition::<AccountWithdrawal>(
+            &self.mir_history,
+            &account.get_hash().into_inner(),
+        )?;
 
         self.merge_pending(account, |e| e.mir_history.as_ref(), &mut immutable_mirs).await;
 
@@ -326,7 +330,7 @@ impl ImmutableHistoricalAccountStore {
 
     fn make_epoch_key(account: &StakeAddress, epoch: u32) -> [u8; 32] {
         let mut key = [0u8; 32];
-        key[..28].copy_from_slice(&account.get_credential().get_hash());
+        key[..28].copy_from_slice(&account.get_credential().get_hash().into_inner());
         key[28..32].copy_from_slice(&epoch.to_be_bytes());
         key
     }

modules/rest_blockfrost/src/handlers/epochs.rs

@@ -15,7 +15,6 @@ use acropolis_common::{
         spdd::{SPDDStateQuery, SPDDStateQueryResponse},
         utils::query_state,
     },
-    serialization::Bech32WithHrp,
     NetworkId, PoolId, StakeAddress, StakeCredential,
 };
 use anyhow::{anyhow, Result};

modules/rest_blockfrost/src/handlers/pools.rs

@@ -17,9 +17,7 @@ use acropolis_common::{
         utils::query_state,
     },
     rest_helper::ToCheckedF64,
-    serialization::Bech32WithHrp,
-    PoolRetirement, PoolUpdateAction, StakeCredential, TxIdentifier,
-    PoolId, PoolRetirement, PoolUpdateAction, StakeCredential, TxHash,
+    PoolId, PoolRetirement, PoolUpdateAction, StakeCredential, TxIdentifier,
 };
 use anyhow::Result;
 use caryatid_sdk::Context;

modules/rest_blockfrost/src/types.rs

@@ -7,7 +7,7 @@ use acropolis_common::{
     rest_helper::ToCheckedF64,
     serialization::{DisplayFromBech32, PoolPrefix},
     AssetAddressEntry, AssetMetadataStandard, AssetMintRecord, KeyHash, PolicyAsset,
-    PoolEpochState, PoolUpdateAction, Relay, TxHash, VrfKeyHash, Vote,
+    PoolEpochState, PoolUpdateAction, Relay, TxHash, Vote, VrfKeyHash,
 };
 use anyhow::Result;
 use num_traits::ToPrimitive;

modules/spo_state/src/state.rs

@@ -10,8 +10,8 @@ use acropolis_common::{
     params::TECHNICAL_PARAMETER_POOL_RETIRE_MAX_EPOCH,
     queries::governance::VoteRecord,
     stake_addresses::StakeAddressMap,
-    BlockInfo, KeyHash, PoolMetadata, PoolRetirement, PoolUpdateEvent, Relay,
-    StakeAddress, TxCertificate, TxIdentifier, Voter, VotingProcedures,
+    BlockInfo, KeyHash, PoolId, PoolMetadata, PoolRegistration, PoolRetirement, PoolUpdateEvent,
+    Relay, StakeAddress, TxCertificate, TxHash, TxIdentifier, Voter, VotingProcedures,
 };
 use anyhow::Result;
 use imbl::HashMap;
@@ -671,7 +671,8 @@ mod tests {
     use acropolis_common::hash::Hash;
     use acropolis_common::{
         state_history::{StateHistory, StateHistoryStore},
-        PoolRetirement, Ratio, StakeAddress, TxCertificate, TxCertificateWithPos, TxIdentifier, VrfKeyHash
+        PoolRetirement, Ratio, StakeAddress, TxCertificate, TxCertificateWithPos, TxIdentifier,
+        VrfKeyHash,
     };
     use tokio::sync::Mutex;
 
@@ -868,7 +869,7 @@ mod tests {
         assert!(state.handle_tx_certs(&block, &msg).is_ok());
 
         assert_eq!(1, state.spos.len());
-        let spo = state.spos.get(pool_id);
+        let spo = state.spos.get(&pool_id);
         assert!(spo.is_some());
 
         block.number = 1;
@@ -908,7 +909,7 @@ mod tests {
         });
         assert!(state.handle_tx_certs(&block, &msg).is_ok());
         assert_eq!(1, state.spos.len());
-        let spo = state.spos.get(pool_id);
+        let spo = state.spos.get(&pool_id);
         assert!(spo.is_some());
         history.lock().await.commit(block.number, state);
 
@@ -1038,7 +1039,7 @@ mod tests {
         let mut state = State::new(&save_blocks_store_config());
         let mut block = new_block(0);
         let mut msg = new_certs_msg();
-        let spo_id = keyhash_224(&[1_u8]);
+        let spo_id = test_pool_id_from_bytes(&[1]);
         msg.certificates.push(TxCertificateWithPos {
             cert: TxCertificate::PoolRegistration(default_pool_registration(spo_id.clone(), None)),
             tx_identifier: TxIdentifier::default(),
@@ -1047,11 +1048,9 @@ mod tests {
         let spo_id = test_pool_id_from_bytes(&[1]);
 
         msg.certificates.push(TxCertificateWithPos {
-            PoolRegistrationWithPos {
-                reg: default_pool_registration(spo_id, None),
-                tx_hash: TxHash::default(),
-                cert_index: 0,
-            },
+            cert: TxCertificate::PoolRegistration(default_pool_registration(spo_id.clone(), None)),
+            tx_identifier: TxIdentifier::default(),
+            cert_index: 0,
         });
         assert!(state.handle_tx_certs(&block, &msg).is_ok());
         block = new_block(2);