Refactor: Simplify Bech32 encoding for StakeAddress in governance and epochs handlers

Author: lowhung

modules/accounts_state/src/spo_distribution_store.rs

@@ -1,5 +1,5 @@
 use acropolis_common::{PoolId, StakeAddress};
-use anyhow::Result;
+use anyhow::{anyhow, Context, Result};
 use fjall::{Config, Keyspace, PartitionCreateOptions};
 use std::collections::HashMap;
 
@@ -30,19 +30,27 @@ fn encode_epoch_pool_prefix(epoch: u64, pool_id: &PoolId) -> Vec<u8> {
 }
 
 fn decode_key(key: &[u8]) -> Result<(u64, PoolId, StakeAddress)> {
-    if key.len() != TOTAL_KEY_LEN {
-        anyhow::bail!(
-            "Invalid key length: expected {}, got {}",
-            TOTAL_KEY_LEN,
-            key.len()
-        );
-    }
-
-    let epoch = u64::from_be_bytes(key[..EPOCH_LEN].try_into()?);
-    let pool_id = key[EPOCH_LEN..EPOCH_LEN + POOL_ID_LENGTH].try_into()?;
+    let epoch_bytes: [u8; EPOCH_LEN] = key[..EPOCH_LEN]
+        .try_into()
+        .map_err(|_| anyhow!("Failed to extract epoch bytes (offset 0-{})", EPOCH_LEN))?;
+    let epoch = u64::from_be_bytes(epoch_bytes);
+
+    let pool_id: PoolId = key[EPOCH_LEN..EPOCH_LEN + POOL_ID_LENGTH].try_into().map_err(|_| {
+        anyhow!(
+            "Failed to extract pool ID bytes (offset {}-{})",
+            EPOCH_LEN,
+            EPOCH_LEN + POOL_ID_LENGTH
+        )
+    })?;
 
     let stake_address_bytes = &key[EPOCH_LEN + POOL_ID_LENGTH..];
-    let stake_address = StakeAddress::from_binary(stake_address_bytes)?;
+    let stake_address = StakeAddress::from_binary(stake_address_bytes).with_context(|| {
+        format!(
+            "Failed to decode stake address from {} bytes at offset {}",
+            stake_address_bytes.len(),
+            EPOCH_LEN + POOL_ID_LENGTH
+        )
+    })?;
 
     Ok((epoch, pool_id, stake_address))
 }

modules/accounts_state/src/state.rs

@@ -15,9 +15,8 @@ use acropolis_common::{
     protocol_params::ProtocolParams,
     stake_addresses::{StakeAddressMap, StakeAddressState},
     BlockInfo, DRepChoice, DRepCredential, DelegatedStake, InstantaneousRewardSource,
-    InstantaneousRewardTarget, Lovelace, MoveInstantaneousReward, PoolId,
-    PoolLiveStakeInfo, PoolRegistration, Pot, SPORewards, StakeAddress, StakeRewardDelta,
-    TxCertificate,
+    InstantaneousRewardTarget, Lovelace, MoveInstantaneousReward, PoolId, PoolLiveStakeInfo,
+    PoolRegistration, Pot, SPORewards, StakeAddress, StakeRewardDelta, TxCertificate,
 };
 use anyhow::Result;
 use imbl::OrdMap;
@@ -962,7 +961,14 @@ impl State {
 mod tests {
     use super::*;
     use acropolis_common::crypto::{keyhash_224, keyhash_256};
-    use acropolis_common::{protocol_params::ConwayParams, rational_number::RationalNumber, Anchor, Committee, Constitution, CostModel, DRepVotingThresholds, KeyHash, NetworkId, PoolVotingThresholds, Pot, PotDelta, Ratio, Registration, StakeAddress, StakeAddressDelta, StakeAndVoteDelegation, StakeCredential, StakeRegistrationAndStakeAndVoteDelegation, StakeRegistrationAndVoteDelegation, TxCertificateWithPos, TxIdentifier, VoteDelegation, VrfKeyHash, Withdrawal};
+    use acropolis_common::{
+        protocol_params::ConwayParams, rational_number::RationalNumber, Anchor, Committee,
+        Constitution, CostModel, DRepVotingThresholds, KeyHash, NetworkId, PoolVotingThresholds,
+        Pot, PotDelta, Ratio, Registration, StakeAddress, StakeAddressDelta,
+        StakeAndVoteDelegation, StakeCredential, StakeRegistrationAndStakeAndVoteDelegation,
+        StakeRegistrationAndVoteDelegation, TxCertificateWithPos, TxIdentifier, VoteDelegation,
+        VrfKeyHash, Withdrawal,
+    };
 
     // Helper to create a StakeAddress from a byte slice
     fn create_address(hash: &[u8]) -> StakeAddress {

modules/rest_blockfrost/src/handlers/epochs.rs

@@ -456,9 +456,12 @@ pub async fn handle_epoch_total_stakes_blockfrost(
     let spdd_response = spdd
         .into_iter()
         .map(|(pool_id, stake_address, amount)| {
+            let bech32 = stake_address
+                .to_string()
+                .map_err(|e| anyhow::anyhow!("Failed to convert stake address to string {}", e))?;
             Ok(SPDDByEpochItemRest {
                 pool_id,
-                stake_address: stake_address.to_string().unwrap(),
+                stake_address: bech32,
                 amount,
             })
         })
@@ -556,8 +559,11 @@ pub async fn handle_epoch_pool_stakes_blockfrost(
     let spdd_response = spdd
         .into_iter()
         .map(|(stake_address, amount)| {
+            let bech32 = stake_address
+                .to_string()
+                .map_err(|e| anyhow::anyhow!("Failed to convert stake address to string {}", e))?;
             Ok(SPDDByEpochAndPoolItemRest {
-                stake_address: stake_address.to_string().unwrap(),
+                stake_address: bech32,
                 amount,
             })
         })

modules/rest_blockfrost/src/handlers/governance.rs

@@ -10,13 +10,13 @@ use acropolis_common::{
         accounts::{AccountsStateQuery, AccountsStateQueryResponse},
         governance::{GovernanceStateQuery, GovernanceStateQueryResponse},
     },
-    Credential, GovActionId, StakeAddress, TxHash, Voter,
+    Credential, GovActionId, TxHash, Voter,
 };
 use anyhow::{anyhow, Result};
 use caryatid_sdk::Context;
 use reqwest::Client;
 use serde_json::Value;
-use std::{collections::HashMap, sync::Arc};
+use std::sync::Arc;
 
 pub async fn handle_dreps_list_blockfrost(
     context: Arc<Context<Message>>,
@@ -191,27 +191,6 @@ pub async fn handle_drep_delegators_blockfrost(
         Message::StateQueryResponse(StateQueryResponse::Governance(
             GovernanceStateQueryResponse::DRepDelegators(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"))?;
-                    Ok((addr.clone(), bech32))
-                })
-                .collect::<Result<HashMap<_, _>>>()
-            {
-                Ok(map) => map,
-                Err(e) => {
-                    return Ok(RESTResponse::with_text(
-                        500,
-                        &format!("Internal error: {e}"),
-                    ));
-                }
-            };
-
             let msg = Arc::new(Message::StateQuery(StateQuery::Accounts(
                 AccountsStateQuery::GetAccountsUtxoValuesMap {
                     stake_addresses: delegators.addresses.clone(),
@@ -226,27 +205,31 @@ pub async fn handle_drep_delegators_blockfrost(
                 Message::StateQueryResponse(StateQueryResponse::Accounts(
                     AccountsStateQueryResponse::AccountsUtxoValuesMap(map),
                 )) => {
-                    let mut response = Vec::new();
-
-                    for (stake_address, amount) in map {
-                        let Some(bech32) = stake_address_to_bech32.get(&stake_address) else {
-                            return Ok(RESTResponse::with_text(
+                    let response: Result<Vec<_>> = map
+                        .into_iter()
+                        .map(|(stake_address, amount)| {
+                            let bech32 = stake_address
+                                .to_string()
+                                .map_err(|e| anyhow!("Failed to encode stake address {}", e))?;
+
+                            Ok(serde_json::json!({
+                                "address": bech32,
+                                "amount": amount.to_string(),
+                            }))
+                        })
+                        .collect();
+
+                    match response {
+                        Ok(response) => match serde_json::to_string_pretty(&response) {
+                            Ok(json) => Ok(RESTResponse::with_json(200, &json)),
+                            Err(e) => Ok(RESTResponse::with_text(
                                 500,
-                                "Internal error: missing Bech32 for stake key",
-                            ));
-                        };
-
-                        response.push(serde_json::json!({
-                            "address": bech32,
-                            "amount": amount.to_string(),
-                        }));
-                    }
-
-                    match serde_json::to_string_pretty(&response) {
-                        Ok(json) => Ok(RESTResponse::with_json(200, &json)),
+                                &format!("Failed to serialize DRep delegators: {e}"),
+                            )),
+                        },
                         Err(e) => Ok(RESTResponse::with_text(
                             500,
-                            &format!("Failed to serialize DRep delegators: {e}"),
+                            &format!("Internal error: {e}"),
                         )),
                     }
                 }
@@ -279,7 +262,6 @@ pub async fn handle_drep_delegators_blockfrost(
         _ => Ok(RESTResponse::with_text(500, "Unexpected message type")),
     }
 }
-
 pub async fn handle_drep_metadata_blockfrost(
     context: Arc<Context<Message>>,
     params: Vec<String>,